Update gradio_mix.py

gradio_mix.py

@@ -42,11 +42,27 @@ langid.set_languages(['es','pt','zh','en','de','fr','it', 'ru', 'id', 'vi'])
 
 
 os.environ['CURL_CA_BUNDLE'] = ''
-DEMO_PATH = os.getenv("DEMO_PATH", "./demo")
-TMP_PATH = os.getenv("TMP_PATH", "./demo/temp")
+DEMO_PATH = os.getenv("DEMO_PATH", "./pretrained_models/demo")
+TMP_PATH = os.getenv("TMP_PATH", "./pretrained_models/demo/temp")
 MODELS_PATH = os.getenv("MODELS_PATH", "./pretrained_models")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Pick device for the TTS editing model. By default we try CUDA, but fall
+# back to CPU if the CUDA stack is not actually usable (e.g. kernel image
+# mismatch on older GPUs). You can override via LEMAS_DEVICE env (e.g. "cpu"
+# or "cuda").
+def _pick_device():
+    forced = os.getenv("LEMAS_DEVICE")
+    if forced:
+        return forced
+    if torch.cuda.is_available():
+        try:
+            torch.zeros(1).to("cuda")
+            return "cuda"
+        except Exception as e:
+            logging.warning("CUDA appears available but failed (%s); falling back to CPU.", e)
+    return "cpu"
+
+device = _pick_device()
 ASR_DEVICE = "cpu"  # force whisperx/pyannote to CPU to avoid cuDNN issues
 whisper_model, align_model = None, None
 tts_edit_model = None
@@ -75,14 +91,18 @@ class UVR5:
     """Small wrapper around the bundled uvr5 implementation for denoising."""
 
     def __init__(self, model_dir):
+        # Code directory is always the local `uvr5` folder in this repo
         code_dir = os.path.join(os.path.dirname(__file__), "uvr5")
         self.model = self.load_model(model_dir, code_dir)
     
     def load_model(self, model_dir, code_dir):
-        import sys, json
+        import sys, json, os
         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:
@@ -93,6 +113,9 @@ class UVR5:
             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
         )
@@ -332,7 +355,10 @@ class MMSAlignModel:
     def __init__(self):
         from torchaudio.pipelines import MMS_FA as bundle
         self.mms_model = bundle.get_model()
-        self.mms_model.to(device)
+        # MMS forced alignment is relatively light; keep it on CPU to avoid
+        # CUDA kernel / arch mismatches on environments where the main TTS
+        # model still uses GPU.
+        self.mms_model.to("cpu")
         self.mms_tokenizer = bundle.get_tokenizer()
         self.mms_aligner = bundle.get_aligner()
         self.text_normalizer = ur.Uroman() 
@@ -354,7 +380,7 @@ class MMSAlignModel:
 
     def compute_alignments(self, waveform: torch.Tensor, tokens):
         with torch.inference_mode():
-            emission, _ = self.mms_model(waveform.to(device))
+            emission, _ = self.mms_model(waveform.to("cpu"))
             token_spans = self.mms_aligner(emission[0], tokens)
         return emission, token_spans
 
@@ -373,7 +399,7 @@ class MMSAlignModel:
         assert len(text_normed) == len(raw_text), f"normalized text len != raw text len: {len(text_normed)} != {len(raw_text)}"
         tokens = self.mms_tokenizer(text_normed)
         with torch.inference_mode():
-            emission, _ = self.mms_model(waveform.to(device))
+            emission, _ = self.mms_model(waveform.to("cpu"))
             token_spans = self.mms_aligner(emission[0], tokens)
         num_frames = emission.size(1)
         ratio = waveform.size(1) / num_frames
@@ -391,18 +417,10 @@ class MMSAlignModel:
 class WhisperxModel:
     def __init__(self, model_name):
         from whisperx import load_model
-        from pathlib import Path
         prompt = None  # "This might be a blend of Simplified Chinese and English speech, do not translate, only transcription be allowed."
 
-        # Prefer a local VAD model (to avoid network download / 301 issues)
-        vad_fp = Path(MODELS_PATH) / "whisperx-vad-segmentation.bin"
-        if not vad_fp.is_file():
-            logging.warning(
-                "Local whisperx VAD not found at %s, falling back to default download path.",
-                vad_fp,
-            )
-            vad_fp = None
-
+        # Use the lighter Silero VAD backend to avoid pyannote checkpoints
+        # and their PyTorch 2.6 `weights_only` pickling issues.
         self.model = load_model(
             model_name,
             ASR_DEVICE,
@@ -417,7 +435,7 @@ class WhisperxModel:
                 "multilingual": True,
                 "hotwords": None
             },
-            vad_model_fp=str(vad_fp) if vad_fp is not None else None,
+            vad_method="silero",
         )
 
     def transcribe(self, audio_info, lang=None):
@@ -515,17 +533,20 @@ def get_audio_slice(audio, words_info, start_time, end_time, max_len=10, sr=1600
 def load_models(lemas_model_name, whisper_model_name, alignment_model_name, denoise_model_name):  # , audiosr_name):
 
     global transcribe_model, align_model, denoise_model, text_norm, tts_edit_model
-    # if voicecraft_model:
-    #     del denoise_model
-    #     del transcribe_model
-    #     del align_model
-    #     del voicecraft_model
-    #     del audiosr
     torch.cuda.empty_cache()
     gc.collect()
 
     if denoise_model_name == "UVR5":
-        denoise_model = UVR5(os.path.join(str(PRETRAINED_ROOT), "uvr5"))
+        # Simple layout: UVR5 assets live directly under:
+        #   <MODELS_PATH>/uvr5
+        # with files:
+        #   Kim_Vocal_1.onnx
+        #   MDX-Net-Kim-Vocal1.json
+        #   model_data.json
+        #   model_name_mapper.json
+        from pathlib import Path
+        uv_root = Path(MODELS_PATH) / "uvr5"
+        denoise_model = UVR5(str(uv_root))
     elif denoise_model_name == "DeepFilterNet":
         denoise_model = DeepFilterNet("./pretrained_models/denoiser_model.onnx")
 
@@ -943,8 +964,7 @@ def get_app():
                 # InvalidPathError with local filesystem paths.
                 _demo_value = None
                 demo_candidates = [
-                    os.path.join(DEMO_PATH, "V-00013_en-US.wav"),
-                    os.path.join(os.path.dirname(__file__), "..", "VoiceCraft", "demo", "V-00013_en-US.wav"),
+                    os.path.join(DEMO_PATH, "test.wav"),
                 ]
                 for demo_path in demo_candidates:
                     try:
@@ -1174,11 +1194,10 @@ def get_app():
 if __name__ == "__main__":
     import argparse
 
-    parser = argparse.ArgumentParser(description="VoiceCraft gradio app.")
+    parser = argparse.ArgumentParser(description="LEMAS-Edit gradio app.")
     
-    parser.add_argument("--demo-path", default="./demo", help="Path to demo directory")
-    parser.add_argument("--tmp-path", default="/cto_labs/vistring/zhaozhiyuan/outputs/voicecraft/tmp", help="Path to tmp directory")
-    parser.add_argument("--models-path", default="/cto_labs/vistring/zhaozhiyuan/outputs/voicecraft/pretrain/VoiceCraft", help="Path to voicecraft models directory")
+    parser.add_argument("--demo-path", default="./pretrained_models/demo", help="Path to demo directory")
+    parser.add_argument("--tmp-path", default="./pretrained_models/tmp", help="Path to tmp directory")
     parser.add_argument("--port", default=41020, type=int, help="App port")
     parser.add_argument("--share", action="store_true", help="Launch with public url")
     parser.add_argument("--server_name", default="0.0.0.0", type=str, help="Server name for launching the app. 127.0.0.1 for localhost; 0.0.0.0 to allow access from other machines in the local network. Might also give access to external users depends on the firewall settings.")