huggingface-projects/diffusiongemma-3d-gen
8
1"""2Gemma Diffusion — text → 3D asset builder (gradio.Server backend + custom frontend).3 4ZeroGPU port. The block-diffusion model designs a standalone SVG illustration of a5described asset; the custom frontend extrudes that SVG into a live, spinning Three.js63D scene. `gradio.Server` (a FastAPI subclass) provides Gradio's queue + SSE streaming7under our hand-written HTML/CSS/JS frontend. The single streaming endpoint `/generate`8yields one JSON frame per denoising step: the raw SVG canvas diffusing on the left, the9extruded 3D object rendering on the right.10 11ZeroGPU specifics:12- `import spaces` happens before `torch`.13- The model is loaded once at module scope with `.to("cuda")` (ZeroGPU registers it).14- The actual `model.generate` call lives inside the `@spaces.GPU` function `_gpu_stream`;15 the `gradio.Server` endpoint only marshals picklable CPU tensors in/out of it.16 17Refs:18- https://huggingface.co/blog/introducing-gradio-server19- https://huggingface.co/docs/hub/spaces-zerogpu20"""21 22import glob23import os24import subprocess25import sys26 27# Set before torch is imported (transformers pulls torch in).28os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")29 30import spaces # must precede torch so ZeroGPU can patch it31 32 33def _ensure_transformers():34 """Install the bundled custom DiffusionGemma `transformers` wheel at runtime.35 36 Spaces installs `requirements.txt` *before* copying the repo files into the image,37 so the wheel can't be referenced by local path there. By the time this app runs the38 file is present in the working directory, so we install it here (only if a stock /39 no transformers is importable) before importing torch/transformers below.40 """41 try:42 import transformers # noqa: F40143 44 if hasattr(transformers, "DiffusionGemmaForBlockDiffusion") or hasattr(45 getattr(transformers, "models", object), "diffusion_gemma"46 ):47 return48 except Exception:49 pass50 wheels = sorted(glob.glob(os.path.join(os.path.dirname(os.path.abspath(__file__)), "transformers-*.whl")))51 if not wheels:52 return53 print(f"[gdiff] Installing bundled transformers wheel: {os.path.basename(wheels[0])}", flush=True)54 subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-cache-dir", wheels[0]])55 import importlib56 57 importlib.invalidate_caches()58 59 60_ensure_transformers()61 62import json63import queue as queue_lib64import re65import threading66import time as _time67 68import torch69from fastapi.responses import HTMLResponse70from gradio import Server71from transformers import AutoTokenizer, DiffusionGemmaForBlockDiffusion72from transformers.generation.streamers import BaseStreamer73 74HERE = os.path.dirname(os.path.abspath(__file__))75MODEL_PATH = os.environ.get("GDIFF_MODEL_PATH", "google/diffusiongemma-26B-A4B-it")76HF_TOKEN = os.environ.get("HF_TOKEN")77MAX_ITERS_CAP = 120 # hard cap on denoising steps per block78# ZeroGPU: the 26B checkpoint (~49 GB bf16) needs the full backing card.79GPU_SIZE = os.environ.get("GDIFF_GPU_SIZE", "xlarge")80 81SYSTEM_PROMPT = (82 "You are an expert vector artist. Given a TEXT description (usually a game asset — a "83 "sword, shield, potion, coin, treasure chest, spaceship, robot, mushroom, key, gem, "84 "etc.) you design an original, polished SVG illustration of it. The SVG will be extruded "85 "into a spinning 3D object, so design it with clean, solid, extrudable shapes.\n"86 "\n"87 "Requirements:\n"88 "- Output ONLY a single standalone SVG document: start your response with `<svg` and end "89 "with `</svg>`. No HTML wrapper, no <?xml?> prologue, no markdown code fences, no "90 "explanation.\n"91 '- The opening tag must include xmlns="http://www.w3.org/2000/svg" and a square viewBox '92 '(e.g. "0 0 100 100").\n'93 "- Draw in a bold, readable, flat 'game asset / icon' style: several distinct shapes "94 "(<path>, <rect>, <circle>, <polygon>) each with a SOLID fill color (the `fill` "95 "attribute) and a coherent, attractive palette. Layer shapes to suggest detail (outline, "96 "body, highlights, shading).\n"97 "- Do NOT add a full-bleed background rectangle — keep the background transparent so each "98 "shape becomes its own clean 3D piece against the dark scene.\n"99 "- Use only solid filled shapes. Avoid gradients, filters, <text>, images, and "100 "stroke-only / fill=\"none\" shapes — they do not extrude.\n"101 "- Use enough shapes to look great while staying clean (roughly 6-16 shapes).\n"102 "- When asked to modify the artwork, return the FULL updated SVG with the change applied, "103 "keeping the same subject unless asked to change it.\n"104)105 106_MARKER_RE = re.compile(107 r"<\|?(?:channel|turn|think|image|audio|video|tool(?:_call|_response)?)\|?>"108)109_FENCE_RE = re.compile(r"```(?:html|svg|xml)?\s*(.*?)\s*```", re.DOTALL)110_SVG_CHILD_RE = re.compile(111 r"<(?:path|rect|circle|ellipse|polygon|polyline|line|g|defs)\b", re.I112)113_SVG_OPEN = '<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 100 100">\n'114 115 116# --------------------------------------------------------------------------- #117# Model (loaded once at module scope; ZeroGPU registers .to("cuda") tensors)118# --------------------------------------------------------------------------- #119DEVICE = "cuda" if torch.cuda.is_available() else "cpu"120print(f"[gdiff] Loading model from {MODEL_PATH} on {DEVICE} ...", flush=True)121tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, token=HF_TOKEN)122model = DiffusionGemmaForBlockDiffusion.from_pretrained(123 MODEL_PATH,124 dtype=torch.bfloat16,125 low_cpu_mem_usage=True,126 token=HF_TOKEN,127).to(DEVICE)128model.eval()129CANVAS_LEN = model.config.canvas_length130PAD_ID = tokenizer.pad_token_id or 0131print(f"[gdiff] Model ready | canvas_length={CANVAS_LEN}", flush=True)132 133# Cache of the last *cleaned* SVG so a follow-up tweak can warm-start in place.134model._last_clean_html = None135 136 137# --------------------------------------------------------------------------- #138# Helpers (CPU-only; safe to run in the gradio.Server main process)139# --------------------------------------------------------------------------- #140def warm_canvas_from_cache():141 """Starting canvas (first block) built from the previous *cleaned* SVG.142 143 Returns a CPU tensor (it is pickled across the ZeroGPU process boundary and moved144 to CUDA inside the GPU worker). We re-tokenize the cleaned SVG rather than reuse145 raw output tokens so a mangled ``<svg`` header can't compound across tweaks.146 """147 svg = getattr(model, "_last_clean_html", None)148 if not svg:149 return None150 ids = tokenizer(svg, add_special_tokens=False).input_ids[:CANVAS_LEN]151 if not ids:152 return None153 if len(ids) < CANVAS_LEN:154 ids = ids + [PAD_ID] * (CANVAS_LEN - len(ids))155 return torch.tensor(ids, dtype=torch.long).unsqueeze(0)156 157 158def last_assistant_html(history_json: str):159 try:160 history = json.loads(history_json) if history_json else []161 except json.JSONDecodeError:162 return None163 for turn in reversed(history):164 if turn.get("role") == "assistant" and turn.get("content"):165 return turn["content"]166 return None167 168 169def clean_text(text: str) -> str:170 return _MARKER_RE.sub("", text).lstrip()171 172 173def extract_svg(text: str) -> str:174 """Pull a clean standalone <svg>…</svg> out of the (possibly mangled) model output.175 176 Warm-start diffusion frequently chews the very front of the document — the opening177 ``<svg`` loses its ``<`` (``svg viewBox=…``) or a few more chars. If we can't find an178 intact ``<svg``, we rebuild a canonical wrapper around the first real child element so179 the output is always valid (the 3D viewer auto-fits the camera, so a default viewBox is180 fine). Repairing here is essential: the cleaned result is what we cache for the next181 tweak's warm-start, which stops corruption from compounding across tweaks.182 """183 text = clean_text(text)184 fenced = _FENCE_RE.search(text)185 if fenced:186 text = fenced.group(1)187 188 lower = text.lower()189 s = lower.find("<svg")190 if s != -1:191 text = text[s:]192 else:193 m = _SVG_CHILD_RE.search(text)194 if m:195 text = _SVG_OPEN + text[m.start():]196 # else: nothing salvageable; fall through and just trim/close it197 198 lower = text.lower()199 e = lower.rfind("</svg>")200 if e != -1:201 text = text[: e + len("</svg>")]202 else:203 text = text.rstrip() + "\n</svg>" # tail eaten mid-stream; close it204 return text.strip()205 206 207class QueueDiffusionStreamer(BaseStreamer):208 def __init__(self, tok, q: "queue_lib.Queue"):209 self.tok = tok210 self.q = q211 self.confirmed_ids: list[int] = []212 self.prompt_skipped = False213 self.block = 0214 self.step = 0215 216 def _decode(self, ids):217 return self.tok.decode(ids, skip_special_tokens=True)218 219 def put(self, value):220 ids = value[0].tolist() if value.dim() > 1 else value.tolist()221 if not self.prompt_skipped:222 self.prompt_skipped = True223 return224 self.confirmed_ids.extend(ids)225 self.block += 1226 self.step = 0227 self.q.put(("commit", self._decode(self.confirmed_ids), self.block, self.step))228 229 def put_draft(self, value):230 self.step += 1231 ids = value[0].tolist() if value.dim() > 1 else value.tolist()232 self.q.put(("draft", self._decode(self.confirmed_ids + ids), self.block + 1, self.step))233 234 def end(self):235 self.q.put(("end", self._decode(self.confirmed_ids), self.block, self.step))236 237 238def build_messages(history_json: str, prompt: str):239 try:240 history = json.loads(history_json) if history_json else []241 except json.JSONDecodeError:242 history = []243 messages = [{"role": "system", "content": SYSTEM_PROMPT}]244 for turn in history:245 role = turn.get("role")246 content = turn.get("content", "")247 if role in ("user", "assistant") and content:248 messages.append({"role": role, "content": content})249 messages.append({"role": "user", "content": prompt})250 return messages251 252 253# --------------------------------------------------------------------------- #254# GPU work — runs in a forked ZeroGPU worker process.255# Inputs/outputs cross the boundary via pickle, so only CPU tensors / plain256# Python objects go in and out (no CUDA tensors are returned).257# --------------------------------------------------------------------------- #258def _estimate_duration(input_ids, max_new_tokens=2048, max_iters=64, full_denoise=False, canvas_ids=None):259 blocks = max(1, int(max_new_tokens) // max(1, CANVAS_LEN))260 secs = 30 + blocks * int(max_iters) * 0.3261 return int(min(120, secs)) # xlarge internally doubles this for the quota check262 263 264@spaces.GPU(duration=_estimate_duration, size=GPU_SIZE)265def _gpu_stream(input_ids, max_new_tokens, max_iters, full_denoise, canvas_ids):266 input_ids = input_ids.to(model.device)267 gen_kwargs = dict(max_new_tokens=int(max_new_tokens), max_denoising_steps=int(max_iters))268 if full_denoise:269 gen_kwargs["confidence_threshold"] = 1e-9270 gen_kwargs["stability_threshold"] = int(max_iters)271 if canvas_ids is not None:272 gen_kwargs["canvas_ids"] = canvas_ids.to(model.device)273 274 q: "queue_lib.Queue" = queue_lib.Queue()275 streamer = QueueDiffusionStreamer(tokenizer, q)276 err = {}277 278 def worker():279 try:280 with torch.inference_mode():281 model.generate(input_ids, streamer=streamer, **gen_kwargs)282 except Exception as exc: # surface to the endpoint283 err["msg"] = f"{type(exc).__name__}: {exc}"284 q.put(("error", str(exc), 0, 0))285 finally:286 q.put(("end", "", 0, 0)) # always unblock the consumer287 288 thread = threading.Thread(target=worker)289 thread.start()290 try:291 while True:292 kind, text, block, step = q.get()293 if kind == "error":294 yield ("error", err.get("msg", text), 0, 0)295 return296 if kind == "end":297 return298 yield (kind, text, block, step)299 finally:300 thread.join()301 302 303# --------------------------------------------------------------------------- #304# Server305# --------------------------------------------------------------------------- #306app = Server(title="Gemma Diffusion 3D Asset Builder")307 308 309@app.api(name="generate", concurrency_limit=1, time_limit=600, stream_every=0.05)310def generate(311 prompt: str,312 history_json: str = "[]",313 max_new_tokens: int = 2048,314 max_iters: int = 64,315 full_denoise: bool = False,316 anim_delay: float = 0.0,317 warm_start: bool = True,318) -> str:319 """Stream the diffusion generation as JSON frames (one per denoising step).320 321 The model emits a raw SVG illustration; the frontend extrudes it into 3D with Three.js.322 """323 prompt = (prompt or "").strip()324 if not prompt:325 yield json.dumps({"kind": "error", "message": "Empty prompt."})326 return327 328 messages = build_messages(history_json, prompt)329 max_iters = max(1, min(int(max_iters), MAX_ITERS_CAP))330 331 # Tweak warm-start: seed the diffusion's first canvas with the previous artwork's own332 # tokens (native `canvas_ids` API) so the model edits the existing SVG in place.333 is_tweak = bool(last_assistant_html(history_json))334 canvas_ids = warm_canvas_from_cache() if (warm_start and is_tweak) else None335 warming = canvas_ids is not None336 337 input_ids = tokenizer.apply_chat_template(338 messages,339 tokenize=True,340 add_generation_prompt=True,341 return_tensors="pt",342 return_dict=True,343 )["input_ids"]344 345 last_text = ""346 for kind, text, block, step in _gpu_stream(347 input_ids, int(max_new_tokens), max_iters, bool(full_denoise), canvas_ids348 ):349 if kind == "error":350 yield json.dumps({"kind": "error", "message": text})351 return352 last_text = text353 yield json.dumps(354 {355 "kind": "draft" if kind == "draft" else "commit",356 "source": clean_text(text),357 "block": block,358 "step": step,359 "canvas": CANVAS_LEN,360 "max_iters": max_iters,361 "warming": warming,362 }363 )364 if anim_delay and kind == "draft":365 _time.sleep(float(anim_delay))366 367 final_source = extract_svg(last_text)368 # Cache the *cleaned* SVG so the next tweak warm-starts from a valid header.369 if final_source.strip():370 model._last_clean_html = final_source371 yield json.dumps({"kind": "done", "source": final_source})372 373 374@app.get("/", response_class=HTMLResponse)375async def homepage():376 with open(os.path.join(HERE, "index.html"), "r", encoding="utf-8") as f:377 return f.read()378 379 380# HF Spaces' gradio runtime looks for a top-level `demo` (or `app`) to launch.381demo = app382 383if __name__ == "__main__":384 app.launch(385 server_name=os.environ.get("GDIFF_HOST", "0.0.0.0"),386 server_port=int(os.environ.get("GDIFF_PORT", "7860")),387 show_error=True,388 )389