Update app.py

app.py

@@ -11,7 +11,10 @@ from vidshow import vidshow_dual_from_png_bytes, vidshow_from_png_bytes
 ACTION_PROBS_FILES = ("action_probs.tar.gz", "action_probs.zip")
 ACTION_PROBS_SUFFIXES = (".tar.gz", ".zip")
 
-JAXGMG_COLLECTION_SLUG = "davidquarel/jaxgmg-69528c4f23d35d7de3eecfa1"
+JAXGMG_COLLECTION_SLUGS = (
+    "davidquarel/jaxgmg-69528c4f23d35d7de3eecfa1",
+    "timaeus/project-rl2-69bd40fed77a62a6eba28fc9",
+)
 DEFAULT_REPO = "davidquarel/jaxgmg_3phase_seed"
 LOCAL_BASE_PATH = Path(__file__).parent.parent
 USE_LOCAL = os.environ.get("JAXGMG_USE_LOCAL", "").lower() in ("1", "true", "yes")
@@ -28,14 +31,20 @@ def get_hf_api() -> HfApi:
 
 @lru_cache(maxsize=1)
 def get_collection_repos() -> tuple[str, ...]:
-    """Fetch all model repos from the jaxgmg collection."""
-    try:
-        collection = get_hf_api().get_collection(JAXGMG_COLLECTION_SLUG)
-        repos = [item.item_id for item in collection.items if item.item_type == "model"]
-        return tuple(repos) if repos else (DEFAULT_REPO,)
-    except Exception as e:
-        print(f"Warning: Could not fetch collection: {e}")
-        return (DEFAULT_REPO,)
+    """Fetch all model repos from all configured collections (deduplicated)."""
+    seen: set[str] = set()
+    repos: list[str] = []
+    api = get_hf_api()
+    for slug in JAXGMG_COLLECTION_SLUGS:
+        try:
+            collection = api.get_collection(slug)
+            for item in collection.items:
+                if item.item_type == "model" and item.item_id not in seen:
+                    seen.add(item.item_id)
+                    repos.append(item.item_id)
+        except Exception as e:
+            print(f"Warning: Could not fetch collection {slug}: {e}")
+    return tuple(repos) if repos else (DEFAULT_REPO,)
 
 
 def get_local_path(repo_id: str) -> Path:
@@ -366,7 +375,8 @@ def create_demo():
 
         gr.Markdown(
             f"<small>Source: {'Local' if USE_LOCAL else 'HuggingFace'} · "
-            f"[Collection](https://huggingface.co/collections/davidquarel/jaxgmg)</small>",
+            f"[davidquarel/jaxgmg](https://huggingface.co/collections/davidquarel/jaxgmg) · "
+            f"[timaeus/project-rl2](https://huggingface.co/collections/timaeus/project-rl2)</small>",
         )
 
         repo_dropdown.change(