Update app.py

app.py

@@ -52,7 +52,8 @@ api = HfApi()
 
 # Get all files in the repo
 all_files = api.list_repo_files(REPO_ID)
-
+print('ALL FILES')
+print(all_files)
 # Group files by model folder
 models = defaultdict(list)
 for f in all_files:
@@ -76,27 +77,56 @@ available_models = list(models.keys())
 print("Available models:", available_models)
 
 # Load homer model
-homer_model = load_model("homer")
-
-
-
+#homer_model = load_model("homer")
+
+
+
+
+def model_data(model_name, use_hf=False, repo_id="simpsonsaiorg/stream-models"):
+    """
+    Load model data either from local weights folder or from HuggingFace Hub.
+    
+    Args:
+        model_name (str): name of the model (e.g. "homer", "bart")
+        use_hf (bool): whether to load from HuggingFace instead of local
+        repo_id (str): HuggingFace repo id
+    
+    Returns:
+        tgt_sr, net_g, vc, version, index_file, if_f0
+    """
+    # -----------------------
+    # 1. Locate .pth file
+    # -----------------------
+    if not use_hf:
+        # Local
+        model_path = os.path.join(model_root, model_name)
+        pth_files = [os.path.join(model_path, f) for f in os.listdir(model_path) if f.endswith(".pth")]
+        index_files = [os.path.join(model_path, f) for f in os.listdir(model_path) if f.endswith(".index")]
+    else:
+        # HuggingFace
+        pth_files = [hf_hub_download(repo_id=repo_id, filename=f"{model_name}/model.pth")]
+        try:
+            index_files = [hf_hub_download(repo_id=repo_id, filename=f"{model_name}/model.index")]
+        except Exception:
+            index_files = []
 
-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}")
+        raise ValueError(f"No .pth file found for model {model_name}")
     pth_path = pth_files[0]
     print(f"Loading {pth_path}")
+
+    # -----------------------
+    # 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)
@@ -109,22 +139,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 = ""