app_space_image.py

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)