Use official-style DVD image pipeline wrapper

README.md

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

app_dvd_image_wrapper.py

@@ -0,0 +1,187 @@
+import os
+import random
+import uuid
+from pathlib import Path
+
+os.environ.setdefault("SPCONV_ALGO", "native")
+os.environ.setdefault("ATTN_BACKEND", "flash_attn")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+os.environ.setdefault("DVD_MODEL_REPO", "Zhengrui/dvd")
+
+import spaces
+import gradio as gr
+import numpy as np
+import torch
+
+from dvd import DVDImageToVoxelPipeline, export_cubified_voxels
+
+
+MAX_SEED = 2**31 - 1
+ROOT_DIR = Path(__file__).resolve().parent
+TMP_DIR = ROOT_DIR / "tmp" / "dvd_image_wrapper"
+TMP_DIR.mkdir(parents=True, exist_ok=True)
+IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
+EXAMPLE_DIR = ROOT_DIR / "assets" / "example_image"
+EXAMPLES = [
+    str(path)
+    for path in sorted(EXAMPLE_DIR.iterdir())
+    if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS
+] if EXAMPLE_DIR.exists() else []
+
+
+def log_event(message: str):
+    print(f"[DVD Wrapper] {message}", flush=True)
+
+
+def worker_path(name: str) -> str:
+    path = TMP_DIR / f"worker-{uuid.uuid4().hex}"
+    path.mkdir(parents=True, exist_ok=True)
+    return str(path / name)
+
+
+def cfg_schedule(mode: str, constant: float, early: float, late: float, split: float):
+    if mode == "Constant":
+        return float(constant)
+    if mode == "Two-stage":
+        split = float(split)
+        early = float(early)
+        late = float(late)
+        return lambda t: early if t < split else late
+    return None
+
+
+repo = os.environ.get("DVD_MODEL_REPO", "Zhengrui/dvd")
+subfolder = os.environ.get("DVD_MODEL_SUBFOLDER") or None
+revision = os.environ.get("DVD_MODEL_REVISION") or None
+token = os.environ.get("DVD_MODEL_TOKEN") or os.environ.get("HF_TOKEN") or None
+
+log_event(f"loading DVD image pipeline from {repo}")
+dvd_pipe = DVDImageToVoxelPipeline.from_pretrained(
+    repo,
+    variant="base",
+    subfolder=subfolder,
+    revision=revision,
+    token=token,
+)
+log_event("moving DVD image pipeline to cuda")
+dvd_pipe.to("cuda")
+log_event("DVD image pipeline ready on cuda")
+
+
+@spaces.GPU(duration=30)
+def zero_gpu_smoke_test():
+    log_event("zero_gpu_smoke_test start")
+    if not torch.cuda.is_available():
+        log_event("zero_gpu_smoke_test no cuda")
+        return "CUDA unavailable inside ZeroGPU worker"
+    value = torch.ones((1,), device="cuda").sum().item()
+    name = torch.cuda.get_device_name(0)
+    log_event(f"zero_gpu_smoke_test done device={name} value={value}")
+    return f"OK: {name}, value={value}"
+
+
+@spaces.GPU(duration=180)
+def generate_voxels(
+    image,
+    seed: int,
+    randomize_seed: bool,
+    preprocess_image: bool,
+    dvd_steps: int,
+    dvd_cfg_mode: str,
+    dvd_cfg_constant: float,
+    dvd_cfg_early: float,
+    dvd_cfg_late: float,
+    dvd_cfg_split: float,
+    progress=gr.Progress(track_tqdm=True),
+):
+    progress(0.01, desc="Starting ZeroGPU callback")
+    log_event(f"generate_voxels start seed={seed} randomize={randomize_seed} steps={dvd_steps}")
+    if image is None:
+        raise gr.Error("Please provide an image.")
+
+    seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+    sampler_kwargs = {"steps": int(dvd_steps)}
+    schedule = cfg_schedule(
+        dvd_cfg_mode,
+        dvd_cfg_constant,
+        dvd_cfg_early,
+        dvd_cfg_late,
+        dvd_cfg_split,
+    )
+    if schedule is not None:
+        sampler_kwargs["cfg_strength"] = schedule
+
+    progress(0.08, desc="Sampling DVD voxels")
+    voxels = dvd_pipe.sample_voxels(
+        image,
+        seed=seed,
+        preprocess_image=preprocess_image,
+        **sampler_kwargs,
+    )
+
+    progress(0.88, desc="Exporting voxel preview")
+    mesh_path = worker_path("generated_voxels.glb")
+    npy_path = worker_path("generated_voxel64_coords.npy")
+    export_cubified_voxels(voxels, mesh_path)
+    np.save(npy_path, voxels.coords_without_batch.detach().cpu().numpy().astype(np.int32))
+    torch.cuda.empty_cache()
+    log_event(f"generate_voxels done seed={seed} mesh={mesh_path} npy={npy_path}")
+    return mesh_path, npy_path, int(seed), f"Done. seed={seed}"
+
+
+with gr.Blocks(title="DVD Image", fill_width=True) as demo:
+    gr.Markdown("## DVD Image Voxel Generation")
+    with gr.Row():
+        smoke_btn = gr.Button("ZeroGPU Smoke Test")
+        smoke_out = gr.Textbox(label="ZeroGPU Status", interactive=False)
+    smoke_btn.click(zero_gpu_smoke_test, outputs=smoke_out)
+
+    with gr.Row(equal_height=False):
+        with gr.Column():
+            image = gr.Image(label="Input Image", format="png", image_mode="RGBA", type="pil", height=320)
+            if EXAMPLES:
+                gr.Examples(examples=EXAMPLES[:12], inputs=image, examples_per_page=6)
+            with gr.Accordion("DVD Settings", open=False):
+                seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
+                randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
+                preprocess_image = gr.Checkbox(value=True, label="DVD preprocess image")
+                dvd_steps = gr.Slider(1, 512, value=256, step=1, label="DVD steps")
+                dvd_cfg_mode = gr.Radio(
+                    ["Default schedule", "Constant", "Two-stage"],
+                    value="Default schedule",
+                    label="DVD CFG mode",
+                )
+                dvd_cfg_constant = gr.Slider(0.0, 5.0, value=0.7, step=0.05, label="Constant CFG")
+                dvd_cfg_early = gr.Slider(0.0, 5.0, value=0.4, step=0.05, label="Early CFG")
+                dvd_cfg_late = gr.Slider(0.0, 5.0, value=0.7, step=0.05, label="Late CFG")
+                dvd_cfg_split = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="CFG switch time")
+            gen_btn = gr.Button("Generate DVD Voxels", variant="primary")
+        with gr.Column():
+            voxel_view = gr.Model3D(
+                label="Generated / Cubified Voxels",
+                height=360,
+                camera_position=(-180, 90, 3),
+            )
+            npy_download = gr.DownloadButton(label="Download Voxel Coords (.npy)", interactive=False)
+            status = gr.Textbox(label="Status", interactive=False)
+
+    gen_btn.click(
+        generate_voxels,
+        inputs=[
+            image,
+            seed,
+            randomize_seed,
+            preprocess_image,
+            dvd_steps,
+            dvd_cfg_mode,
+            dvd_cfg_constant,
+            dvd_cfg_early,
+            dvd_cfg_late,
+            dvd_cfg_split,
+        ],
+        outputs=[voxel_view, npy_download, seed, status],
+    ).then(lambda: gr.DownloadButton(interactive=True), outputs=[npy_download])
+
+
+if __name__ == "__main__":
+    demo.launch(show_api=False, show_error=True, ssr_mode=False)