Update app.py

app.py

@@ -1,9 +1,7 @@
 #!/usr/bin/env python3
 """gQIR Gradio demo
-
 Single-frame mode follows infer_sd2GAN_stage2.py (color path only).
 Burst mode follows infer_burst_realistic.py for 77->11 aggregation and reconstruction.
-
 Local run cmd:
 python gradio_app.py 
   --single-config configs/inference/eval_sd2GAN.yaml \
@@ -49,7 +47,13 @@ try:
 except Exception:
     h5py = None
 
+try:
+    from huggingface_hub import hf_hub_download
+except Exception:
+    hf_hub_download = None
+
 
+APP_ROOT = Path(__file__).resolve().parent
 DEFAULT_SINGLE_CONFIG_COLOR =  "configs/inference/eval_3bit_color.yaml"
 DEFAULT_SINGLE_CONFIG_MONO  =  "configs/inference/eval_3bit_mono.yaml"
 DEFAULT_BURST_CONFIG_COLOR  =  "configs/inference/eval_burst_mosaic.yaml"
@@ -60,6 +64,23 @@ BURST_WINDOW = 77
 PIPELINE_COLOR = "Color"
 PIPELINE_MONO = "Monochrome"
 PIPELINE_OPTIONS = [PIPELINE_COLOR, PIPELINE_MONO]
+HF_DEFAULT_REPO_ID = "aRy4n/gQIR"
+HF_MODEL_FILES = {
+    PIPELINE_COLOR: {
+        "single_qvae": "0105000.pt",
+        "single_lora": "state_dict.pth",
+        "burst_qvae": "0105000.pt",
+        "burst_lora": "state_dict.pth",
+        "burst_fusion": "fusion_vit_0050000.pt",
+    },
+    PIPELINE_MONO: {
+        "single_qvae": "mono/0150000.pt",
+        "single_lora": "mono/state_dict.pth",
+        "burst_qvae": "mono/0150000.pt",
+        "burst_lora": "mono/state_dict.pth",
+        "burst_fusion": "mono/fusion_vit_0020000.pt",
+    },
+}
 
 SINGLE_MODE_GT = "GT image (simulate 3-bit SPAD)"
 SINGLE_MODE_REAL = "Real SPAD frame"
@@ -77,6 +98,9 @@ RUNTIME_SINGLE_CONFIGS: dict[str, Path] = {}
 RUNTIME_BURST_CONFIGS: dict[str, Path] = {}
 RUNTIME_DEVICE: str = "cuda"
 RUNTIME_BURST_OUT_SIZES: dict[str, int] = {}
+RUNTIME_HF_REPO_ID: str = HF_DEFAULT_REPO_ID
+RUNTIME_HF_CACHE_DIR: Optional[str] = None
+RUNTIME_HF_TOKEN: Optional[str] = None
 _TEMP_VIDEO_DIRS: list[str] = []
 
 
@@ -129,12 +153,114 @@ def parse_args() -> argparse.Namespace:
         default="cuda" if torch.cuda.is_available() else "cpu",
         help="Inference device, e.g. cuda, cuda:0, cpu",
     )
+    parser.add_argument(
+        "--hf-repo-id",
+        type=str,
+        default=HF_DEFAULT_REPO_ID,
+        help="Hugging Face repo containing gQIR checkpoints used when config paths are not local files.",
+    )
+    parser.add_argument(
+        "--hf-cache-dir",
+        type=str,
+        default=None,
+        help="Optional Hugging Face cache directory for checkpoint downloads.",
+    )
+    parser.add_argument(
+        "--hf-token",
+        type=str,
+        default=None,
+        help="Optional HF token. If omitted, app reads HF_TOKEN or HUGGINGFACE_HUB_TOKEN env vars.",
+    )
     parser.add_argument("--port", type=int, default=7860)
     parser.add_argument("--local", action="store_true", help="Bind to 127.0.0.1 instead of 0.0.0.0")
     parser.add_argument("--share", action="store_true")
     return parser.parse_args()
 
 
+def _resolve_existing_file(path_value: Optional[str]) -> Optional[str]:
+    if not path_value:
+        return None
+    raw = Path(str(path_value)).expanduser()
+    candidates = [raw]
+    if not raw.is_absolute():
+        candidates.append(APP_ROOT / raw)
+    for p in candidates:
+        if p.is_file():
+            return str(p.resolve())
+    return None
+
+
+def _download_hf_checkpoint(filename: str) -> str:
+    if hf_hub_download is None:
+        raise RuntimeError(
+            "huggingface_hub is required to download checkpoints. Install it or provide local model paths."
+        )
+
+    kwargs: dict[str, Any] = {
+        "repo_id": RUNTIME_HF_REPO_ID,
+        "filename": filename,
+    }
+    if RUNTIME_HF_CACHE_DIR:
+        kwargs["cache_dir"] = RUNTIME_HF_CACHE_DIR
+    if RUNTIME_HF_TOKEN:
+        kwargs["token"] = RUNTIME_HF_TOKEN
+
+    try:
+        downloaded = hf_hub_download(**kwargs)
+    except Exception as exc:
+        raise RuntimeError(
+            f"Failed to download '{filename}' from '{RUNTIME_HF_REPO_ID}'. "
+            "Check repo, token permissions, and network availability."
+        ) from exc
+    return str(Path(downloaded).resolve())
+
+
+def _resolve_checkpoint_path(config_value: Optional[str], pipeline_type: str, file_key: str) -> str:
+    if pipeline_type not in HF_MODEL_FILES:
+        raise ValueError(f"Unknown pipeline type for checkpoint resolution: {pipeline_type}")
+    if file_key not in HF_MODEL_FILES[pipeline_type]:
+        raise ValueError(f"Unknown checkpoint key '{file_key}' for pipeline type '{pipeline_type}'")
+
+    existing = _resolve_existing_file(config_value)
+    if existing is not None:
+        return existing
+
+    hf_file = HF_MODEL_FILES[pipeline_type][file_key]
+    print(f"[gQIR] Missing local checkpoint; downloading {RUNTIME_HF_REPO_ID}/{hf_file}")
+    return _download_hf_checkpoint(hf_file)
+
+
+def _prepare_single_cfg_paths(cfg: Any, pipeline_type: str) -> Any:
+    cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=False))
+    if "model" not in cfg or "vae_cfg" not in cfg.model:
+        raise ValueError("Single-frame config missing model.vae_cfg")
+
+    qvae_path_cfg = None
+    if "qvae_path" in cfg.model.vae_cfg:
+        qvae_path_cfg = cfg.model.vae_cfg.qvae_path
+    if not qvae_path_cfg and "qvae_path" in cfg:
+        qvae_path_cfg = cfg.qvae_path
+
+    cfg.weight_path = _resolve_checkpoint_path(cfg.get("weight_path"), pipeline_type, "single_lora")
+    resolved_qvae = _resolve_checkpoint_path(qvae_path_cfg, pipeline_type, "single_qvae")
+    cfg.model.vae_cfg.qvae_path = resolved_qvae
+    if "qvae_path" in cfg:
+        cfg.qvae_path = resolved_qvae
+    return cfg
+
+
+def _prepare_burst_cfg_paths(cfg: Any, pipeline_type: str) -> Any:
+    cfg = OmegaConf.create(OmegaConf.to_container(cfg, resolve=False))
+    cfg.qvae_path = _resolve_checkpoint_path(cfg.get("qvae_path"), pipeline_type, "burst_qvae")
+    cfg.unet_weight_path = _resolve_checkpoint_path(cfg.get("unet_weight_path"), pipeline_type, "burst_lora")
+    cfg.fusion_vit_weight_path = _resolve_checkpoint_path(
+        cfg.get("fusion_vit_weight_path"), pipeline_type, "burst_fusion"
+    )
+    if "model" in cfg and "vae_cfg" in cfg.model:
+        cfg.model.vae_cfg.qvae_path = cfg.qvae_path
+    return cfg
+
+
 def _ensure_rgb_image(arr: np.ndarray) -> np.ndarray:
     arr = np.asarray(arr)
     if arr.ndim == 2:
@@ -378,9 +504,10 @@ def differentiable_warp(x: torch.Tensor, flow: torch.Tensor) -> torch.Tensor:
 
 
 class SingleColorPipeline:
-    def __init__(self, config_path: Path, device: str):
+    def __init__(self, config_path: Path, device: str, pipeline_type: str):
         self.config_path = config_path
         self.device = device
+        self.pipeline_type = pipeline_type
         self.max_size = DEFAULT_MAX_SIZE
         self.model: Optional[SD2Enhancer] = None
 
@@ -388,6 +515,7 @@ class SingleColorPipeline:
         if self.model is not None:
             return
         cfg = OmegaConf.load(str(self.config_path))
+        cfg = _prepare_single_cfg_paths(cfg, self.pipeline_type)
         if cfg.base_model_type != "sd2":
             raise ValueError(f"Unsupported base_model_type for single pipeline: {cfg.base_model_type}")
         self.model = SD2Enhancer(
@@ -462,9 +590,10 @@ class SingleColorPipeline:
 
 
 class BurstColorPipeline:
-    def __init__(self, config_path: Path, device: str):
+    def __init__(self, config_path: Path, device: str, pipeline_type: str):
         self.config_path = config_path
         self.device = device
+        self.pipeline_type = pipeline_type
         self.cfg = None
         self.out_size = 512
         self.weight_dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32
@@ -482,6 +611,7 @@ class BurstColorPipeline:
             return
 
         cfg = OmegaConf.load(str(self.config_path))
+        cfg = _prepare_burst_cfg_paths(cfg, self.pipeline_type)
         self.cfg = cfg
         self.out_size = int(cfg.dataset.val.params.out_size)
 
@@ -504,7 +634,7 @@ class BurstColorPipeline:
             dropout = False
 
         raft_model = RAFT(RAFTArgs())
-        raft_path =  "pretrained_ckpts/models/raft-things.pth"
+        raft_path = APP_ROOT / "pretrained_ckpts" / "models" / "raft-things.pth"
         raft_dict = torch.load(str(raft_path), map_location="cpu")
         corrected = {}
         for k, v in raft_dict.items():
@@ -687,6 +817,7 @@ def _get_single_pipeline(pipeline_type: str) -> SingleColorPipeline:
             _SINGLE_PIPELINES[pipeline_type] = SingleColorPipeline(
                 RUNTIME_SINGLE_CONFIGS[pipeline_type],
                 RUNTIME_DEVICE,
+                pipeline_type,
             )
             _SINGLE_PIPELINES[pipeline_type].load()
     return _SINGLE_PIPELINES[pipeline_type]
@@ -702,6 +833,7 @@ def _get_burst_pipeline(pipeline_type: str) -> BurstColorPipeline:
             _BURST_PIPELINES[pipeline_type] = BurstColorPipeline(
                 RUNTIME_BURST_CONFIGS[pipeline_type],
                 RUNTIME_DEVICE,
+                pipeline_type,
             )
             _BURST_PIPELINES[pipeline_type].load()
     return _BURST_PIPELINES[pipeline_type]
@@ -1201,18 +1333,15 @@ def build_demo() -> gr.Blocks:
   <a href="https://arxiv.org/abs/2602.20417">ArXiv</a> |
   <a href="https://github.com/Aryan-Garg/gQIR">GitHub</a>
 </p>
-
 ### What You Can Run
 - **Single Frame (Stage-2):** Reconstruct one frame from either a clean GT image (internally simulated to SPAD) or a real SPAD frame.
 - **Burst (Stage-3):** Reconstruct from a fixed **77-frame** window using either GT videos/cubes or real photon cubes.
 - **Pipelines:** Toggle between **Color** and **Monochrome** reconstruction in both tabs.
-
 ### Supported Inputs
 - **Single GT:** Standard image uploads.
 - **Single Real:** Real SPAD frame image uploads.
 - **Burst GT:** Public-friendly videos (`.mp4`, `.mov`, `.wmv`, `.avi`, `.mkv`, `.webm`) plus research cube formats (`.npy`, `.npz`, `.pt`, `.h5`) or image folders.
 - **Burst Real:** Photon cubes (`.npy`, `.npz`, `.pt`, `.h5`) or image folders.
-
 ### Quick Usage
 1. Pick pipeline and input mode.
 2. Load input and select burst start index (for Stage-3).
@@ -1341,12 +1470,16 @@ def build_demo() -> gr.Blocks:
 
 def main() -> None:
     global RUNTIME_SINGLE_CONFIGS, RUNTIME_BURST_CONFIGS, RUNTIME_DEVICE, RUNTIME_BURST_OUT_SIZES
+    global RUNTIME_HF_REPO_ID, RUNTIME_HF_CACHE_DIR, RUNTIME_HF_TOKEN
 
     args = parse_args()
     single_color_cfg = Path(args.single_config if args.single_config else args.single_config_color).resolve()
     burst_color_cfg = Path(args.burst_config if args.burst_config else args.burst_config_color).resolve()
     single_mono_cfg = Path(args.single_config_mono).resolve()
     burst_mono_cfg = Path(args.burst_config_mono).resolve()
+    hf_token = args.hf_token
+    if not hf_token:
+        hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
 
     RUNTIME_SINGLE_CONFIGS = {
         PIPELINE_COLOR: single_color_cfg,
@@ -1356,6 +1489,10 @@ def main() -> None:
         PIPELINE_COLOR: burst_color_cfg,
         PIPELINE_MONO: burst_mono_cfg,
     }
+    RUNTIME_DEVICE = args.device
+    RUNTIME_HF_REPO_ID = str(args.hf_repo_id or HF_DEFAULT_REPO_ID).strip()
+    RUNTIME_HF_CACHE_DIR = str(Path(args.hf_cache_dir).expanduser()) if args.hf_cache_dir else None
+    RUNTIME_HF_TOKEN = hf_token
     RUNTIME_BURST_OUT_SIZES = {}
     for key, cfg_path in RUNTIME_BURST_CONFIGS.items():
         burst_cfg = OmegaConf.load(str(cfg_path))