"""TripoSplat – gradio.Server with custom frontend. Usage: python app.py """ import base64 import os import subprocess import tempfile import time from pathlib import Path from uuid import uuid4 import spaces import torch 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 triposplat import TripoSplatPipeline import example_inputs_b64 as _b64 # ---------------------------------------------------------------------------- # Download checkpoints from HuggingFace Hub (VAST-AI/TripoSplat) # ---------------------------------------------------------------------------- subprocess.run( [ "hf", "download", "VAST-AI/TripoSplat", "--local-dir", "ckpts" ], check=True, ) # ---------------------------------------------------------------------------- # 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")}, ] # ---------------------------------------------------------------------------- # 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) # ---------------------------------------------------------------------------- # 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)