"""Text-to-Splat – gradio.Server with custom frontend. Same image→3D-Gaussian pipeline as TripoSplat, with a text-to-image front step powered by Ideogram 4 (called as a remote Gradio client) so users can go straight from a text prompt to a 3D splat. Usage: python app.py """ import base64 import os import tempfile import time from pathlib import Path from uuid import uuid4 import spaces import torch from huggingface_hub import snapshot_download from PIL import Image from fastapi.responses import HTMLResponse, FileResponse, JSONResponse import gradio as gr from gradio import Server from gradio.data_classes import FileData from gradio_client import Client as GradioClient from triposplat import TripoSplatPipeline import example_inputs_b64 as _b64 # ---------------------------------------------------------------------------- # Download checkpoints from HuggingFace Hub (VAST-AI/TripoSplat) # ---------------------------------------------------------------------------- snapshot_download(repo_id="VAST-AI/TripoSplat", local_dir="ckpts") # ---------------------------------------------------------------------------- # Pipeline (loaded once at startup) # ---------------------------------------------------------------------------- PIPE = TripoSplatPipeline( ckpt_path = "ckpts/diffusion_models/triposplat_fp16.safetensors", decoder_path = "ckpts/vae/triposplat_vae_decoder_fp16.safetensors", dinov3_path = "ckpts/clip_vision/dino_v3_vit_h.safetensors", flux2_vae_encoder_path = "ckpts/vae/flux2-vae.safetensors", rmbg_path = "ckpts/background_removal/birefnet.safetensors", device = "cuda", ) OUT_ROOT = Path("gradio_outputs").resolve() OUT_ROOT.mkdir(parents=True, exist_ok=True) # Decode example images from base64 into a persistent temp directory so that # the custom frontend can serve them via FastAPI routes. _EXAMPLES_TMPDIR = tempfile.mkdtemp(prefix="triposplat_examples_") def _write_example(varname: str, filename: str) -> str: path = Path(_EXAMPLES_TMPDIR) / filename path.write_bytes(base64.b64decode(getattr(_b64, varname))) return str(path) EXAMPLES = [ {"name": "Creature Butterfly", "file": _write_example("CREATURE_BUTTERFLY", "creature_butterfly.webp")}, {"name": "Building Stone House","file": _write_example("BUILDING_STONE_HOUSE", "building_stone_house.webp")}, {"name": "Vehicle Pirate Ship", "file": _write_example("VEHICLE_PIRATE_SHIP", "vehicle_pirate_ship.webp")}, {"name": "Plant Water Lily", "file": _write_example("PLANT_WATER_LILY", "plant_water_lily.webp")}, ] # ---------------------------------------------------------------------------- # Ideogram 4 text-to-image (remote Gradio client) # ---------------------------------------------------------------------------- # Text prompts are turned into an image by calling the public Ideogram 4 Space # as a Gradio client. This is a plain network call (it runs on Ideogram's own # Space, not this GPU), so it is *not* wrapped in @spaces.GPU and never touches # the ZeroGPU allocation used by the 3D pipeline below. # # Set an ``HF_TOKEN`` secret on this Space so the remote call is attributed to # your account's ZeroGPU quota instead of the shared anonymous pool. IDEOGRAM_SPACE = os.environ.get("IDEOGRAM_SPACE", "ideogram-ai/ideogram4") _HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") _ideogram_client = None def _get_ideogram_client() -> GradioClient: """Connect to the Ideogram 4 Space lazily and cache the client. The auth keyword has changed names across gradio_client versions (``hf_token`` vs ``token``), so try the token-bearing constructors first and fall back to an anonymous client if neither kwarg is accepted. """ global _ideogram_client if _ideogram_client is not None: return _ideogram_client attempts = [] if _HF_TOKEN: attempts += [{"hf_token": _HF_TOKEN}, {"token": _HF_TOKEN}] attempts += [{}] last_err = None for kwargs in attempts: try: _ideogram_client = GradioClient(IDEOGRAM_SPACE, **kwargs) return _ideogram_client except TypeError as err: last_err = err # unsupported kwarg name → try the next form raise last_err # ---------------------------------------------------------------------------- # gradio.Server # ---------------------------------------------------------------------------- app = Server() # ---- Static pages ---------------------------------------------------------- @app.get("/") async def homepage(): """Serve the custom frontend.""" html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html") with open(html_path, "r", encoding="utf-8") as f: return HTMLResponse(f.read()) @app.get("/viewer") async def viewer_page(): """Serve the Spark.js 3D viewer (loaded inside an iframe).""" viewer_path = os.path.join( os.path.dirname(os.path.abspath(__file__)), "static", "viewer", "viewer.html", ) with open(viewer_path, "r", encoding="utf-8") as f: return HTMLResponse(f.read()) # ---- Example images -------------------------------------------------------- @app.get("/api/examples") async def get_examples(): """Return a JSON list of example images the frontend can display.""" return JSONResponse([ {"name": ex["name"], "url": f"/api/example/{i}"} for i, ex in enumerate(EXAMPLES) ]) @app.get("/api/example/{idx}") async def get_example(idx: int): """Serve an individual example image by index.""" if 0 <= idx < len(EXAMPLES): return FileResponse(EXAMPLES[idx]["file"], media_type="image/webp") return JSONResponse({"error": "not found"}, status_code=404) # ---------------------------------------------------------------------------- # Text-to-image endpoint (Ideogram 4 via remote Gradio client) # ---------------------------------------------------------------------------- @app.api() def generate_image( prompt: str, mode: str = "Default · 20 steps", width: int = 1024, height: int = 1024, seed: int = 0, randomize_seed: bool = True, ) -> tuple[FileData, str]: """Generate an image from a text prompt with Ideogram 4. The frontend feeds the resulting image into ``/generate`` to produce the 3D splat, giving an end-to-end text → image → 3D flow. Returns (image, seed_string); the frontend reads these as result.data[0..1]. """ if not prompt or not prompt.strip(): raise gr.Error("Enter a prompt to generate an image.") client = _get_ideogram_client() try: result = client.predict( prompt, # prompt mode, # mode "Ideogram (remote)", # upsampler int(width), # width int(height), # height int(seed), # seed bool(randomize_seed), # randomize_seed api_name="/generate", ) except Exception as err: # Surface the remote message (e.g. Ideogram's own ZeroGPU quota notice) # to the user instead of a generic 500. raise gr.Error(f"Ideogram 4 image generation failed: {err}") # Ideogram's /generate returns (image_path, used_seed, caption_json). image_path = result[0] if isinstance(result, (list, tuple)) else result used_seed = result[1] if isinstance(result, (list, tuple)) and len(result) > 1 else seed return FileData(path=str(image_path)), str(used_seed) # ---------------------------------------------------------------------------- # GPU pipeline helper # ---------------------------------------------------------------------------- @spaces.GPU def _run_pipeline(pil_image, seed, steps, guidance_scale, num_gaussians, out_dir, output_format, progress=None): """Run the full pipeline (preprocess → encode → sample → decode → save) in a single GPU acquisition. All file I/O happens here so the unpicklable Gaussian object never crosses the ZeroGPU multiprocessing boundary. ``progress`` is an optional ``gradio.Progress`` tracker. When supplied, the bar tracks the sampling loop only: stages before sampling (preprocess, encode) report 0% and stages after it (decode, save) report 100%, so the bar fills 0% → 100% across the sampling steps via a per-step callback. """ def _report(frac, desc): if progress is not None: progress(frac, desc=desc) t0 = time.time() _report(0.0, "Preprocessing image") prepared = PIPE.preprocess_image(pil_image) _report(0.0, "Encoding image") gen = torch.Generator(device=PIPE._device).manual_seed(int(seed)) cond = PIPE.encode_image(prepared, generator=gen) total_steps = int(steps) def _on_step(step, total): _report(step / total, f"Sampling · step {step}/{total}") out = PIPE.sample_latent( cond, steps=total_steps, guidance_scale=float(guidance_scale), generator=gen, show_progress=True, callback=_on_step, ) _report(1.0, "Decoding gaussians") gaussian = PIPE.decode_latent(out["latent"], num_gaussians=int(num_gaussians)) gen_dt = time.time() - t0 _report(1.0, "Saving output") # Save preprocessed image prep_path = out_dir / "preprocessed.png" prepared.save(str(prep_path)) # Save PLY (always needed for the viewer) ply_path = out_dir / "splat.ply" gaussian.save_ply(str(ply_path)) # Save in the requested download format fmt = output_format.lower() if fmt == "splat": download_path = out_dir / "splat.splat" gaussian.save_splat(str(download_path)) else: download_path = ply_path n_gaussians = gaussian.get_xyz.shape[0] _report(1.0, "Done") # Return only picklable primitives / paths return str(prep_path), str(ply_path), str(download_path), n_gaussians, gen_dt # ---------------------------------------------------------------------------- # Main API endpoint (queued via Gradio's engine) # ---------------------------------------------------------------------------- @app.api() def generate( image: FileData, seed: int = 42, steps: int = 20, guidance_scale: float = 3.0, num_gaussians: int = 262144, output_format: str = "ply", ) -> tuple[FileData, FileData, FileData, str]: """Generate 3D Gaussians from an input image. Returns (preprocessed_image, ply_file, download_file, info_string). The frontend receives these as result.data[0..3]. Sampling progress is streamed to the client over Gradio's SSE queue via a ``gr.Progress`` tracker. The tracker is created inside the function body (rather than declared as a parameter) because ``@app.api()`` derives the endpoint's input schema from the signature and would otherwise treat a ``progress`` parameter as a required API input. """ pil_image = Image.open(image["path"]).convert("RGBA") progress = gr.Progress() out_dir = OUT_ROOT / uuid4().hex[:12] out_dir.mkdir(parents=True, exist_ok=True) prep_path, ply_path, download_path, n_gaussians, gen_dt = _run_pipeline( pil_image, seed, steps, guidance_scale, num_gaussians, out_dir, output_format, progress=progress, ) info = ( f"{n_gaussians:,} gaussians · " f"generation: {gen_dt:.1f}s · saved: {Path(download_path).name}" ) return ( FileData(path=prep_path), FileData(path=ply_path), FileData(path=download_path), info, ) # ---------------------------------------------------------------------------- # Launch # ---------------------------------------------------------------------------- if __name__ == "__main__": app.launch(show_error=True)