Deploy Streamlit Space app

app.py

@@ -307,6 +307,18 @@ def _resolve_weight_source_for_model(model_name, requested_source):
     return "base", f"{short_name} has no '{requested_source}' weights. Using base."
 
 
+def _finetuned_available_for_model(model_name, requested_source):
+    if requested_source == "base":
+        return True
+    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)
+
+
 # ─────────────────────────────────────────────────────────────────────────────
 # Cached Model Loaders (with weight_source support)
 # ─────────────────────────────────────────────────────────────────────────────
@@ -760,29 +772,30 @@ with tab_caption:
 
     with col_result:
         if uploaded_file and generate_btn:
-            resolved_source, warning_msg = _resolve_weight_source_for_model(
-                selected_model, weight_source
-            )
-            if warning_msg:
-                st.warning(warning_msg)
-            with st.spinner(
-                f"Loading {MODEL_SHORT[selected_model]} ({resolved_source}) + generating…"
-            ):
-                try:
-                    caption = generate_caption(
-                        selected_model, selected_mode, image,
-                        num_beams=num_beams,
-                        max_new_tokens=max_new_tokens,
-                        length_penalty=length_penalty,
-                        weight_source=resolved_source,
-                    )
-                except Exception as e:
-                    st.error(f"Generation error: {e}")
-                    caption = None
+            if not _finetuned_available_for_model(selected_model, weight_source):
+                st.error(
+                    f"{MODEL_SHORT[selected_model]} does not have '{weight_source}' weights."
+                )
+                caption = None
+            else:
+                with st.spinner(
+                    f"Loading {MODEL_SHORT[selected_model]} ({weight_source}) + generating…"
+                ):
+                    try:
+                        caption = generate_caption(
+                            selected_model, selected_mode, image,
+                            num_beams=num_beams,
+                            max_new_tokens=max_new_tokens,
+                            length_penalty=length_penalty,
+                            weight_source=weight_source,
+                        )
+                    except Exception as e:
+                        st.error(f"Generation error: {e}")
+                        caption = None
 
             if caption:
                 render_caption_card(
-                    selected_model, caption, resolved_source,
+                    selected_model, caption, weight_source,
                     num_beams, length_penalty, max_new_tokens,
                     container=st,
                 )
@@ -844,16 +857,20 @@ with tab_compare:
         compare_image = Image.open(compare_file).convert("RGB")
 
         resolved_sources = {}
-        warnings = []
         for model_key in MODEL_KEYS:
-            resolved, warning_msg = _resolve_weight_source_for_model(
-                model_key, weight_source
-            )
-            resolved_sources[model_key] = resolved
-            if warning_msg:
-                warnings.append(warning_msg)
-        for msg in sorted(set(warnings)):
-            st.warning(msg)
+            resolved_sources[model_key] = weight_source
+        if weight_source != "base":
+            missing = [
+                MODEL_SHORT[m]
+                for m in MODEL_KEYS
+                if not _finetuned_available_for_model(m, weight_source)
+            ]
+            if missing:
+                st.warning(
+                    "Missing fine-tuned weights for: "
+                    + ", ".join(missing)
+                    + ". Marking those results as unavailable."
+                )
 
         # Generate captions from all 4 models
         results = {}
@@ -879,17 +896,24 @@ with tab_compare:
             else:
                 mode = "Baseline (Full Attention)"
 
-            try:
-                cap = generate_caption(
-                    model_key, mode, compare_image,
-                    num_beams=num_beams,
-                    max_new_tokens=max_new_tokens,
-                    length_penalty=length_penalty,
-                    weight_source=resolved_sources.get(model_key, weight_source),
+            if not _finetuned_available_for_model(model_key, weight_source):
+                results[model_key] = (
+                    f"[Fine-tuned '{weight_source}' weights not available]"
+                    if weight_source != "base"
+                    else "[Not available]"
                 )
-                results[model_key] = cap
-            except Exception as e:
-                results[model_key] = f"[Error: {e}]"
+            else:
+                try:
+                    cap = generate_caption(
+                        model_key, mode, compare_image,
+                        num_beams=num_beams,
+                        max_new_tokens=max_new_tokens,
+                        length_penalty=length_penalty,
+                        weight_source=weight_source,
+                    )
+                    results[model_key] = cap
+                except Exception as e:
+                    results[model_key] = f"[Error: {e}]"
 
         progress.progress(1.0, text="✅ All models complete!")