app_dvd_image_gpu_lazy.py

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("SPARSE_ATTN_BACKEND", "flash_attn")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("DVD_MODEL_REPO", "Zhengrui/dvd")

import gradio as gr
import numpy as np
import spaces
import torch


MAX_SEED = 2**31 - 1
ROOT_DIR = Path(__file__).resolve().parent
TMP_DIR = ROOT_DIR / "tmp" / "dvd_image_gpu_lazy"
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 []

_dvd_pipe = None
_export_cubified_voxels = None


def log_event(message: str):
    print(f"[DVD GPULazy] {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

    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
    for filename in ("dvd_img.json", "dvd_img.safetensors"):
        log_event(f"prefetching {repo}/{filename}")
        hf_hub_download(
            repo_id=repo,
            filename=filename,
            subfolder=subfolder,
            revision=revision,
            token=token,
        )


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}")


prefetch_assets()


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 ensure_dvd_pipe():
    global _dvd_pipe, _export_cubified_voxels
    if _dvd_pipe is not None:
        log_event("reusing cached DVD image pipeline")
        return _dvd_pipe, _export_cubified_voxels

    # Intentionally import DVD/TRELLIS only after @spaces.GPU has acquired a worker.
    log_event("importing DVD inside GPU callback")
    from dvd import DVDImageToVoxelPipeline, export_cubified_voxels

    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} on cuda")
    _dvd_pipe = DVDImageToVoxelPipeline.from_pretrained(
        repo,
        variant="base",
        device="cuda",
        subfolder=subfolder,
        revision=revision,
        token=token,
    )
    _export_cubified_voxels = export_cubified_voxels
    log_event("DVD image pipeline ready on cuda")
    return _dvd_pipe, _export_cubified_voxels


@spaces.GPU(duration=60)
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=600)
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.")

    progress(0.04, desc="Loading DVD pipeline")
    pipeline, export_cubified_voxels = ensure_dvd_pipe()

    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.18, desc="Sampling DVD voxels")
    log_event(f"sampling seed={seed} steps={dvd_steps}")
    voxels = pipeline.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.queue().launch(show_api=False, show_error=True, ssr_mode=False)