From 15e40406ae441017321ed9785341f9c0bc61f10d Mon Sep 17 00:00:00 2001
From: YellowRoseCx <80486540+YellowRoseCx@users.noreply.github.com>
Date: Sat, 8 Mar 2025 12:31:50 -0600
Subject: [PATCH] Fix model unloading issue on AMD GPUs

---
 gradio_interface.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/gradio_interface.py b/gradio_interface.py
index c112259c..4981fd63 100644
--- a/gradio_interface.py
+++ b/gradio_interface.py
@@ -1,3 +1,4 @@
+import gc
 import torch
 import torchaudio
 import gradio as gr
@@ -14,12 +15,17 @@
 SPEAKER_AUDIO_PATH = None
 
 
+def unload_model():
+    gc.collect()
+    torch.cuda.empty_cache()
+    return 'Unloaded'
+    
 def load_model_if_needed(model_choice: str):
     global CURRENT_MODEL_TYPE, CURRENT_MODEL
     if CURRENT_MODEL_TYPE != model_choice:
         if CURRENT_MODEL is not None:
             del CURRENT_MODEL
-            torch.cuda.empty_cache()
+            unload_model()
         print(f"Loading {model_choice} model...")
         CURRENT_MODEL = Zonos.from_pretrained(model_choice, device=device)
         CURRENT_MODEL.requires_grad_(False).eval()