Lock Space to B2-only demo

Dockerfile

@@ -7,6 +7,7 @@ ENV DEBIAN_FRONTEND=noninteractive \
     HF_HOME=/data/.huggingface \
     HUGGINGFACE_HUB_CACHE=/data/.huggingface/hub \
     TRANSFORMERS_CACHE=/data/.huggingface/transformers \
+    MEDVQA_ACTIVE_VARIANTS=B2 \
     WEB_PRELOAD_MODELS=0 \
     ANSWER_REWRITE_ENABLED=0
 

web/README.md

@@ -5,7 +5,8 @@ Thư mục này chứa FastAPI + web UI để:
 - upload ảnh
 - nhập câu hỏi VQA
 - chạy dự đoán
-- so sánh 6 model: `A1`, `A2`, `B1`, `B2`, `DPO`, `PPO`
+- chạy mặc định model `B2` trên Hugging Face Space
+- nếu cần, vẫn có thể bật lại các model khác bằng biến môi trường
 
 ### Chạy server
 
@@ -23,6 +24,14 @@ WEB_PRELOAD_MODELS=1 uvicorn web.main:app --host 0.0.0.0 --port 8000
 
 Mặc định hiện tại là `WEB_PRELOAD_MODELS=0` để Space khởi động nhẹ hơn. Chỉ bật `1` khi GPU đủ mạnh và bạn muốn preload trước.
 
+Mặc định Space chỉ mở chế độ `B2` để giảm RAM/VRAM:
+
+```bash
+MEDVQA_ACTIVE_VARIANTS=B2
+```
+
+Nếu muốn chạy nhiều model hơn, đặt `MEDVQA_ACTIVE_VARIANTS` thành danh sách ngăn cách bởi dấu phẩy, ví dụ `A1,A2,B2`.
+
 Khi chạy trên GPU, nên để `--workers 1` để tránh mỗi worker nạp một bản model riêng.
 
 ### Chạy bằng Docker
@@ -78,8 +87,8 @@ http://localhost:8000
   - form-data:
     - `question`: câu hỏi VQA
     - `image`: ảnh đầu vào
-    - `model_name` hoặc `model_names`:
-      - nếu bỏ trống thì chạy toàn bộ 6 model
+- `model_name` hoặc `model_names`:
+      - nếu bỏ trống thì chạy các model đang bật trong `MEDVQA_ACTIVE_VARIANTS`
       - `model_names` nhận chuỗi JSON list hoặc chuỗi phân tách bằng dấu phẩy
 
 ### Artifact cần có

web/main.py

@@ -133,6 +133,12 @@ class VQAServerState:
         self.preload_models = os.getenv("WEB_PRELOAD_MODELS", "0") == "1"
         # Chạy lần lượt và giải phóng model sau mỗi lượt để giảm đỉnh RAM/VRAM.
         self.release_after_predict = os.getenv("WEB_RELEASE_AFTER_PREDICT", "1") == "1"
+        raw_active_variants = os.getenv("MEDVQA_ACTIVE_VARIANTS", "B2")
+        self.active_variants = {
+            variant.strip()
+            for variant in raw_active_variants.split(",")
+            if variant.strip() in VARIANT_ORDER
+        } or {"B2"}
         self.progress_state: dict[str, Any] = {
             "job_id": "",
             "active": False,
@@ -507,6 +513,8 @@ def _resolve_variant_artifact(variant: str) -> dict[str, Any]:
 def _llava_adapter_specs() -> list[tuple[str, Path]]:
     specs: list[tuple[str, Path]] = []
     for variant in ("B2", "DPO", "PPO"):
+        if variant not in state.active_variants:
+            continue
         artifact = _resolve_variant_artifact(variant)["path"]
         if isinstance(artifact, Path) and artifact.exists():
             specs.append((variant, artifact))
@@ -1051,26 +1059,26 @@ def _parse_model_selection(raw_model_name: Optional[str], raw_model_names: Optio
             parsed = [part.strip() for part in raw_model_names.split(",") if part.strip()]
         if isinstance(parsed, str):
             parsed = [parsed]
-        selected = [name for name in parsed if name in VARIANT_ORDER]
+        selected = [name for name in parsed if name in VARIANT_ORDER and name in state.active_variants]
         if selected:
             return selected
 
-    if raw_model_name and raw_model_name in VARIANT_ORDER:
+    if raw_model_name and raw_model_name in VARIANT_ORDER and raw_model_name in state.active_variants:
         return [raw_model_name]
 
-    return VARIANT_ORDER[:]
+    return [variant for variant in VARIANT_ORDER if variant in state.active_variants]
 
 
 def _variant_availability() -> dict[str, dict[str, Any]]:
     b2_checkpoint = _select_best_b2_checkpoint(ROOT_DIR / "checkpoints" / "B2")
     cuda_ready = torch.cuda.is_available()
     return {
-        "A1": {"available": (_artifact_exists(ROOT_DIR / "checkpoints" / "medical_vqa_A1_best.pth") or bool(state.hub_model_ids.get("A1"))), "artifact": str(ROOT_DIR / "checkpoints" / "medical_vqa_A1_best.pth") if _artifact_exists(ROOT_DIR / "checkpoints" / "medical_vqa_A1_best.pth") else state.hub_model_ids.get("A1", "")},
-        "A2": {"available": (_artifact_exists(ROOT_DIR / "checkpoints" / "medical_vqa_A2_best.pth") or bool(state.hub_model_ids.get("A2"))), "artifact": str(ROOT_DIR / "checkpoints" / "medical_vqa_A2_best.pth") if _artifact_exists(ROOT_DIR / "checkpoints" / "medical_vqa_A2_best.pth") else state.hub_model_ids.get("A2", "")},
-        "B1": {"available": cuda_ready, "artifact": state.llava_model_id},
-        "B2": {"available": cuda_ready and (b2_checkpoint is not None or bool(state.hub_model_ids.get("B2"))), "artifact": str(b2_checkpoint) if b2_checkpoint else state.hub_model_ids.get("B2", "")},
-        "DPO": {"available": cuda_ready and (_artifact_exists(ROOT_DIR / "checkpoints" / "DPO" / "final_adapter") or _artifact_exists(ROOT_DIR / "checkpoints" / "DPO" / "checkpoint-25") or bool(state.hub_model_ids.get("DPO"))), "artifact": "checkpoints/DPO/final_adapter" if _artifact_exists(ROOT_DIR / "checkpoints" / "DPO" / "final_adapter") else state.hub_model_ids.get("DPO", "")},
-        "PPO": {"available": cuda_ready and (_artifact_exists(ROOT_DIR / "checkpoints" / "PPO" / "final_adapter") or bool(state.hub_model_ids.get("PPO"))), "artifact": "checkpoints/PPO/final_adapter" if _artifact_exists(ROOT_DIR / "checkpoints" / "PPO" / "final_adapter") else state.hub_model_ids.get("PPO", "")},
+        "A1": {"available": ("A1" in state.active_variants) and (_artifact_exists(ROOT_DIR / "checkpoints" / "medical_vqa_A1_best.pth") or bool(state.hub_model_ids.get("A1"))), "artifact": str(ROOT_DIR / "checkpoints" / "medical_vqa_A1_best.pth") if _artifact_exists(ROOT_DIR / "checkpoints" / "medical_vqa_A1_best.pth") else state.hub_model_ids.get("A1", "")},
+        "A2": {"available": ("A2" in state.active_variants) and (_artifact_exists(ROOT_DIR / "checkpoints" / "medical_vqa_A2_best.pth") or bool(state.hub_model_ids.get("A2"))), "artifact": str(ROOT_DIR / "checkpoints" / "medical_vqa_A2_best.pth") if _artifact_exists(ROOT_DIR / "checkpoints" / "medical_vqa_A2_best.pth") else state.hub_model_ids.get("A2", "")},
+        "B1": {"available": ("B1" in state.active_variants) and cuda_ready, "artifact": state.llava_model_id},
+        "B2": {"available": ("B2" in state.active_variants) and cuda_ready and (b2_checkpoint is not None or bool(state.hub_model_ids.get("B2"))), "artifact": str(b2_checkpoint) if b2_checkpoint else state.hub_model_ids.get("B2", "")},
+        "DPO": {"available": ("DPO" in state.active_variants) and cuda_ready and (_artifact_exists(ROOT_DIR / "checkpoints" / "DPO" / "final_adapter") or _artifact_exists(ROOT_DIR / "checkpoints" / "DPO" / "checkpoint-25") or bool(state.hub_model_ids.get("DPO"))), "artifact": "checkpoints/DPO/final_adapter" if _artifact_exists(ROOT_DIR / "checkpoints" / "DPO" / "final_adapter") else state.hub_model_ids.get("DPO", "")},
+        "PPO": {"available": ("PPO" in state.active_variants) and cuda_ready and (_artifact_exists(ROOT_DIR / "checkpoints" / "PPO" / "final_adapter") or bool(state.hub_model_ids.get("PPO"))), "artifact": "checkpoints/PPO/final_adapter" if _artifact_exists(ROOT_DIR / "checkpoints" / "PPO" / "final_adapter") else state.hub_model_ids.get("PPO", "")},
     }
 
 

web/static/index.html

@@ -177,7 +177,7 @@ X2 Vision
 <div class="flex flex-col items-center text-center max-w-4xl mx-auto mb-14">
 <div class="mb-4 flex items-center gap-2">
 <div class="h-[1px] w-12 bg-china-gold"></div>
-<span class="text-china-gold font-display text-sm tracking-[0.2em] uppercase">6-model comparison</span>
+<span class="text-china-gold font-display text-sm tracking-[0.2em] uppercase">B2-only comparison</span>
 <div class="h-[1px] w-12 bg-china-gold"></div>
 </div>
 <h1 class="text-imperial-red text-[42px] md:text-[64px] font-display font-bold leading-[1.1] tracking-tight mb-6 drop-shadow-sm">
@@ -298,7 +298,7 @@ Reset
 <span class="material-symbols-outlined absolute right-6 text-[28px] opacity-20 group-hover:opacity-40 transition-opacity text-gold-light">chess_knight</span>
 </button>
 
-<div class="text-center text-sm font-serif italic text-ink-black/60" id="status-text">Select an image, enter a question, then run all six models.</div>
+<div class="text-center text-sm font-serif italic text-ink-black/60" id="status-text">Select an image, enter a question, then run B2.</div>
 </div>
 </div>
 </div>
@@ -359,7 +359,7 @@ Alignment and RL variants now have equal room in the grid, making the comparison
 <span class="font-display font-bold text-lg tracking-wider">VQA RESEARCH</span>
 </div>
 <div class="text-[13px] text-paper-white/60 font-serif">
-Medical VQA web demo for six-model comparison.
+Medical VQA web demo for B2-only inference.
 </div>
 </div>
 <div class="flex gap-8 text-[13px] text-paper-white/80 font-display tracking-widest uppercase">
@@ -399,9 +399,10 @@ Medical VQA web demo for six-model comparison.
   };
 
   let currentImageFile = null;
-  let selectedModels = new Set(MODEL_ORDER);
+  let selectedModels = new Set(["B2"]);
   let questionSuggestions = [];
   let progressTimer = null;
+  let modelAvailability = {};
 
   function escapeHtml(value) {
     return String(value ?? "")
@@ -597,10 +598,18 @@ Medical VQA web demo for six-model comparison.
   function updateModelChips() {
     document.querySelectorAll(".model-chip").forEach((chip) => {
       const variant = chip.dataset.model;
+      const available = modelAvailability[variant] !== false;
       const active = selectedModels.has(variant);
+      chip.disabled = !available;
+      chip.style.opacity = available ? "1" : "0.35";
+      chip.style.cursor = available ? "pointer" : "not-allowed";
       chip.style.background = active ? "#A8181B" : "#fff";
       chip.style.color = active ? "#FDFBF7" : "#1A1A1A";
       chip.style.borderColor = active ? "#A8181B" : "rgba(212,175,55,0.35)";
+      if (!available && !active) {
+        chip.style.background = "#faf7f0";
+        chip.style.color = "rgba(26,26,26,0.45)";
+      }
     });
   }
 
@@ -626,8 +635,14 @@ Medical VQA web demo for six-model comparison.
     try {
       const res = await fetch("/v1/models");
       const data = await res.json();
+      modelAvailability = Object.fromEntries((data.models || []).map((item) => [item.name, Boolean(item.available)]));
+      if (!modelAvailability.B2) {
+        selectedModels = new Set();
+      } else if (!selectedModels.has("B2")) {
+        selectedModels = new Set(["B2"]);
+      }
       updateModelChips();
-      setStatus("Ready. Upload an image and run all six models.");
+      setStatus("Ready. Upload an image and run B2.");
     } catch (err) {
       setStatus(`Failed to load model metadata: ${err.message}`);
     }
@@ -666,17 +681,20 @@ Medical VQA web demo for six-model comparison.
   document.querySelectorAll(".model-chip").forEach((chip) => {
     chip.addEventListener("click", () => {
       const variant = chip.dataset.model;
+      if (modelAvailability[variant] === false) {
+        return;
+      }
       if (selectedModels.has(variant)) selectedModels.delete(variant);
-      else selectedModels.add(variant);
+      else selectedModels = new Set([variant]);
       if (selectedModels.size === 0) {
-        selectedModels = new Set(MODEL_ORDER);
+        selectedModels = new Set(["B2"]);
       }
       updateModelChips();
     });
   });
 
   el.resetBtn.addEventListener("click", () => {
-    selectedModels = new Set(MODEL_ORDER);
+    selectedModels = new Set(["B2"]);
     el.question.value = "";
     el.imageInput.value = "";
     setPreview(null);
@@ -696,13 +714,13 @@ Medical VQA web demo for six-model comparison.
       return;
     }
     if (selectedModels.size === 0) {
-      setStatus("Please select at least one model.");
+      setStatus("Please select B2.");
       return;
     }
 
     el.runBtn.disabled = true;
     el.runBtn.querySelector("span").textContent = "Running...";
-    setStatus("Running all selected models...");
+    setStatus("Running B2...");
     renderRunningModelGrid();
     applyTiltEffect(".tilt-card", 5);
     startProgressPolling();
@@ -730,7 +748,7 @@ Medical VQA web demo for six-model comparison.
 
       renderModelGrid(resultData?.payload?.results || []);
       applyTiltEffect(".tilt-card", 5);
-      setStatus(`Done. ${resultData?.payload?.summary?.success_count ?? 0} models succeeded.`);
+      setStatus(`Done. B2 succeeded.`);
     } catch (err) {
       setStatus(err.message || "Prediction failed");
     } finally {