import gradio as gr import torch import spaces import json import os import subprocess import sys import uuid from pathlib import Path import numpy as np 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") MAX_SEED = np.iinfo(np.int32).max TMP_DIR = Path(__file__).resolve().parent / "tmp" / "dvd_space_image" TMP_DIR.mkdir(parents=True, exist_ok=True) IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"} EXAMPLE_DIR = Path(__file__).resolve().parent / "assets" / "example_image" GENERATION_IMAGE_EXAMPLES = [ str(path) for path in sorted(EXAMPLE_DIR.iterdir()) if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS ] def log_event(message: str): print(f"[DVD Space Lean] {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 def save_voxel_coords(voxels, path: str) -> str: coords = voxels.coords.detach().cpu().numpy().astype(np.int32) np.save(path, coords) return path def get_seed(randomize_seed: bool, seed: int) -> int: return int(np.random.randint(0, MAX_SEED)) if randomize_seed else int(seed) @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=300) 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 = get_seed(randomize_seed, seed) image_path = worker_path("input.png") mesh_path = worker_path("generated_voxels.glb") npy_path = worker_path("generated_voxel64_coords.npy") config_path = worker_path("generate_config.json") image.save(image_path) with open(config_path, "w") as f: json.dump( { "image_path": image_path, "mesh_path": mesh_path, "npy_path": npy_path, "seed": int(seed), "preprocess_image": bool(preprocess_image), "dvd_steps": int(dvd_steps), "dvd_cfg_mode": dvd_cfg_mode, "dvd_cfg_constant": float(dvd_cfg_constant), "dvd_cfg_early": float(dvd_cfg_early), "dvd_cfg_late": float(dvd_cfg_late), "dvd_cfg_split": float(dvd_cfg_split), }, f, ) progress(0.08, desc="Running isolated DVD worker") env = os.environ.copy() env.setdefault("PYTHONUNBUFFERED", "1") cmd = [sys.executable, str(Path(__file__).resolve().parent / "space_image_worker.py"), "generate", config_path] log_event("starting isolated DVD worker") proc = subprocess.run(cmd, env=env, text=True, capture_output=True) if proc.stdout: print(proc.stdout, flush=True) if proc.stderr: print(proc.stderr, flush=True) if proc.returncode != 0: raise gr.Error(f"DVD worker failed with exit code {proc.returncode}. Check Space logs.") progress(0.98, desc="Done") log_event(f"generate_voxels done seed={seed} npy={npy_path}") return mesh_path, npy_path, seed, f"Done. seed={seed}" with gr.Blocks(title="DVD Image Generation", 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(scale=1): image = gr.Image(label="Input Image", format="png", image_mode="RGBA", type="pil", height=320) gr.Examples( examples=GENERATION_IMAGE_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=128, 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(scale=1): 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.queue().launch(show_api=False, show_error=True, ssr_mode=False)