File size: 4,574 Bytes
31f6f71 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | import gradio as gr
import torch
import spaces
import json
import os
import random
import subprocess
import sys
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")
MAX_SEED = 2**31 - 1
TMP_DIR = Path(__file__).resolve().parent / "tmp" / "dvd_image_min"
TMP_DIR.mkdir(parents=True, exist_ok=True)
def log_event(message: str):
print(f"[DVD Image Min] {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)
@spaces.GPU(duration=30)
def zero_gpu_smoke_test():
log_event("zero_gpu_smoke_test start")
if not torch.cuda.is_available():
log_event("cuda unavailable")
return "CUDA unavailable"
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=120)
def generate_voxels(image, seed, randomize_seed, preprocess_image, dvd_steps, 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)
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": "Default schedule",
"dvd_cfg_constant": 0.7,
"dvd_cfg_early": 0.4,
"dvd_cfg_late": 0.7,
"dvd_cfg_split": 0.5,
},
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]
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.")
log_event(f"generate_voxels done seed={seed}")
return mesh_path, npy_path, int(seed), f"Done. seed={seed}"
with gr.Blocks(title="DVD Image", fill_width=True) as demo:
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)
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="Preprocess image")
dvd_steps = gr.Slider(1, 512, value=256, step=1, label="DVD steps")
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],
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)
|