app.py

app.py

@@ -15,6 +15,20 @@
 
 import os
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
+
+# Redirect all model caches to persistent storage if /data is mounted.
+# Without this, the 28GB of model downloads happen on every cold start
+# and we hit the 30-minute HF Spaces healthcheck timeout.
+if os.path.isdir("/data") and os.access("/data", os.W_OK):
+    os.makedirs("/data/hf_cache", exist_ok=True)
+    os.makedirs("/data/torch_cache", exist_ok=True)
+    os.environ["HF_HOME"] = "/data/hf_cache"
+    os.environ["HUGGINGFACE_HUB_CACHE"] = "/data/hf_cache"
+    os.environ["TORCH_HOME"] = "/data/torch_cache"
+    print("[persist] Using /data for HF and torch caches")
+else:
+    print("[persist] /data not available, using ephemeral storage")
+
 os.system("rm -rf /data-nvme/zerogpu-offload/")
 os.system("pip install --no-build-isolation chumpy")
 os.system("pip uninstall -y basicsr")
@@ -390,15 +404,53 @@ def _build_model(cfg):
 
 def launch_pretrained():
     from huggingface_hub import snapshot_download, hf_hub_download
-    hf_hub_download(repo_id="3DAIGC/LHM", repo_type='model', filename='assets.tar', local_dir="./")
-    os.system("tar -xf assets.tar && rm assets.tar")
-    # hf_hub_download(repo_id="3DAIGC/LHM", repo_type='model', filename='LHM-0.5B.tar', local_dir="./")
-    # os.system("tar -xf LHM-0.5B.tar && rm LHM-0.5B.tar")
-    hf_hub_download(repo_id="3DAIGC/LHM", repo_type='model', filename='LHM_prior_model.tar', local_dir="./")
-    os.system("tar -xf LHM_prior_model.tar && rm LHM_prior_model.tar")
-    # replace the weight of full body
-    hf_hub_download(repo_id="3DAIGC/LHM-500M-HF", repo_type='model', filename='config.json', local_dir="./exps/releases/video_human_benchmark/human-lrm-500M/step_060000/")
-    hf_hub_download(repo_id="3DAIGC/LHM-500M-HF", repo_type='model', filename='model.safetensors', local_dir="./exps/releases/video_human_benchmark/human-lrm-500M/step_060000/")
+
+    persist = "/data" if (os.path.isdir("/data") and os.access("/data", os.W_OK)) else None
+
+    if persist:
+        # Cache extracted models on persistent storage and symlink into app dir.
+        # First boot still does the full download; subsequent boots are fast.
+        extract_root = f"{persist}/extracted"
+        os.makedirs(extract_root, exist_ok=True)
+
+        # assets.tar -> /data/extracted/assets
+        if not os.path.isdir(f"{extract_root}/assets"):
+            print("[persist] Downloading assets.tar to /data...")
+            hf_hub_download(repo_id="3DAIGC/LHM", repo_type='model', filename='assets.tar', local_dir=persist)
+            os.system(f"cd {extract_root} && tar -xf {persist}/assets.tar && rm {persist}/assets.tar")
+        else:
+            print("[persist] assets/ already cached, skipping download")
+        os.system(f"rm -rf ./assets && ln -sfn {extract_root}/assets ./assets")
+
+        # LHM_prior_model.tar -> /data/extracted/pretrained_models (and possibly other dirs)
+        if not os.path.isdir(f"{extract_root}/pretrained_models"):
+            print("[persist] Downloading LHM_prior_model.tar to /data...")
+            hf_hub_download(repo_id="3DAIGC/LHM", repo_type='model', filename='LHM_prior_model.tar', local_dir=persist)
+            os.system(f"cd {extract_root} && tar -xf {persist}/LHM_prior_model.tar && rm {persist}/LHM_prior_model.tar")
+        else:
+            print("[persist] pretrained_models/ already cached, skipping download")
+        os.system(f"rm -rf ./pretrained_models && ln -sfn {extract_root}/pretrained_models ./pretrained_models")
+
+        # 500M-HF weights -> /data/extracted/exps/.../step_060000
+        target_rel = "exps/releases/video_human_benchmark/human-lrm-500M/step_060000"
+        target_abs = f"{extract_root}/{target_rel}"
+        if not os.path.isfile(f"{target_abs}/model.safetensors"):
+            print("[persist] Downloading 500M-HF weights to /data...")
+            os.makedirs(target_abs, exist_ok=True)
+            hf_hub_download(repo_id="3DAIGC/LHM-500M-HF", repo_type='model', filename='config.json', local_dir=target_abs)
+            hf_hub_download(repo_id="3DAIGC/LHM-500M-HF", repo_type='model', filename='model.safetensors', local_dir=target_abs)
+        else:
+            print("[persist] 500M-HF weights already cached, skipping download")
+        os.makedirs("./exps/releases/video_human_benchmark/human-lrm-500M", exist_ok=True)
+        os.system(f"rm -rf ./{target_rel} && ln -sfn {target_abs} ./{target_rel}")
+    else:
+        # Original behavior (ephemeral storage)
+        hf_hub_download(repo_id="3DAIGC/LHM", repo_type='model', filename='assets.tar', local_dir="./")
+        os.system("tar -xf assets.tar && rm assets.tar")
+        hf_hub_download(repo_id="3DAIGC/LHM", repo_type='model', filename='LHM_prior_model.tar', local_dir="./")
+        os.system("tar -xf LHM_prior_model.tar && rm LHM_prior_model.tar")
+        hf_hub_download(repo_id="3DAIGC/LHM-500M-HF", repo_type='model', filename='config.json', local_dir="./exps/releases/video_human_benchmark/human-lrm-500M/step_060000/")
+        hf_hub_download(repo_id="3DAIGC/LHM-500M-HF", repo_type='model', filename='model.safetensors', local_dir="./exps/releases/video_human_benchmark/human-lrm-500M/step_060000/")
 
 def launch_env_not_compile_with_cuda():
     os.system("pip install --no-build-isolation chumpy")