Load checkpoints from Hugging Face Hub

web/main.py

@@ -13,6 +13,7 @@ import torch
 from fastapi import FastAPI, File, Form, HTTPException, UploadFile
 from fastapi.responses import FileResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
+from huggingface_hub import snapshot_download
 from PIL import Image
 from peft import PeftModel
 from transformers import AutoTokenizer, LlavaForConditionalGeneration, LlavaProcessor
@@ -106,6 +107,17 @@ class VQAServerState:
         self.model_b_cfg = CFG.get("model_b", {})
         self.eval_cfg = CFG.get("eval", {})
         self.models_dir = ROOT_DIR / "checkpoints"
+        self.artifact_cache_dir = Path(
+            os.getenv("MEDVQA_ARTIFACT_CACHE", str(ROOT_DIR / ".cache" / "hub_artifacts"))
+        )
+        self.artifact_cache_dir.mkdir(parents=True, exist_ok=True)
+        self.hub_model_ids = {
+            "A1": os.getenv("MEDVQA_A1_MODEL_ID", "SpringWang08/medical-vqa-a1"),
+            "A2": os.getenv("MEDVQA_A2_MODEL_ID", "SpringWang08/medical-vqa-a2"),
+            "B2": os.getenv("MEDVQA_B2_MODEL_ID", "SpringWang08/medical-vqa-b2"),
+            "DPO": os.getenv("MEDVQA_DPO_MODEL_ID", "SpringWang08/medical-vqa-dpo"),
+            "PPO": os.getenv("MEDVQA_PPO_MODEL_ID", "SpringWang08/medical-vqa-ppo"),
+        }
         self.qa_tokenizer = None
         self.translator = MedicalTranslator(device="cpu")
         self.answer_rewriter = MedicalAnswerRewriter()
@@ -134,6 +146,19 @@ def _artifact_exists(path: Path) -> bool:
     return path.exists()
 
 
+def _download_hub_snapshot(repo_id: str, cache_subdir: str, allow_patterns: Optional[list[str]] = None) -> Path:
+    target_dir = state.artifact_cache_dir / cache_subdir
+    target_dir.mkdir(parents=True, exist_ok=True)
+    snapshot_download(
+        repo_id=repo_id,
+        repo_type="model",
+        local_dir=str(target_dir),
+        local_dir_use_symlinks=False,
+        allow_patterns=allow_patterns,
+    )
+    return target_dir
+
+
 def _as_bool(value: Any) -> bool:
     if isinstance(value, bool):
         return value
@@ -352,7 +377,20 @@ def _resolve_variant_artifact(variant: str) -> dict[str, Any]:
         ckpt_path = ROOT_DIR / "checkpoints" / f"medical_vqa_{variant}_best.pth"
         if not ckpt_path.exists():
             resume_path = ROOT_DIR / "checkpoints" / f"medical_vqa_{variant}_resume.pth"
-            ckpt_path = resume_path if resume_path.exists() else ckpt_path
+            if resume_path.exists():
+                ckpt_path = resume_path
+            else:
+                repo_id = state.hub_model_ids.get(variant, "")
+                if repo_id:
+                    downloaded_dir = _download_hub_snapshot(
+                        repo_id=repo_id,
+                        cache_subdir=variant.lower(),
+                        allow_patterns=["README.md", "*.pth"],
+                    )
+                    downloaded_ckpt = downloaded_dir / f"medical_vqa_{variant}_best.pth"
+                    if not downloaded_ckpt.exists():
+                        downloaded_ckpt = downloaded_dir / f"medical_vqa_{variant}_resume.pth"
+                    ckpt_path = downloaded_ckpt
         return {"type": "direction_a", "path": ckpt_path}
 
     if variant == "B1":
@@ -360,15 +398,49 @@ def _resolve_variant_artifact(variant: str) -> dict[str, Any]:
 
     if variant == "B2":
         ckpt_dir = _select_best_b2_checkpoint(ROOT_DIR / "checkpoints" / "B2")
+        if ckpt_dir is None:
+            repo_id = state.hub_model_ids.get("B2", "")
+            if repo_id:
+                ckpt_dir = _download_hub_snapshot(
+                    repo_id=repo_id,
+                    cache_subdir="b2",
+                    allow_patterns=["README.md", "adapter_model.safetensors", "adapter_config.json", "tokenizer.json", "tokenizer_config.json", "processor_config.json", "chat_template.jinja"],
+                )
         return {"type": "llava_adapter", "path": ckpt_dir}
 
     if variant == "DPO":
         final_adapter = ROOT_DIR / "checkpoints" / "DPO" / "final_adapter"
         fallback = ROOT_DIR / "checkpoints" / "DPO" / "checkpoint-25"
-        return {"type": "llava_adapter", "path": final_adapter if final_adapter.exists() else fallback}
+        if final_adapter.exists():
+            return {"type": "llava_adapter", "path": final_adapter}
+        if fallback.exists():
+            return {"type": "llava_adapter", "path": fallback}
+        repo_id = state.hub_model_ids.get("DPO", "")
+        if repo_id:
+            return {
+                "type": "llava_adapter",
+                "path": _download_hub_snapshot(
+                    repo_id=repo_id,
+                    cache_subdir="dpo",
+                    allow_patterns=["README.md", "adapter_model.safetensors", "adapter_config.json", "tokenizer.json", "tokenizer_config.json", "processor_config.json", "chat_template.jinja"],
+                ),
+            }
+        return {"type": "llava_adapter", "path": final_adapter}
 
     if variant == "PPO":
         final_adapter = ROOT_DIR / "checkpoints" / "PPO" / "final_adapter"
+        if final_adapter.exists():
+            return {"type": "llava_adapter", "path": final_adapter}
+        repo_id = state.hub_model_ids.get("PPO", "")
+        if repo_id:
+            return {
+                "type": "llava_adapter",
+                "path": _download_hub_snapshot(
+                    repo_id=repo_id,
+                    cache_subdir="ppo",
+                    allow_patterns=["README.md", "adapter_model.safetensors", "adapter_config.json", "tokenizer.json", "tokenizer_config.json", "processor_config.json", "chat_template.jinja"],
+                ),
+            }
         return {"type": "llava_adapter", "path": final_adapter}
 
     raise ValueError(f"Unknown variant: {variant}")
@@ -857,12 +929,12 @@ 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")), "artifact": "checkpoints/medical_vqa_A1_best.pth"},
-        "A2": {"available": (_artifact_exists(ROOT_DIR / "checkpoints" / "medical_vqa_A2_best.pth")), "artifact": "checkpoints/medical_vqa_A2_best.pth"},
+        "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, "artifact": str(b2_checkpoint) if b2_checkpoint else ""},
-        "DPO": {"available": cuda_ready and (_artifact_exists(ROOT_DIR / "checkpoints" / "DPO" / "final_adapter") or _artifact_exists(ROOT_DIR / "checkpoints" / "DPO" / "checkpoint-25")), "artifact": "checkpoints/DPO/final_adapter"},
-        "PPO": {"available": cuda_ready and _artifact_exists(ROOT_DIR / "checkpoints" / "PPO" / "final_adapter"), "artifact": "checkpoints/PPO/final_adapter"},
+        "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", "")},
     }