Update app.py

app.py

@@ -104,7 +104,7 @@ def model_data(model_name, model_root="weights", repo_id="simpsonsaiorg/stream-m
         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))
@@ -112,30 +112,24 @@ def model_data(model_name, model_root="weights", repo_id="simpsonsaiorg/stream-m
         ]
     else:
         # --- HuggingFace load ---
-        # List all files in the repo first
         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 (pick the first one)
+        # 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 optional index files (match any pattern like added_IVF*_v2.index)
+        # 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}")
     
-    if len(pth_files) == 0:
-        raise ValueError(f"No .pth file found for model {model_name}")
-    pth_path = pth_files[0]
-    print(f"Loading {pth_path}")
+    print(f"Loading {pth_path}")  # <-- safe to do for both cases
+
 
     # -----------------------
     # 2. Load checkpoint