Update app.py

app.py

@@ -10,6 +10,7 @@ import gradio as gr
 import librosa
 import torch
 from fairseq import checkpoint_utils
+from fairseq.data import dictionary
 
 from config import Config
 from lib.infer_pack.models import (
@@ -20,6 +21,8 @@ from lib.infer_pack.models import (
 )
 from rmvpe import RMVPE
 from vc_infer_pipeline import VC
+from huggingface_hub import HfApi, hf_hub_download
+from collections import defaultdict  
 
 logging.getLogger("fairseq").setLevel(logging.WARNING)
 logging.getLogger("numba").setLevel(logging.WARNING)
@@ -36,30 +39,110 @@ tts_voice_list = asyncio.get_event_loop().run_until_complete(edge_tts.list_voice
 tts_voices = [f"{v['ShortName']}-{v['Gender']}" for v in tts_voice_list]
 
 model_root = "weights"
-models = [
+
+# ---- Local models ----
+local_models = [
     d for d in os.listdir(model_root) if os.path.isdir(os.path.join(model_root, d))
 ]
-if len(models) == 0:
-    raise ValueError("No model found in `weights` folder")
-models.sort()
-
-
-def model_data(model_name):
-    # global n_spk, tgt_sr, net_g, vc, cpt, version, index_file
-    pth_files = [
-        os.path.join(model_root, model_name, f)
-        for f in os.listdir(os.path.join(model_root, model_name))
-        if f.endswith(".pth")
-    ]
-    if len(pth_files) == 0:
-        raise ValueError(f"No pth file found in {model_root}/{model_name}")
-    pth_path = pth_files[0]
-    print(f"Loading {pth_path}")
+if not local_models:
+    print("⚠️ No model found in local `weights` folder")
+local_models.sort()
+
+# ---- HF models ----
+REPO_ID = "simpsonsaiorg/stream-models"
+api = HfApi()
+
+all_files = api.list_repo_files(REPO_ID)
+
+hf_models = defaultdict(list)
+for f in all_files:
+    parts = f.split("/")
+    if len(parts) == 3 and parts[0] == "weights":
+        model_name, filename = parts[1], parts[2]
+        hf_models[model_name].append(filename)
+
+# ---- Merge / display ----
+av_models = sorted(set(local_models) | set(hf_models.keys()))
+
+print("Local models:", local_models)
+print("HF models:", list(hf_models.keys()))
+print("Available models (combined):", av_models)
+
+models = hf_models
+
+
+# Example: load a specific model (like your hubert loader)
+"""
+def load_model(model_name):
+    if model_name not in models:
+        raise ValueError(f"Model '{model_name}' not found in repo")
+    files = models[model_name]
+    loaded = {}
+    for file in files:
+        path = hf_hub_download(repo_id=REPO_ID, filename=f"{model_name}/{file}")
+        loaded[file] = torch.load(path, map_location="cpu")
+    return loaded
+
+# Load homer model
+#homer_model = load_model("homer")
+"""
+
+
+
+def model_data(model_name, model_root="weights", repo_id="simpsonsaiorg/stream-models", use_hf=True):
+    """
+    Load a model either from local disk or HuggingFace repo.
+    """
+
+    if not use_hf:
+        # --- Local load ---
+        pth_files = [
+            os.path.join(model_root, model_name, f)
+            for f in os.listdir(os.path.join(model_root, model_name))
+            if f.endswith(".pth")
+        ]
+        if len(pth_files) == 0:
+            raise ValueError(f"No .pth file found in {model_root}/{model_name}")
+        pth_path = pth_files[0]
+    
+        index_files = [
+            os.path.join(model_root, model_name, f)
+            for f in os.listdir(os.path.join(model_root, model_name))
+            if f.endswith(".index")
+        ]
+    else:
+        # --- HuggingFace load ---
+        all_files = api.list_repo_files(repo_id)
+        model_files = [f for f in all_files if f.startswith(f"weights/{model_name}/")]
+        
+        # Find .pth file
+        pth_files = [f for f in model_files if f.endswith(".pth")]
+        if not pth_files:
+            raise ValueError(f"No .pth file found for model {model_name} in repo")
+        pth_path = hf_hub_download(repo_id=repo_id, filename=pth_files[0])
+        
+        # Find index files
+        index_files = [
+            hf_hub_download(repo_id=repo_id, filename=f)
+            for f in model_files
+            if f.endswith(".index") and "added_IVF" in f
+        ]
+    
+    print(f"Loading {pth_path}")  # <-- safe to do for both cases
+
+
+    # -----------------------
+    # 2. Load checkpoint
+    # -----------------------
     cpt = torch.load(pth_path, map_location="cpu")
     tgt_sr = cpt["config"][-1]
     cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]  # n_spk
     if_f0 = cpt.get("f0", 1)
     version = cpt.get("version", "v1")
+
+    # -----------------------
+    # 3. Init network
+    # -----------------------
     if version == "v1":
         if if_f0 == 1:
             net_g = SynthesizerTrnMs256NSFsid(*cpt["config"], is_half=config.is_half)
@@ -72,22 +155,21 @@ def model_data(model_name):
             net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"])
     else:
         raise ValueError("Unknown version")
+
     del net_g.enc_q
     net_g.load_state_dict(cpt["weight"], strict=False)
-    print("Model loaded")
     net_g.eval().to(config.device)
+
     if config.is_half:
         net_g = net_g.half()
     else:
         net_g = net_g.float()
+
     vc = VC(tgt_sr, config)
-    # n_spk = cpt["config"][-3]
 
-    index_files = [
-        os.path.join(model_root, model_name, f)
-        for f in os.listdir(os.path.join(model_root, model_name))
-        if f.endswith(".index")
-    ]
+    # -----------------------
+    # 4. Index file
+    # -----------------------
     if len(index_files) == 0:
         print("No index file found")
         index_file = ""
@@ -100,19 +182,17 @@ def model_data(model_name):
 
 def load_hubert():
     global hubert_model
-    models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
-        ["hubert_base.pt"],
-        suffix="",
-    )
+    safe_globals = [dictionary.Dictionary]  # allow this class
+    with torch.serialization.safe_globals(safe_globals):
+        models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
+            ["hubert_base.pt"],
+            suffix="",
+        )
     hubert_model = models[0]
     hubert_model = hubert_model.to(config.device)
-    if config.is_half:
-        hubert_model = hubert_model.half()
-    else:
-        hubert_model = hubert_model.float()
+    hubert_model = hubert_model.half() if config.is_half else hubert_model.float()
     return hubert_model.eval()
 
-
 print("Loading hubert model...")
 hubert_model = load_hubert()
 print("Hubert model loaded.")