Update gradio_mix.py

gradio_mix.py

@@ -51,17 +51,24 @@ MODELS_PATH = os.getenv("MODELS_PATH", "./pretrained_models")
 # Mirrors LEMAS-TTS `inference_gradio.py`.
 HF_PRETRAINED_ROOT = "hf://LEMAS-Project/LEMAS-TTS/pretrained_models"
 
+# Detect whether we are running inside a HF Space with stateless GPU.
+IS_SPACES = os.getenv("SYSTEM") == "spaces"
+
 # Pick device for the TTS editing model.
-# - Default: "cuda" if available, else "cpu"
-# - You can override via LEMAS_DEVICE env (e.g. "cpu" or "cuda").
+# - On Spaces (SYSTEM=spaces): always use CPU in the main process to respect
+#   stateless GPU constraints.
+# - Elsewhere: "cuda" if available, else "cpu", unless overridden via
+#   LEMAS_DEVICE env (e.g. "cpu" or "cuda").
 def _pick_device():
+    if IS_SPACES:
+        return "cpu"
     forced = os.getenv("LEMAS_DEVICE")
     if forced:
         return forced
     return "cuda" if torch.cuda.is_available() else "cpu"
 
 device = _pick_device()
-ASR_DEVICE = "cuda"  # force whisperx/pyannote to CPU to avoid cuDNN issues
+ASR_DEVICE = "cpu"  # force whisperx/pyannote to CPU to avoid cuDNN issues
 whisper_model, align_model = None, None
 tts_edit_model = None
 
@@ -94,32 +101,39 @@ class UVR5:
         self.model = self.load_model(model_dir, code_dir)
     
     def load_model(self, model_dir, code_dir):
-        import sys, json, os
+        import sys, json, os, torch
         if code_dir not in sys.path:
             sys.path.append(code_dir)
         from multiprocess_cuda_infer import ModelData, Inference
         # In the minimal LEMAS-TTS layout, UVR5 weights live under:
-        #   <pretrained_models>/uvr5/models/MDX_Net_Models/model_data/
-        # Here `model_dir` points to that `model_data` directory.
         model_path = os.path.join(model_dir, "Kim_Vocal_1.onnx")
         config_path = os.path.join(model_dir, "MDX-Net-Kim-Vocal1.json")
         with open(config_path, "r", encoding="utf-8") as f:
             configs = json.load(f)
         model_data = ModelData(
             model_path=model_path,
-            audio_path = model_dir,
-            result_path = model_dir,
-            device = _pick_device(),
-            process_method = "MDX-Net",
+            audio_path=model_dir,
+            result_path=model_dir,
+            device="cpu",
+            process_method="MDX-Net",
             # Keep base_dir and model_dir the same so all UVR5 metadata
             # (model_data.json, model_name_mapper.json, etc.) are resolved
             # under `pretrained_models/uvr5`, matching LEMAS-TTS inference.
             base_dir=model_dir,
-            **configs
+            **configs,
         )
 
-        uvr5_model = Inference(model_data, _pick_device())
-        uvr5_model.load_model(model_path, 1)
+        uvr5_model = Inference(model_data, "cpu")
+        # On HF Spaces with stateless GPU, we must not initialize CUDA in the
+        # main process. UVR5's internal `load_model` checks `torch.cuda.is_available()`
+        # and may touch CUDA APIs. Temporarily spoof this to force CPU-only
+        # providers during UVR5 init.
+        orig_is_available = torch.cuda.is_available
+        torch.cuda.is_available = lambda: False
+        try:
+            uvr5_model.load_model(model_path, 1)
+        finally:
+            torch.cuda.is_available = orig_is_available
         return uvr5_model
         
     def denoise(self, audio_info):
@@ -322,11 +336,9 @@ class TextNorm():
 def chunk_text(text, max_chars=135):
     """
     Splits the input text into chunks, each with a maximum number of characters.
-
     Args:
         text (str): The text to be split.
         max_chars (int): The maximum number of characters per chunk.
-
     Returns:
         List[str]: A list of text chunks.
     """
@@ -688,9 +700,7 @@ def align(transcript, audio_info, state):
         state
     ]
 
-@spaces.GPU
-@torch.no_grad()
-@torch.inference_mode()
+
 def denoise(audio_info):
     denoised_audio, sr = denoise_model.denoise(audio_info)
     denoised_audio = denoised_audio # .squeeze().numpy()
@@ -1249,4 +1259,4 @@ if __name__ == "__main__":
     MODELS_PATH = args.models_path
     
     app = get_app()
-    app.queue().launch(share=args.share, server_name=args.server_name, server_port=args.port)
+    app.queue().launch(share=args.share, server_name=args.server_name, server_port=args.port)