Restore full image app with GPU-lazy DVD imports

README.md

@@ -5,11 +5,12 @@ colorFrom: blue
 colorTo: green
 sdk: gradio
 sdk_version: 5.34.2
-app_file: app_dvd_image_gpu_lazy.py
+app_file: app_dvd_image.py
 pinned: false
 license: mit
 models:
   - Zhengrui/dvd
+  - microsoft/TRELLIS-image-large
 ---
 
 # DVD: Discrete Voxel Diffusion for 3D Generation and Editing

app_dvd_image.py

@@ -90,6 +90,89 @@ def log_event(message: str):
     print(f"[DVD Space] {message}", flush=True)
 
 
+def _download_file(url: str, dest: Path):
+    import urllib.request
+
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    if dest.exists() and dest.stat().st_size > 0:
+        log_event(f"using cached {dest.name}")
+        return
+    tmp = dest.with_name(dest.name + ".tmp")
+    if tmp.exists():
+        tmp.unlink()
+    log_event(f"downloading {url} -> {dest}")
+    urllib.request.urlretrieve(url, tmp)
+    tmp.replace(dest)
+
+
+def prefetch_dvd_weights():
+    from huggingface_hub import hf_hub_download, snapshot_download
+
+    if DVD_MODEL_REPO:
+        filenames = (
+            "dvd_img.json",
+            "dvd_img.safetensors",
+            "dvd_img_BSP_ft.json",
+            "dvd_img_BSP_ft.safetensors",
+        )
+        for filename in filenames:
+            log_event(f"prefetching {DVD_MODEL_REPO}/{filename}")
+            hf_hub_download(
+                repo_id=DVD_MODEL_REPO,
+                filename=filename,
+                subfolder=DVD_MODEL_SUBFOLDER,
+                revision=DVD_MODEL_REVISION,
+                token=DVD_MODEL_TOKEN,
+            )
+
+    if os.environ.get("DVD_PREFETCH_TRELLIS", "0").lower() not in {"0", "false", "no", "off"}:
+        log_event(f"prefetching TRELLIS stage2 model {TRELLIS_IMAGE_MODEL}")
+        snapshot_download(repo_id=TRELLIS_IMAGE_MODEL)
+
+
+def prefetch_dinov2():
+    import shutil
+    import zipfile
+
+    hub_dir = Path(torch.hub.get_dir())
+    hub_dir.mkdir(parents=True, exist_ok=True)
+    repo_dir = hub_dir / "facebookresearch_dinov2_main"
+    if repo_dir.exists():
+        log_event(f"using cached DINOv2 repo {repo_dir}")
+    else:
+        zip_path = hub_dir / "main.zip"
+        _download_file("https://github.com/facebookresearch/dinov2/zipball/main", zip_path)
+        extract_tmp = hub_dir / "_dvd_dinov2_extract"
+        if extract_tmp.exists():
+            shutil.rmtree(extract_tmp)
+        extract_tmp.mkdir(parents=True, exist_ok=True)
+        log_event(f"extracting DINOv2 repo to {repo_dir}")
+        with zipfile.ZipFile(zip_path) as zf:
+            zf.extractall(extract_tmp)
+            top_level = zf.namelist()[0].split("/", 1)[0]
+        shutil.move(str(extract_tmp / top_level), str(repo_dir))
+        shutil.rmtree(extract_tmp, ignore_errors=True)
+
+    ckpt_dir = hub_dir / "checkpoints"
+    _download_file(
+        "https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_reg4_pretrain.pth",
+        ckpt_dir / "dinov2_vitl14_reg4_pretrain.pth",
+    )
+
+
+def prefetch_assets():
+    if os.environ.get("DVD_PREFETCH_ASSETS", "1").lower() in {"0", "false", "no", "off"}:
+        log_event("asset prefetch disabled")
+        return
+    try:
+        log_event("asset prefetch start")
+        prefetch_dvd_weights()
+        prefetch_dinov2()
+        log_event("asset prefetch done")
+    except Exception as exc:
+        log_event(f"asset prefetch failed; continuing without prefetch: {exc}")
+
+
 def ensure_dvd_imports():
     global DVDImageToVoxelPipeline
     global TrellisImageTo3DPipeline
@@ -203,7 +286,7 @@ EDIT_IMAGE_EXAMPLES = [
     (asset_label(path), str(path)) for path in list_asset_files("assets/example_image_edit", IMAGE_EXTENSIONS)
 ]
 EDIT_VOXEL_EXAMPLES = [
-    (asset_label(path), str(path)) for path in list_asset_files("assets/example_voxel_edit", {".npy", ".pt", ".pth"})
+    (asset_label(path), str(path)) for path in list_asset_files("assets/example_voxel_edit", {".npy", ".npz", ".pt", ".pth"})
 ]
 
 
@@ -290,7 +373,7 @@ def ensure_trellis_pipeline(device: str | None = None):
 
 def preload_zero_gpu_models(device: str | None = None, preload: str | None = None):
     device = device or os.environ.get("DVD_SPACE_DEVICE", "cuda")
-    preload = preload or os.environ.get("DVD_STARTUP_PRELOAD", "gen")
+    preload = preload or os.environ.get("DVD_STARTUP_PRELOAD", "none")
     requested = {
         item.strip().lower()
         for item in preload.replace(";", ",").split(",")
@@ -530,7 +613,7 @@ def mask_inputs():
     ]
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=600)
 def generate_voxels(
     image: Image.Image,
     seed: int,
@@ -651,7 +734,7 @@ def visualize_edit_mask(
     return mesh_path
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=600)
 def run_editing(
     target_image: Image.Image,
     voxels,
@@ -1025,7 +1108,8 @@ def parse_args():
 if __name__ == "__main__":
     args = parse_args()
 
-    # ZeroGPU's CUDA emulation expects CUDA model placement at startup.
+    # Keep startup free of DVD/TRELLIS imports; only prefetch raw assets.
+    prefetch_assets()
     preload_zero_gpu_models(os.environ.get("DVD_SPACE_DEVICE", "cuda"))
     demo.queue().launch(
         share=args.share,