app.py

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