update

.gitattributes

@@ -2,3 +2,5 @@
 *.png filter=lfs diff=lfs merge=lfs -text
 *.npz filter=lfs diff=lfs merge=lfs -text
 *.npy filter=lfs diff=lfs merge=lfs -text
+# checkpoints/depth/*.ckpt filter=lfs diff=lfs merge=lfs -text
+# checkpoints/depth/*.pt filter=lfs diff=lfs merge=lfs -text

InfiniDepth/utils/hf_demo_utils.py

@@ -1,6 +1,7 @@
 import os
 import tempfile
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Optional
 
 import cv2
@@ -23,12 +24,109 @@ from .sampling_utils import make_2d_uniform_coord
 from .vis_utils import clip_outliers_by_percentile, colorize_depth_maps
 
 
+PROJECT_ROOT = Path(__file__).resolve().parents[2]
+REPO_CKPT_ROOT = PROJECT_ROOT / "checkpoints" / "depth"
+SPACE_RUNTIME = bool(os.getenv("SPACE_ID"))
+HF_LOCAL_FILES_ONLY = os.getenv("INFINIDEPTH_HF_LOCAL_FILES_ONLY", "1" if SPACE_RUNTIME else "0") == "1"
+
+DEFAULT_MODEL_FILENAMES = {
+    "InfiniDepth": os.getenv("INFINIDEPTH_RGB_CKPT", "infinidepth.ckpt"),
+    "InfiniDepth_DC": os.getenv("INFINIDEPTH_DC_CKPT", "infinidepth_dc.ckpt"),
+}
 DEFAULT_MODEL_PATHS = {
-    "InfiniDepth": "checkpoints/depth/infinidepth.ckpt",
-    "InfiniDepth_DC": "checkpoints/depth/infinidepth_dc.ckpt",
+    model_type: (REPO_CKPT_ROOT / filename).as_posix()
+    for model_type, filename in DEFAULT_MODEL_FILENAMES.items()
 }
-os.environ.setdefault("INFINIDEPTH_MOGE2_PRETRAINED", "checkpoints/depth/moge2.pt")
-DEFAULT_MOGE2_PRETRAINED = os.getenv("INFINIDEPTH_MOGE2_PRETRAINED", "Ruicheng/moge-2-vitl-normal")
+DEFAULT_MOGE2_FILENAME = os.getenv("INFINIDEPTH_MOGE2_FILENAME", "moge2.pt")
+DEFAULT_MOGE2_LOCAL_PATH = REPO_CKPT_ROOT / DEFAULT_MOGE2_FILENAME
+os.environ.setdefault("INFINIDEPTH_MOGE2_PRETRAINED", DEFAULT_MOGE2_LOCAL_PATH.as_posix())
+
+
+def _resolve_hub_artifact(
+    repo_id: Optional[str],
+    filename: str,
+    artifact_label: str,
+    repo_type: str = "model",
+) -> str:
+    if not repo_id:
+        raise FileNotFoundError(
+            f"{artifact_label} is not available locally, and the env var pointing to the Hugging Face repo is not set. "
+            f"Configure the repo env var and add the file to `preload_from_hub`: {filename}."
+        )
+
+    Log.info(
+        f"Resolving {artifact_label} from Hugging Face cache: repo_id={repo_id}, "
+        f"filename={filename}, repo_type={repo_type}, local_files_only={HF_LOCAL_FILES_ONLY}"
+    )
+    return hf_hub_download(
+        repo_id=repo_id,
+        filename=filename,
+        repo_type=repo_type,
+        local_files_only=HF_LOCAL_FILES_ONLY,
+    )
+
+
+def _resolve_local_or_hub_artifact(
+    local_path: Path,
+    filename: str,
+    repo_id: Optional[str],
+    artifact_label: str,
+    repo_type: str = "model",
+) -> str:
+    if local_path.exists():
+        Log.info(f"Using local {artifact_label}: {local_path}")
+        return local_path.as_posix()
+    return _resolve_hub_artifact(
+        repo_id=repo_id,
+        filename=filename,
+        artifact_label=artifact_label,
+        repo_type=repo_type,
+    )
+
+
+def resolve_moge2_pretrained() -> str:
+    configured = os.getenv("INFINIDEPTH_MOGE2_PRETRAINED", DEFAULT_MOGE2_LOCAL_PATH.as_posix())
+    configured_path = Path(configured)
+    if configured_path.exists():
+        Log.info(f"Using configured MoGe-2 weights: {configured_path}")
+        return configured_path.as_posix()
+
+    if configured != DEFAULT_MOGE2_LOCAL_PATH.as_posix():
+        if configured.count("/") == 1 and configured_path.suffix == "":
+            repo_type = os.getenv("INFINIDEPTH_MOGE2_REPO_TYPE", "model")
+            Log.info(
+                f"Resolving MoGe-2 repo from INFINIDEPTH_MOGE2_PRETRAINED via Hugging Face cache: "
+                f"repo_id={configured}, filename={DEFAULT_MOGE2_FILENAME}, "
+                f"repo_type={repo_type}, local_files_only={HF_LOCAL_FILES_ONLY}"
+            )
+            return hf_hub_download(
+                repo_id=configured,
+                filename=DEFAULT_MOGE2_FILENAME,
+                repo_type=repo_type,
+                local_files_only=HF_LOCAL_FILES_ONLY,
+            )
+        Log.info(f"Using non-local MoGe-2 source from INFINIDEPTH_MOGE2_PRETRAINED: {configured}")
+        return configured
+
+    repo_id = os.getenv("INFINIDEPTH_MOGE2_REPO", os.getenv("INFINIDEPTH_CKPT_REPO"))
+    repo_type = os.getenv("INFINIDEPTH_MOGE2_REPO_TYPE", os.getenv("INFINIDEPTH_CKPT_REPO_TYPE", "model"))
+    return _resolve_local_or_hub_artifact(
+        local_path=DEFAULT_MOGE2_LOCAL_PATH,
+        filename=DEFAULT_MOGE2_FILENAME,
+        repo_id=repo_id,
+        artifact_label="MoGe-2 checkpoint",
+        repo_type=repo_type,
+    )
+
+
+def prepare_runtime_assets() -> None:
+    Log.info(
+        f"Preparing runtime assets: checkpoint_root={REPO_CKPT_ROOT}, "
+        f"space_runtime={SPACE_RUNTIME}, hf_local_files_only={HF_LOCAL_FILES_ONLY}"
+    )
+    for model_type in DEFAULT_MODEL_FILENAMES:
+        resolve_checkpoint_path(model_type)
+    resolve_moge2_pretrained()
 
 
 @dataclass
@@ -106,12 +204,13 @@ def _resolve_depth_inputs(
     device: torch.device,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, str]:
     input_depth_path = depth_path if depth_path else None
+    moge2_pretrained = resolve_moge2_pretrained()
     gt_depth, prompt_depth, gt_depth_mask, used_input_depth = prepare_metric_depth_inputs(
         input_depth_path=input_depth_path,
         input_size=input_size,
         image=image,
         device=device,
-        moge2_pretrained=DEFAULT_MOGE2_PRETRAINED,
+        moge2_pretrained=moge2_pretrained,
         depth_load_kwargs={"enable_noise_filter": False},
     )
     prompt_mask = prompt_depth > 0
@@ -185,31 +284,17 @@ def resolve_checkpoint_path(model_type: str) -> str:
     if model_type not in DEFAULT_MODEL_PATHS:
         raise ValueError(f"Unsupported model_type: {model_type}")
 
-    local_path = DEFAULT_MODEL_PATHS[model_type]
-    if os.path.exists(local_path):
-        Log.info(f"Using local checkpoint for {model_type}: {local_path}")
-        return local_path
-
+    local_path = Path(DEFAULT_MODEL_PATHS[model_type])
     repo_id = os.getenv("INFINIDEPTH_CKPT_REPO")
-    if not repo_id:
-        raise FileNotFoundError(
-            f"Local checkpoint not found at {local_path}, and env INFINIDEPTH_CKPT_REPO is not set for HF Hub download."
-        )
-
-    filename_key = "INFINIDEPTH_RGB_CKPT" if model_type == "InfiniDepth" else "INFINIDEPTH_DC_CKPT"
-    default_name = os.path.basename(local_path)
-    filename = os.getenv(filename_key, default_name)
-    Log.info(
-        f"Downloading checkpoint for {model_type} from Hugging Face Hub: repo_id={repo_id}, filename={filename}"
-    )
-
-    downloaded = hf_hub_download(
-        repo_id=repo_id,
+    repo_type = os.getenv("INFINIDEPTH_CKPT_REPO_TYPE", "model")
+    filename = DEFAULT_MODEL_FILENAMES[model_type]
+    return _resolve_local_or_hub_artifact(
+        local_path=local_path,
         filename=filename,
-        local_dir=os.path.dirname(local_path),
-        local_dir_use_symlinks=False,
+        repo_id=repo_id,
+        artifact_label=f"{model_type} checkpoint",
+        repo_type=repo_type,
     )
-    return downloaded
 
 
 def run_single_image_demo(

README.md

@@ -6,8 +6,14 @@ colorTo: purple
 sdk: gradio
 sdk_version: 5.47.2
 app_file: app.py
+preload_from_hub:
+  - ritianyu/InfiniDepth infinidepth.ckpt,infinidepth_dc.ckpt,moge2.pt
 pinned: false
 license: apache-2.0
 short_description: InfiniDepth Huggingface Demo
 ---
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+For Hugging Face Spaces, `preload_from_hub` now preloads the three checkpoint files from
+`ritianyu/InfiniDepth` during build time. Set `INFINIDEPTH_CKPT_REPO=ritianyu/InfiniDepth`
+in the Space variables so runtime loading resolves the same cached files.

app.py

@@ -22,7 +22,7 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 
-from InfiniDepth.utils.hf_demo_utils import ModelCache, run_single_image_demo
+from InfiniDepth.utils.hf_demo_utils import ModelCache, prepare_runtime_assets, run_single_image_demo
 from InfiniDepth.utils.logger import Log
 
 try:
@@ -423,6 +423,7 @@ demo = demo.queue()
 
 
 if __name__ == "__main__":
+    prepare_runtime_assets()
     server_name = "0.0.0.0" if os.getenv("SPACE_ID") else "127.0.0.1"
     # Hugging Face Spaces with Gradio typically expects port 7860.
     # Respect explicit Gradio/PORT env overrides when provided.