Deploy Streamlit Space app

app.py

@@ -190,11 +190,21 @@ OUTPUT_ROOT = DEFAULT_OUTPUT_ROOT
 
 
 @st.cache_resource(show_spinner=False)
-def _download_weights(need_outputs: bool, need_shakespeare: bool) -> str:
+def _download_weights(
+    need_outputs: bool,
+    need_shakespeare: bool,
+    output_model_dir: str | None = None,
+) -> str:
     from huggingface_hub import snapshot_download
     allow_patterns = []
     if need_outputs:
-        allow_patterns += ["outputs/*", "outputs/**/*"]
+        if output_model_dir:
+            allow_patterns += [
+                f"outputs/{output_model_dir}/*",
+                f"outputs/{output_model_dir}/**/*",
+            ]
+        else:
+            allow_patterns += ["outputs/*", "outputs/**/*"]
     if need_shakespeare:
         allow_patterns += ["input.txt", "shakespeare_transformer.pt"]
     if not allow_patterns:
@@ -210,39 +220,60 @@ def _download_weights(need_outputs: bool, need_shakespeare: bool) -> str:
 
 @st.cache_resource(show_spinner=False)
 def _download_model_outputs(model_dir: str) -> str:
-    from huggingface_hub import snapshot_download
-    allow_patterns = [
-        f"outputs/{model_dir}/*",
-        f"outputs/{model_dir}/**/*",
-    ]
-    return snapshot_download(
-        repo_id=WEIGHTS_REPO_ID,
-        repo_type="model",
-        local_dir=WEIGHTS_CACHE_DIR,
-        local_dir_use_symlinks=False,
-        allow_patterns=allow_patterns,
+    return _download_weights(
+        need_outputs=True,
+        need_shakespeare=False,
+        output_model_dir=model_dir,
     )
 
 
 def _ensure_model_outputs_available(model_dir: str) -> None:
-    if not model_dir:
-        return
-    local = os.path.isdir(os.path.join(DEFAULT_OUTPUT_ROOT, model_dir))
-    cached = os.path.isdir(os.path.join(WEIGHTS_CACHE_DIR, "outputs", model_dir))
-    if local or cached:
-        return
+    # Intentionally no eager snapshot download here.
+    # We only fetch checkpoints when a user explicitly selects a fine-tuned weight.
+    _ = model_dir
+
+
+@st.cache_data(show_spinner=False, ttl=900)
+def _weights_repo_file_set() -> set[str]:
+    from huggingface_hub import HfApi
+    api = HfApi()
     try:
-        _download_model_outputs(model_dir)
-    except Exception as e:
-        print(f"⚠️ Could not prefetch outputs for {model_dir}: {e}")
+        files = api.list_repo_files(repo_id=WEIGHTS_REPO_ID, repo_type="model")
+        return set(files)
+    except Exception:
+        return set()
 
 
-def _resolve_weight_paths(need_outputs: bool, need_shakespeare: bool):
+def _remote_has_finetuned(model_dir: str, subdir: str) -> bool:
+    files = _weights_repo_file_set()
+    if not files:
+        return False
+    prefix = f"outputs/{model_dir}/{subdir}/"
+    return any(path.startswith(prefix) for path in files)
+
+
+def _resolve_weight_paths(
+    need_outputs: bool,
+    need_shakespeare: bool,
+    output_model_dir: str | None = None,
+):
     output_root = DEFAULT_OUTPUT_ROOT
     shakespeare_file = DEFAULT_SHAKESPEARE_FILE
     shakespeare_weights = DEFAULT_SHAKESPEARE_WEIGHTS
 
-    have_outputs = os.path.isdir(output_root) and len(os.listdir(output_root)) > 0
+    if need_outputs:
+        if output_model_dir:
+            local_model = os.path.join(output_root, output_model_dir)
+            cache_model = os.path.join(WEIGHTS_CACHE_DIR, "outputs", output_model_dir)
+            have_outputs = (
+                (os.path.isdir(local_model) and len(os.listdir(local_model)) > 0)
+                or (os.path.isdir(cache_model) and len(os.listdir(cache_model)) > 0)
+            )
+        else:
+            have_outputs = os.path.isdir(output_root) and len(os.listdir(output_root)) > 0
+    else:
+        have_outputs = True
+
     have_shakespeare = (
         os.path.exists(shakespeare_file) and os.path.exists(shakespeare_weights)
     )
@@ -250,7 +281,11 @@ def _resolve_weight_paths(need_outputs: bool, need_shakespeare: bool):
         return output_root, shakespeare_file, shakespeare_weights
 
     try:
-        cache_dir = _download_weights(need_outputs, need_shakespeare)
+        cache_dir = _download_weights(
+            need_outputs,
+            need_shakespeare,
+            output_model_dir=output_model_dir,
+        )
         candidate_output_root = os.path.join(cache_dir, "outputs")
         candidate_shakespeare_file = os.path.join(cache_dir, "input.txt")
         candidate_shakespeare_weights = os.path.join(
@@ -291,7 +326,7 @@ def _has_finetuned(model_dir, subdir):
     for path in candidates:
         if os.path.isdir(path) and len(os.listdir(path)) > 0:
             return True
-    return False
+    return _remote_has_finetuned(model_dir, subdir)
 
 
 def _ckpt_path(output_root, model_dir, subdir):
@@ -312,6 +347,7 @@ def _resolve_weight_source_for_model(model_name, requested_source):
     _resolve_weight_paths(
         need_outputs=True,
         need_shakespeare=(model_dir == "custom_vlm"),
+        output_model_dir=model_dir,
     )
     if _has_finetuned(model_dir, requested_source):
         return requested_source, None
@@ -325,9 +361,6 @@ def _finetuned_available_for_model(model_name, requested_source):
     model_dir = MODEL_DIR.get(model_name)
     if not model_dir or model_dir in DISABLE_FINETUNE_FOR:
         return False
-    if _has_finetuned(model_dir, requested_source):
-        return True
-    _ensure_model_outputs_available(model_dir)
     return _has_finetuned(model_dir, requested_source)
 
 
@@ -346,7 +379,9 @@ def load_blip(weight_source="base"):
 
     if weight_source != "base":
         output_root, _, _ = _resolve_weight_paths(
-            need_outputs=True, need_shakespeare=False
+            need_outputs=True,
+            need_shakespeare=False,
+            output_model_dir="blip",
         )
         ckpt = _ckpt_path(output_root, "blip", weight_source)
         if os.path.isdir(ckpt) and os.listdir(ckpt):
@@ -375,7 +410,9 @@ def load_vit_gpt2(weight_source="base"):
 
     if weight_source != "base":
         output_root, _, _ = _resolve_weight_paths(
-            need_outputs=True, need_shakespeare=False
+            need_outputs=True,
+            need_shakespeare=False,
+            output_model_dir="vit_gpt2",
         )
         ckpt = _ckpt_path(output_root, "vit_gpt2", weight_source)
         if os.path.isdir(ckpt) and os.listdir(ckpt):
@@ -400,7 +437,9 @@ def load_git(weight_source="base"):
 
     if weight_source != "base":
         output_root, _, _ = _resolve_weight_paths(
-            need_outputs=True, need_shakespeare=False
+            need_outputs=True,
+            need_shakespeare=False,
+            output_model_dir="git",
         )
         ckpt = _ckpt_path(output_root, "git", weight_source)
         if os.path.isdir(ckpt) and os.listdir(ckpt):
@@ -422,7 +461,9 @@ def load_custom_vlm(weight_source="base"):
     device = get_device()
     cfg = CFG()
     output_root, shakespeare_file, shakespeare_weights = _resolve_weight_paths(
-        need_outputs=(weight_source != "base"), need_shakespeare=True
+        need_outputs=(weight_source != "base"),
+        need_shakespeare=True,
+        output_model_dir="custom_vlm" if weight_source != "base" else None,
     )
     cfg.output_root = output_root
     cfg.shakespeare_file = shakespeare_file
@@ -501,7 +542,9 @@ def load_blip_attention_model(weight_source="base"):
 
     if weight_source != "base":
         output_root, _, _ = _resolve_weight_paths(
-            need_outputs=True, need_shakespeare=False
+            need_outputs=True,
+            need_shakespeare=False,
+            output_model_dir="blip",
         )
         ckpt = _ckpt_path(output_root, "blip", weight_source)
         if os.path.isdir(ckpt) and os.listdir(ckpt):