"""RenderFormer-V2 — neural rendering of triangle-mesh scenes, on ZeroGPU.

RenderFormer replaces a path tracer with a transformer: it takes a triangle
soup plus per-triangle material/emission tokens and a camera, and directly
predicts the HDR image — global illumination, soft shadows, reflections and
refraction included, with no ray tracing at inference time.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
# No sm_120 flash-attn wheel matches the runtime torch ABI; SDPA is the
# repository's own supported fallback and is selected explicitly so the import
# never has to fail first.
os.environ.setdefault("ATTN_IMPL", "sdpa")

import spaces  # noqa: E402  (must precede any CUDA-touching import)

import json  # noqa: E402
import shutil  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
import traceback  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402

from renderformer.data import SceneData  # noqa: E402
import simple_exr  # noqa: E402  (the repository's own EXR writer)
from renderformer.pipelines import RenderFormerV2Pipeline  # noqa: E402

import scene_prep  # noqa: E402

MODEL_ID = "RenderFormer/renderformer-v2"
RESOLUTION = 512
HERE = Path(__file__).parent
WORK = Path(tempfile.gettempdir()) / "rf2-work"
WORK.mkdir(parents=True, exist_ok=True)


# --------------------------------------------------------------------------- #
# model
# --------------------------------------------------------------------------- #

PIPE = RenderFormerV2Pipeline.from_pretrained(
    MODEL_ID,
    resolution=RESOLUTION,
    device="cpu",
    apply_fused_kernels=True,
    # Both roles are the same Qwen-Image VAE; matching dtypes makes the
    # component registry hand back one shared instance instead of two copies.
    environment_encoder_dtype=torch.float16,
    texture_encoder_dtype=torch.float16,
)
PIPE._ensure_texture_vae()
PIPE.to("cuda")
print("[rf2] pipeline ready:", json.dumps(PIPE.component_info, default=str), flush=True)


# --------------------------------------------------------------------------- #
# bundled paper scenes
# --------------------------------------------------------------------------- #

SCENES = {
    "Transparent torus": {
        "dir": HERE / "scenes" / "transparent-torus",
        "hero": "Torus",
        "tone": "none",
    },
    "Cube pile": {
        "dir": HERE / "scenes" / "cube-pile",
        "hero": "object1",
        "tone": "none",
    },
    "Dragon": {
        "dir": HERE / "scenes" / "dragon",
        "hero": "dragon",
        "tone": "agx",
    },
}

PREPARED = {}
for _name, _spec in SCENES.items():
    _prep, _cfg = scene_prep.load_prepared_frame(
        _spec["dir"], WORK / _spec["dir"].name
    )
    PREPARED[_name] = (_prep, _cfg)
    print(
        f"[rf2] staged {_name}: "
        f"{sum(int(o.triangles.shape[0]) for o in _prep.objects)} triangles",
        flush=True,
    )

STUDIO_TEMPLATE = SCENES["Transparent torus"]["dir"]


# --------------------------------------------------------------------------- #
# material / tone-mapping helpers
# --------------------------------------------------------------------------- #

MATERIAL_MODES = ["Scene default", "Opaque (metallic/roughness)", "Glass (transmissive)"]
TONE_MAPPERS = ["Scene default", "none", "agx", "filmic", "pbr_neutral"]


def _srgb_hex_to_linear(value: str) -> list:
    text = (value or "").strip()
    if text.startswith("rgba") or text.startswith("rgb"):
        parts = text[text.find("(") + 1 : text.find(")")].split(",")
        srgb = [float(p) / 255.0 for p in parts[:3]]
    else:
        text = text.lstrip("#")
        if len(text) == 3:
            text = "".join(c * 2 for c in text)
        if len(text) != 6:
            raise gr.Error(f"could not parse colour {value!r}")
        srgb = [int(text[i : i + 2], 16) / 255.0 for i in (0, 2, 4)]
    return [
        c / 12.92 if c <= 0.04045 else ((c + 0.055) / 1.055) ** 2.4 for c in srgb
    ]


def _material_override(mode, colour, metallic, roughness, ior):
    """Build a MaterialConfig attribute patch for the hero object."""
    if mode == MATERIAL_MODES[0]:
        return None
    base_color = _srgb_hex_to_linear(colour)
    patch = {
        "texture_map_path": None,
        "measured_brdf_npy_path": None,
        "diffuse_specular_params": None,
        "metallic_roughness_params": None,
        "metallic_roughness_transmission_params": None,
    }
    if mode == MATERIAL_MODES[1]:
        patch["material_type"] = "homo_metallic_roughness"
        patch["metallic_roughness_params"] = {
            "base_color": base_color,
            "metallic": float(metallic),
            "roughness": float(roughness),
        }
    else:
        patch["material_type"] = "homo_metallic_roughness_transmission"
        patch["metallic_roughness_transmission_params"] = {
            "base_color": base_color,
            "metallic": float(metallic),
            "roughness": float(roughness),
            "transmission_weight": 1.0,
            "ior": float(ior),
        }
    return patch


def _tone_map(hdr: np.ndarray, name: str) -> np.ndarray:
    if name == "none":
        return np.clip(hdr, 0.0, 1.0)
    from simple_ocio import ToneMapper

    ocio_name = "Khronos PBR Neutral" if name == "pbr_neutral" else name
    return ToneMapper(ocio_name).hdr_to_ldr(hdr)


# --------------------------------------------------------------------------- #
# GPU worker
# --------------------------------------------------------------------------- #


def _prime_dynamo_handlers():
    """Make ``torch.compile`` usable inside a ZeroGPU worker.

    RenderFormer-V2's triangle transformer is built on ``flex_attention`` with a
    sliding-window + summary/register-sink block mask, and the repository wraps
    both ``flex_attention`` and ``create_block_mask`` in ``torch.compile`` — the
    eager fallback would materialise a dense NxN score matrix, which is not an
    option at 41k triangles.

    ZeroGPU's torch shim aliases ``torch.cuda._maybe_exchange_device`` onto
    ``torch.cuda._exchange_device``. Dynamo builds a ``function -> handler``
    table keyed by those objects and refuses a duplicate key, so the very first
    compile dies with ``AssertionError: Handler already registered for
    <built-in function _cuda_exchangeDevice>``. The table is ``functools.cache``d,
    so building it once with the alias broken fixes every later compile.
    """
    from torch._dynamo.variables.torch import TorchInGraphFunctionVariable

    build = TorchInGraphFunctionVariable._get_handlers
    info = getattr(build, "cache_info", None)
    if info is not None and info().currsize:
        return

    saved = torch.cuda._maybe_exchange_device
    if saved is torch.cuda._exchange_device:

        def _distinct_maybe_exchange_device(device):
            return saved(device)

        torch.cuda._maybe_exchange_device = _distinct_maybe_exchange_device
    try:
        build()
    except AssertionError as exc:  # pragma: no cover - diagnostic only
        print(f"[rf2] dynamo handler priming failed: {exc}", flush=True)
    finally:
        torch.cuda._maybe_exchange_device = saved


def _tame_inductor():
    """Stop Inductor from autotuning kernels by timing them with CUDA events.

    Inside a ZeroGPU worker the autotuner's timing loop raises
    ``RuntimeError: Both events must be completed before calculating elapsed
    time``, which aborts the compile. Picking the heuristic config instead of a
    benchmarked one costs a little kernel throughput and saves the compile, so
    turn the benchmarking off and keep a defensive fallback for any autotuner
    that still slips through.
    """
    from torch._inductor import config as inductor_config

    inductor_config.max_autotune = False
    inductor_config.max_autotune_pointwise = False
    inductor_config.max_autotune_gemm = False
    inductor_config.coordinate_descent_tuning = False
    inductor_config.benchmark_kernel = False
    inductor_config.triton.autotune_pointwise = False
    inductor_config.triton.autotune_cublasLt = False

    from torch._inductor.runtime import benchmarking

    cls = type(benchmarking.benchmarker)
    original = cls.benchmark_gpu
    if getattr(original, "_rf2_patched", False):
        return

    def benchmark_gpu(self, _callable, *args, **kwargs):
        try:
            return original(self, _callable, *args, **kwargs)
        except RuntimeError as exc:
            print(f"[rf2] inductor autotune timing unavailable: {exc}", flush=True)
            _callable()
            torch.cuda.synchronize()
            return 0.0

    benchmark_gpu._rf2_patched = True
    cls.benchmark_gpu = benchmark_gpu


def _install_compiled_flex():
    """Force flex attention through Inductor instead of its dense fallback.

    ``torch.nn.attention.flex_attention`` only emits the fused sparse kernel when
    it is *inlined into* a compiled region — ``torch.compiler.is_dynamo_compiling()``
    has to be true inside its own frame. The repository rebinds
    ``flex_attention = torch.compile(flex_attention)``, and in the ZeroGPU worker
    that decoration graph-breaks before the check, so every call lands in
    ``sdpa_dense`` and materialises a dense ``(B, H, N, N)`` score matrix. That is
    survivable at 4k triangles and about 54 GB at 41k, which is what killed the
    dragon scene with an allocator assert.

    Compiling a thin wrapper *around* the call instead keeps ``flex_attention``
    inside the traced graph, so the block-sparse Triton kernel is actually used.
    """
    from renderformer.models import attention as rf_attention

    if getattr(rf_attention, "_rf2_flex_shim", False):
        return

    import torch._dynamo as dynamo
    import torch.nn.attention.flex_attention as flex_module
    from torch.nn.attention.flex_attention import flex_attention as eager_flex

    dynamo.config.recompile_limit = max(getattr(dynamo.config, "recompile_limit", 8), 32)

    # ``flex_attention`` opens with a block of pure input validation, and Dynamo
    # cannot trace ``query.device != key.device`` there ("SourcelessBuilder.create
    # does not know how to wrap <class 'torch.device'>"). That graph break is what
    # drops the call out of the compiled region in the first place. The checks
    # guard against mismatched devices/dtypes/shapes that the pipeline has already
    # ruled out, so retire them for the duration of the render.
    # The same trace stops again on ``query.device != block_mask.kv_num_blocks.device``,
    # which is inline in ``flex_attention`` and cannot be patched out. Teach the
    # sourceless builder to wrap ``torch.device`` instead (newer PyTorch does this
    # itself). Best effort: if it does not take, the shim below still falls back.
    try:
        from torch._dynamo.variables import constant as dynamo_constant
        from torch._dynamo.variables.builder import SourcelessBuilder

        handlers = SourcelessBuilder._type_handlers
        if torch.device not in handlers:
            dynamo_constant.common_constant_types.add(torch.device)
            handlers[torch.device] = (
                lambda tx, value: dynamo_constant.ConstantVariable.create(value)
            )
    except Exception as exc:  # pragma: no cover - diagnostic only
        print(f"[rf2] could not teach dynamo about torch.device: {exc}", flush=True)

    def _skip_validation(*_args, **_kwargs):
        return None

    for name in (
        "_validate_sdpa_input",
        "_validate_device",
        "_validate_embed_dim",
        "_validate_no_nested_tensors",
        "_validate_block_mask_shape",
    ):
        if hasattr(flex_module, name):
            setattr(flex_module, name, _skip_validation)

    @torch.compile(fullgraph=True, dynamic=False)
    def compiled_flex(query, key, value, block_mask):
        return eager_flex(query, key, value, block_mask=block_mask)

    def flex(query, key, value, block_mask=None, **kwargs):
        if block_mask is not None and not kwargs:
            try:
                return compiled_flex(query, key, value, block_mask)
            except Exception as exc:  # pragma: no cover - diagnostic only
                print(f"[rf2] compiled flex attention unavailable: {exc}", flush=True)
        return eager_flex(query, key, value, block_mask=block_mask, **kwargs)

    rf_attention.flex_attention = flex
    rf_attention._rf2_flex_shim = True


def _gpu_duration(*_args, **_kwargs):
    """Measured on this Space: 1.6-6 s warm, ~86 s on the first render after a restart.

    Inductor caches the flex-attention kernels to disk, so only the first visitor
    after a reboot pays the compile; the reservation has to cover it anyway.
    """
    return 120


@spaces.GPU(duration=_gpu_duration)
def _gpu_render(inputs: "scene_prep.SceneInputs", c2w: np.ndarray, fov: np.ndarray):
    """Run the RenderFormer-V2 transformer on one prepared scene + one camera."""
    _prime_dynamo_handlers()
    _tame_inductor()
    _install_compiled_flex()
    started = time.perf_counter()

    raw = torch.from_numpy(inputs.raw_texture_unique)
    encoded = PIPE.encode_texture(raw, dtype=torch.float16)  # (O, 80, 4, 4)
    counts = torch.as_tensor(
        inputs.object_triangle_counts, device=encoded.device, dtype=torch.long
    )
    texture = encoded.repeat_interleave(counts, dim=0)  # (N, 80, 4, 4)

    scene = SceneData(
        path=Path(inputs.scene_name),
        triangles=torch.from_numpy(inputs.triangles),
        texture=texture,
        vn=torch.from_numpy(inputs.vn),
        c2w=torch.from_numpy(c2w),
        fov=torch.from_numpy(fov),
        mvp=torch.from_numpy(np.repeat(np.eye(4, dtype=np.float32)[None], 1, axis=0)),
        light_strength=torch.from_numpy(inputs.light_strength),
        input_metadata={
            "texture": "preencoded",
            "light_strength": "h5",
            "environment": "missing",
        },
    )

    out = PIPE.render_scene(scene, resolution=RESOLUTION, precision="fp16")
    hdr = out.hdr[0, 0].float().cpu().numpy()
    torch.cuda.synchronize()
    return np.ascontiguousarray(hdr), time.perf_counter() - started


# --------------------------------------------------------------------------- #
# shared render path (CPU prep + GPU call + tone mapping)
# --------------------------------------------------------------------------- #


def _render(prepared, camera, overrides, yaw, pitch, distance, tone_name, label):
    inputs = scene_prep.build_scene_inputs(prepared, material_overrides=overrides)
    position, target, up = scene_prep.orbit_camera(
        camera, float(yaw), float(pitch), float(distance)
    )
    c2w, fov, _ = scene_prep.camera_matrices(position, target, up, camera.fov)

    hdr, seconds = _gpu_render(inputs, c2w, fov)
    hdr = np.nan_to_num(np.asarray(hdr, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
    hdr = np.clip(hdr, 0.0, 1e4)

    ldr = (np.clip(_tone_map(hdr, tone_name), 0.0, 1.0) * 255.0).astype(np.uint8)

    exr_dir = Path(tempfile.mkdtemp(prefix="rf2-out-"))
    exr_path = exr_dir / f"{label}.exr"
    # The OpenCV wheels ship without the OpenEXR codec, so use the writer the
    # RenderFormer data pipeline itself uses.
    simple_exr.write_exr(str(exr_path), hdr)

    info = (
        f"**{inputs.num_triangles:,} triangles** · {RESOLUTION}x{RESOLUTION} · "
        f"tone mapper `{tone_name}` · **{seconds:.2f} s** on the GPU"
    )
    return ldr, info, str(exr_path)


# --------------------------------------------------------------------------- #
# handlers
# --------------------------------------------------------------------------- #


def render_paper_scene(
    scene_name: str,
    yaw: float = 0.0,
    pitch: float = 0.0,
    distance: float = 1.0,
    tone_mapper: str = "Scene default",
    material_mode: str = "Scene default",
    base_colour: str = "#c75a3c",
    metallic: float = 0.0,
    roughness: float = 0.3,
    ior: float = 1.33,
):
    """Render one of the bundled RenderFormer-V2 paper scenes.

    Args:
        scene_name: Which bundled scene to render.
        yaw: Orbit the camera left/right around the scene, in degrees.
        pitch: Orbit the camera up/down, in degrees.
        distance: Multiplier on the camera's distance to its look-at point.
        tone_mapper: HDR-to-LDR transform; "Scene default" uses the authored one.
        material_mode: Optionally replace the hero object's material.
        base_colour: Base colour of the replacement material (sRGB hex).
        metallic: Metallic parameter of the replacement material.
        roughness: Roughness parameter of the replacement material.
        ior: Index of refraction, used by the glass material only.

    Returns:
        The tone-mapped render, a short info line, and the linear HDR .exr file.
    """
    if scene_name not in SCENES:
        raise gr.Error(f"unknown scene {scene_name!r}")
    spec = SCENES[scene_name]
    prepared, config = PREPARED[scene_name]
    tone = spec["tone"] if tone_mapper == "Scene default" else tone_mapper

    patch = _material_override(material_mode, base_colour, metallic, roughness, ior)
    overrides = {spec["hero"]: patch} if patch else {}

    try:
        return _render(
            prepared,
            config.cameras[0],
            overrides,
            yaw,
            pitch,
            distance,
            tone,
            spec["dir"].name,
        )
    except gr.Error:
        raise
    except Exception as exc:  # surface a readable message in the UI
        traceback.print_exc()
        raise gr.Error(f"render failed: {type(exc).__name__}: {exc}") from exc


def render_uploaded_mesh(
    mesh_file,
    yaw: float = 0.0,
    pitch: float = 0.0,
    distance: float = 1.0,
    tone_mapper: str = "none",
    material_mode: str = "Opaque (metallic/roughness)",
    base_colour: str = "#c75a3c",
    metallic: float = 0.0,
    roughness: float = 0.3,
    ior: float = 1.33,
):
    """Drop a triangle mesh into a bundled studio set and render it.

    The studio (floor, two coloured walls, one area light) is lifted verbatim
    from the authors' "transparent torus" scene, so the lighting matches a
    configuration the model was demonstrated on. Meshes above ~14k triangles
    are quadric-decimated, then centred, scaled and seated on the floor.

    Args:
        mesh_file: Path to an .obj / .ply / .glb / .stl triangle mesh.
        yaw: Orbit the camera left/right around the studio, in degrees.
        pitch: Orbit the camera up/down, in degrees.
        distance: Multiplier on the camera's distance to its look-at point.
        tone_mapper: HDR-to-LDR transform.
        material_mode: Material assigned to the uploaded mesh.
        base_colour: Base colour of that material (sRGB hex).
        metallic: Metallic parameter of that material.
        roughness: Roughness parameter of that material.
        ior: Index of refraction, used by the glass material only.

    Returns:
        The tone-mapped render, a short info line, and the linear HDR .exr file.
    """
    if not mesh_file:
        raise gr.Error("upload a mesh first")
    path = mesh_file if isinstance(mesh_file, str) else getattr(mesh_file, "name", None)
    if not path or not Path(path).is_file():
        raise gr.Error("could not read the uploaded mesh")

    mode = material_mode if material_mode != "Scene default" else MATERIAL_MODES[1]
    patch = _material_override(mode, base_colour, metallic, roughness, ior)
    tone = "none" if tone_mapper == "Scene default" else tone_mapper

    out_dir = Path(tempfile.mkdtemp(prefix="rf2-studio-", dir=WORK))
    try:
        scene_prep.build_studio_frame(path, STUDIO_TEMPLATE, out_dir)
        prepared, config = scene_prep.load_prepared_frame(out_dir, out_dir)
        return _render(
            prepared,
            config.cameras[0],
            {"hero": patch} if patch else {},
            yaw,
            pitch,
            distance,
            tone,
            "studio",
        )
    except gr.Error:
        raise
    except Exception as exc:
        traceback.print_exc()
        raise gr.Error(f"render failed: {type(exc).__name__}: {exc}") from exc
    finally:
        shutil.rmtree(out_dir, ignore_errors=True)


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #

CSS = """
#col-container { max-width: 1180px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

HEADER = """
# RenderFormer&nbsp;V2 — neural rendering with heterogeneous scene primitives

A transformer that **replaces the path tracer**: give it a triangle soup with
per-triangle material and emission tokens plus a camera, and it directly
predicts the HDR image — global illumination, soft shadows, glossy reflections
and refraction included, with no ray tracing and no per-scene training.

[paper](https://huggingface.co/papers/2609.05738) ·
[model](https://huggingface.co/RenderFormer/renderformer-v2) ·
[code](https://github.com/iamNCJ/RenderFormer-Studio)
"""


def _camera_controls():
    yaw = gr.Slider(-180, 180, value=0, step=1, label="Orbit (yaw°)")
    pitch = gr.Slider(-60, 60, value=0, step=1, label="Elevation (pitch°)")
    distance = gr.Slider(0.5, 2.0, value=1.0, step=0.05, label="Distance ×")
    return yaw, pitch, distance


def _material_controls(default_mode: str):
    mode = gr.Radio(MATERIAL_MODES, value=default_mode, label="Material")
    colour = gr.ColorPicker(value="#c75a3c", label="Base colour")
    metallic = gr.Slider(0.0, 1.0, value=0.0, step=0.01, label="Metallic")
    roughness = gr.Slider(0.02, 1.0, value=0.3, step=0.01, label="Roughness")
    ior = gr.Slider(1.0, 2.5, value=1.33, step=0.01, label="Index of refraction (glass)")
    return mode, colour, metallic, roughness, ior


with gr.Blocks(title="RenderFormer V2") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(HEADER)

        with gr.Tabs():
            # ---------------------------------------------------------- #
            with gr.Tab("Paper scenes"):
                with gr.Row():
                    with gr.Column(scale=1):
                        p_scene = gr.Dropdown(
                            list(SCENES), value="Transparent torus", label="Scene"
                        )
                        p_yaw, p_pitch, p_distance = _camera_controls()
                        p_run = gr.Button("Render", variant="primary")
                        with gr.Accordion("Material & tone mapping", open=False):
                            p_tone = gr.Dropdown(
                                TONE_MAPPERS, value="Scene default", label="Tone mapper"
                            )
                            (
                                p_mode,
                                p_colour,
                                p_metallic,
                                p_roughness,
                                p_ior,
                            ) = _material_controls(MATERIAL_MODES[0])
                    with gr.Column(scale=1):
                        p_image = gr.Image(label="Render", type="numpy", height=512)
                        p_info = gr.Markdown()
                        p_exr = gr.File(label="Linear HDR (.exr)")

                p_inputs = [
                    p_scene,
                    p_yaw,
                    p_pitch,
                    p_distance,
                    p_tone,
                    p_mode,
                    p_colour,
                    p_metallic,
                    p_roughness,
                    p_ior,
                ]
                p_outputs = [p_image, p_info, p_exr]
                p_run.click(render_paper_scene, p_inputs, p_outputs, api_name="render_scene")

                gr.Examples(
                    examples=[
                        ["Transparent torus", 0, 0],
                        ["Cube pile", 25, 10],
                        ["Dragon", 0, 0],
                    ],
                    inputs=[p_scene, p_yaw, p_pitch],
                    outputs=p_outputs,
                    fn=render_paper_scene,
                    cache_examples=True,
                    cache_mode="lazy",
                )

            # ---------------------------------------------------------- #
            with gr.Tab("Your own mesh"):
                with gr.Row():
                    with gr.Column(scale=1):
                        u_mesh = gr.Model3D(label="Triangle mesh (.obj / .glb / .ply)")
                        u_yaw, u_pitch, u_distance = _camera_controls()
                        u_run = gr.Button("Render", variant="primary")
                        with gr.Accordion("Material & tone mapping", open=False):
                            u_tone = gr.Dropdown(
                                TONE_MAPPERS, value="none", label="Tone mapper"
                            )
                            (
                                u_mode,
                                u_colour,
                                u_metallic,
                                u_roughness,
                                u_ior,
                            ) = _material_controls(MATERIAL_MODES[1])
                    with gr.Column(scale=1):
                        u_image = gr.Image(label="Render", type="numpy", height=512)
                        u_info = gr.Markdown()
                        u_exr = gr.File(label="Linear HDR (.exr)")

                u_inputs = [
                    u_mesh,
                    u_yaw,
                    u_pitch,
                    u_distance,
                    u_tone,
                    u_mode,
                    u_colour,
                    u_metallic,
                    u_roughness,
                    u_ior,
                ]
                u_outputs = [u_image, u_info, u_exr]
                u_run.click(render_uploaded_mesh, u_inputs, u_outputs, api_name="render_mesh")

                gr.Examples(
                    examples=[[str(HERE / "assets" / "constant-width.obj")]],
                    inputs=[u_mesh],
                    outputs=u_outputs,
                    fn=render_uploaded_mesh,
                    cache_examples=True,
                    cache_mode="lazy",
                )

        gr.Markdown(
            "Scenes are built entirely on CPU (trimesh + the checkpoint's BRDF "
            "mappers), so no Blender install is needed — the transformer is the "
            "only thing that touches the GPU.\n\n"
            "**Assets** — *Transparent torus* and *Cube pile* are the authors' own "
            "scenes (MIT). *Dragon* is “Dragon” by Delatronic, CC&nbsp;BY&nbsp;3.0, "
            "via Benedikt Bitterli's Rendering Resources. The constant-width "
            "surface example is CC0 1.0 by Keenan Crane."
        )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True)