initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+static/example_inputs/
+gradio_outputs/
+ckpts/
+__pycache__/
+*.pyc

README.md

@@ -0,0 +1,14 @@
+---
+title: TripoSplat
+emoji: 👁
+colorFrom: purple
+colorTo: blue
+sdk: gradio
+sdk_version: 6.15.2
+python_version: '3.12'
+app_file: app.py
+pinned: false
+license: mit
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,198 @@
+"""TripoSplat Gradio demo with Spark.js in-browser viewer.
+Usage: python app.py
+"""
+import base64
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from uuid import uuid4
+
+import gradio as gr
+import spaces
+import torch
+
+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)
+VIEWER_HTML = Path("static/viewer/viewer.html").resolve()
+
+# Decode example images from base64 into a persistent temp directory so that
+# gr.Examples (which needs file paths) works without binary files in the repo.
+_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 = [
+    _write_example("CREATURE_BUTTERFLY",   "creature_butterfly.webp"),
+    _write_example("BUILDING_STONE_HOUSE", "building_stone_house.webp"),
+    _write_example("VEHICLE_PIRATE_SHIP",  "vehicle_pirate_ship.webp"),
+    _write_example("PLANT_WATER_LILY",     "plant_water_lily.webp"),
+]
+
+PLACEHOLDER_HTML = (
+    "<div style='display:flex;align-items:center;justify-content:center;height:520px;"
+    "color:#94a3b8;font:16px system-ui;background:#111318;border-radius:12px'>"
+    "3D viewer will appear here after generation</div>"
+)
+
+
+def _gr_file(path: Path) -> str:
+    """Gradio serves any file under `allowed_paths` at `/gradio_api/file=<abspath>`."""
+    return f"/gradio_api/file={path.as_posix()}"
+
+
+def _viewer_iframe(ply_path: Path) -> str:
+    ts = time.time()  # cache-bust so the iframe reloads each generation
+    src = f"{_gr_file(VIEWER_HTML)}?ply={_gr_file(ply_path)}&ts={ts}"
+    return (
+        f"<iframe src='{src}' "
+        "style='width:100%;height:520px;border:0;border-radius:12px;background:#0a0b0e'></iframe>"
+    )
+
+
+# ----------------------------------------------------------------------------
+# Event handlers
+# ----------------------------------------------------------------------------
+
+def on_image_change(image):
+    """Run preprocessing as soon as the input changes — gives the user instant
+    feedback on the matte/crop without waiting for the full generation."""
+    if image is None:
+        return None
+    return PIPE.preprocess_image(image)
+
+
+@spaces.GPU
+def generate(prepared, seed: int, steps: int, guidance_scale: float,
+             num_gaussians: int, output_format: str,
+             progress=gr.Progress(track_tqdm=True)):
+    if prepared is None:
+        raise gr.Error("Please upload an image and wait for preprocessing to finish.")
+
+    progress(0, desc="Generating...")
+    t0 = time.time()
+    gen = torch.Generator(device=PIPE._device).manual_seed(int(seed))
+    cond = PIPE.encode_image(prepared, generator=gen)
+    out  = PIPE.sample_latent(cond, steps=int(steps),
+                              guidance_scale=float(guidance_scale),
+                              generator=gen, show_progress=True)
+    gaussian = PIPE.decode_latent(out["latent"], num_gaussians=int(num_gaussians))
+    gen_dt = time.time() - t0
+
+    out_dir = OUT_ROOT / uuid4().hex[:12]
+    out_dir.mkdir(parents=True, exist_ok=True)
+    ply_path = out_dir / "splat.ply"
+    gaussian.save_ply(str(ply_path))
+
+    fmt = output_format.lower()
+    if fmt == "ply":
+        download_path = ply_path
+    elif fmt == "splat":
+        download_path = out_dir / "splat.splat"
+        gaussian.save_splat(str(download_path))
+    else:
+        raise gr.Error(f"Unknown output format: {output_format}")
+
+    info = (f"{gaussian.get_xyz.shape[0]:,} gaussians  ·  "
+            f"generation: {gen_dt:.1f}s  ·  saved: {download_path.name}")
+    return _viewer_iframe(ply_path), gr.update(value=str(download_path), interactive=True), info
+
+
+# ----------------------------------------------------------------------------
+# Gradio UI
+# ----------------------------------------------------------------------------
+
+with gr.Blocks(title="TripoSplat") as demo:
+    gr.Markdown("# TripoSplat")
+    gr.Markdown(
+        "TripoSplat converts a single 2D image into high-quality and variable number of 3D Gaussians, developed by [TripoAI](https://www.tripo3d.ai/). "
+        "It can serve as a powerful pipeline tool for asset creation, AR/VR, game development, simulation environments, and beyond.\n\n"
+        "[Read Paper](https://arxiv.org/abs/2605.16355) | [Research Blog](https://www.tripo3d.ai/research/triposplat)"
+    )
+
+    image_in = gr.Image(label="Input image", type="pil", image_mode="RGBA",
+                        height=320, render=False)
+
+    gr.Examples(
+        examples=[[p] for p in EXAMPLES],
+        inputs=[image_in],
+        label="Examples (click to load)",
+        examples_per_page=10,
+        cache_examples=False,
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_in.render()
+
+            with gr.Accordion("Sampling settings", open=False):
+                seed_in = gr.Number(label="Seed", value=42, precision=0)
+                steps_in = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=20)
+                cfg_in = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, step=0.5, value=3.0)
+                num_g_in = gr.Dropdown(
+                    label="Number of gaussians",
+                    choices=["32768", "65536", "131072", "262144"],
+                    value="262144",
+                )
+                fmt_in = gr.Dropdown(label="Download format", choices=["ply", "splat"], value="ply")
+
+            run_btn = gr.Button("Generate", variant="primary")
+            prepared_out = gr.Image(label="Preprocessed input", interactive=False, height=240)
+            info_out = gr.Markdown()
+
+        with gr.Column(scale=2):
+            viewer_out = gr.HTML(value=PLACEHOLDER_HTML, label="Spark.js viewer")
+            file_out = gr.DownloadButton(label="Download", value=None, interactive=False)
+
+    image_in.change(
+        fn=on_image_change,
+        inputs=[image_in],
+        outputs=[prepared_out],
+    )
+
+    run_btn.click(
+        fn=generate,
+        inputs=[prepared_out, seed_in, steps_in, cfg_in, num_g_in, fmt_in],
+        outputs=[viewer_out, file_out, info_out],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch(
+        allowed_paths=[
+            str(VIEWER_HTML.parent),
+            str(OUT_ROOT),
+            _EXAMPLES_TMPDIR,
+        ],
+    )

model.py

@@ -0,0 +1,1725 @@
+from typing import Optional
+import math
+import re
+
+import numpy as np
+import safetensors.torch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.ops import deform_conv2d
+
+
+# ---------------------------------------------------------------------------
+# DINOv3 ViT-H/16+
+# ---------------------------------------------------------------------------
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class DinoV3PatchEmbed(nn.Module):
+    def __init__(self, patch_size=16, in_chans=3, embed_dim=1280):
+        super().__init__()
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True)
+
+    def forward(self, x):
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class DinoV3RotaryEmbedding2D(nn.Module):
+    def __init__(self, dim: int, base: float = 100.0):
+        super().__init__()
+        inv_freq = 1.0 / (base ** torch.arange(0, 1, 4.0 / dim, dtype=torch.float32))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def forward(self, height: int, width: int, device: torch.device, dtype: torch.dtype):
+        coords_h = torch.arange(0.5, height, dtype=torch.float32, device=device) / height
+        coords_w = torch.arange(0.5, width, dtype=torch.float32, device=device) / width
+        coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1)
+        coords = (2.0 * coords - 1.0).flatten(0, 1)
+        angles = (2 * math.pi * coords[:, :, None] * self.inv_freq[None, None, :]).flatten(1, 2).tile(2)
+        cos = angles.cos().unsqueeze(0).unsqueeze(0)
+        sin = angles.sin().unsqueeze(0).unsqueeze(0)
+        return cos.to(dtype=dtype), sin.to(dtype=dtype)
+
+
+class DinoV3Attention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, qkv_bias: tuple = (True, False, True)):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        q_bias, k_bias, v_bias = qkv_bias
+        self.q_proj = nn.Linear(dim, dim, bias=q_bias)
+        self.k_proj = nn.Linear(dim, dim, bias=k_bias)
+        self.v_proj = nn.Linear(dim, dim, bias=v_bias)
+        self.o_proj = nn.Linear(dim, dim, bias=True)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+                num_prefix_tokens: int = 0) -> torch.Tensor:
+        B, N, C = x.shape
+        q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2)
+        if num_prefix_tokens > 0:
+            q_pre, q_pat = q.split((num_prefix_tokens, N - num_prefix_tokens), dim=-2)
+            k_pre, k_pat = k.split((num_prefix_tokens, N - num_prefix_tokens), dim=-2)
+            q = torch.cat((q_pre, q_pat * cos + _rotate_half(q_pat) * sin), dim=-2)
+            k = torch.cat((k_pre, k_pat * cos + _rotate_half(k_pat) * sin), dim=-2)
+        else:
+            q = q * cos + _rotate_half(q) * sin
+            k = k * cos + _rotate_half(k) * sin
+        out = F.scaled_dot_product_attention(q, k, v)
+        return self.o_proj(out.transpose(1, 2).reshape(B, N, C))
+
+
+class DinoV3MLP(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, bias: bool = True):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DinoV3Block(nn.Module):
+    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0,
+                 qkv_bias: tuple = (True, False, True), layerscale_init: float = 1.0,
+                 mlp_bias: bool = True, eps: float = 1e-5):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, eps=eps)
+        self.attn = DinoV3Attention(dim, num_heads, qkv_bias=qkv_bias)
+        self.ls1 = nn.Parameter(torch.ones(dim) * layerscale_init)
+        self.norm2 = nn.LayerNorm(dim, eps=eps)
+        self.mlp = DinoV3MLP(dim, int(dim * mlp_ratio), bias=mlp_bias)
+        self.ls2 = nn.Parameter(torch.ones(dim) * layerscale_init)
+
+    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+                num_prefix_tokens: int = 0) -> torch.Tensor:
+        x = x + self.ls1 * self.attn(self.norm1(x), cos, sin, num_prefix_tokens=num_prefix_tokens)
+        x = x + self.ls2 * self.mlp(self.norm2(x))
+        return x
+
+
+class DinoV3ViT(nn.Module):
+    def __init__(self, hidden_size: int = 1280, num_heads: int = 20, num_layers: int = 32,
+                 patch_size: int = 16, num_register_tokens: int = 4,
+                 intermediate_size: int = 5120, layerscale_init: float = 1.0,
+                 query_bias: bool = True, key_bias: bool = False, value_bias: bool = True,
+                 mlp_bias: bool = True, rope_theta: float = 100.0, layer_norm_eps: float = 1e-5):
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.patch_embed = DinoV3PatchEmbed(patch_size=patch_size, embed_dim=hidden_size)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, hidden_size))
+        self.register_tokens = nn.Parameter(torch.zeros(1, num_register_tokens, hidden_size))
+        self.rope = DinoV3RotaryEmbedding2D(dim=hidden_size // num_heads, base=rope_theta)
+        qkv_bias = (query_bias, key_bias, value_bias)
+        self.blocks = nn.ModuleList([
+            DinoV3Block(hidden_size, num_heads, mlp_ratio=intermediate_size / hidden_size,
+                        qkv_bias=qkv_bias, layerscale_init=layerscale_init,
+                        mlp_bias=mlp_bias, eps=layer_norm_eps)
+            for _ in range(num_layers)
+        ])
+        self.norm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+
+    @property
+    def device(self) -> torch.device:
+        return self.cls_token.device
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        B, _, H, W = pixel_values.shape
+        x = self.patch_embed(pixel_values)
+        hp, wp = H // self.patch_size, W // self.patch_size
+        cos, sin = self.rope(hp, wp, x.device, x.dtype)
+        x = torch.cat([self.cls_token.expand(B, -1, -1),
+                        self.register_tokens.expand(B, -1, -1), x], dim=1)
+        num_prefix = 1 + self.num_register_tokens
+        for block in self.blocks:
+            x = block(x, cos, sin, num_prefix_tokens=num_prefix)
+        return self.norm(x)
+
+    def load_safetensors(self, path: str) -> None:
+        state_dict = safetensors.torch.load_file(path)
+        our_sd = self.state_dict()
+        loaded = {}
+        for hf_key in state_dict:
+            k = (hf_key
+                 .replace("embeddings.patch_embeddings.", "patch_embed.proj.")
+                 .replace("embeddings.cls_token", "cls_token")
+                 .replace("embeddings.mask_token", "mask_token")
+                 .replace("embeddings.register_tokens", "register_tokens"))
+            m = re.match(r"layer\.(\d+)\.(.+)", k)
+            if m:
+                rest = m.group(2)
+                for proj in ["q_proj", "k_proj", "v_proj", "o_proj"]:
+                    rest = rest.replace(f"attention.{proj}", f"attn.{proj}")
+                rest = (rest.replace("layer_scale1.lambda1", "ls1")
+                            .replace("layer_scale2.lambda1", "ls2"))
+                k = f"blocks.{m.group(1)}.{rest}"
+            if k in our_sd:
+                assert state_dict[hf_key].shape == our_sd[k].shape, \
+                    f"Shape mismatch {k}: {state_dict[hf_key].shape} vs {our_sd[k].shape}"
+                loaded[k] = state_dict[hf_key]
+        check_sd = {k: v for k, v in our_sd.items() if k != "mask_token"}
+        missing = set(check_sd) - set(loaded)
+        unexpected = set(loaded) - set(check_sd)
+        if missing:
+            raise KeyError(f"[DINOv3] Missing keys: {missing}")
+        if unexpected:
+            raise KeyError(f"[DINOv3] Unexpected keys: {unexpected}")
+        self.load_state_dict(loaded, strict=True)
+
+
+# ---------------------------------------------------------------------------
+# Flux2 VAE Encoder
+# ---------------------------------------------------------------------------
+
+class Flux2ResnetBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, use_shortcut=False):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(32, in_channels, eps=1e-6)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, 1)
+        self.norm2 = nn.GroupNorm(32, out_channels, eps=1e-6)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1)
+        self.conv_shortcut = nn.Conv2d(in_channels, out_channels, 1, 1, 0) if use_shortcut else None
+
+    def forward(self, x):
+        h = F.silu(self.norm1(x))
+        h = F.silu(self.norm2(self.conv1(h)))
+        h = self.conv2(h)
+        return h + (self.conv_shortcut(x) if self.conv_shortcut is not None else x)
+
+
+class Flux2Downsampler(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 3, 2, 0)
+
+    def forward(self, x):
+        return self.conv(F.pad(x, (0, 1, 0, 1)))
+
+
+class Flux2Attention(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.group_norm = nn.GroupNorm(32, channels, eps=1e-6)
+        self.to_q = nn.Linear(channels, channels)
+        self.to_k = nn.Linear(channels, channels)
+        self.to_v = nn.Linear(channels, channels)
+        self.to_out = nn.ModuleList([nn.Linear(channels, channels), nn.Identity()])
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        h = self.group_norm(x).reshape(B, C, H * W).transpose(1, 2)
+        q = self.to_q(h).reshape(B, -1, 1, C).permute(0, 2, 1, 3)
+        k = self.to_k(h).reshape(B, -1, 1, C).permute(0, 2, 1, 3)
+        v = self.to_v(h).reshape(B, -1, 1, C).permute(0, 2, 1, 3)
+        out = F.scaled_dot_product_attention(q, k, v)
+        out = self.to_out[0](out.permute(0, 2, 1, 3).reshape(B, -1, C))
+        return x + out.transpose(1, 2).reshape(B, C, H, W)
+
+
+class Flux2Encoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv_in = nn.Conv2d(3, 128, 3, 1, 1)
+        self.down_0_resnets = nn.ModuleList([Flux2ResnetBlock(128, 128), Flux2ResnetBlock(128, 128)])
+        self.down_0_sampler = Flux2Downsampler(128)
+        self.down_1_resnets = nn.ModuleList([Flux2ResnetBlock(128, 256, use_shortcut=True), Flux2ResnetBlock(256, 256)])
+        self.down_1_sampler = Flux2Downsampler(256)
+        self.down_2_resnets = nn.ModuleList([Flux2ResnetBlock(256, 512, use_shortcut=True), Flux2ResnetBlock(512, 512)])
+        self.down_2_sampler = Flux2Downsampler(512)
+        self.down_3_resnets = nn.ModuleList([Flux2ResnetBlock(512, 512), Flux2ResnetBlock(512, 512)])
+        self.mid_attn = Flux2Attention(512)
+        self.mid_resnets = nn.ModuleList([Flux2ResnetBlock(512, 512), Flux2ResnetBlock(512, 512)])
+        self.conv_norm_out = nn.GroupNorm(32, 512, eps=1e-6)
+        self.conv_out = nn.Conv2d(512, 64, 3, 1, 1)
+
+    def forward(self, x):
+        x = self.conv_in(x)
+        for r in self.down_0_resnets: x = r(x)
+        x = self.down_0_sampler(x)
+        for r in self.down_1_resnets: x = r(x)
+        x = self.down_1_sampler(x)
+        for r in self.down_2_resnets: x = r(x)
+        x = self.down_2_sampler(x)
+        for r in self.down_3_resnets: x = r(x)
+        x = self.mid_resnets[0](x)
+        x = self.mid_attn(x)
+        x = self.mid_resnets[1](x)
+        return self.conv_out(F.silu(self.conv_norm_out(x)))
+
+
+class Flux2VAEEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = Flux2Encoder()
+        self.quant_conv = nn.Conv2d(64, 64, 1, 1, 0)
+        self.bn = nn.BatchNorm1d(128, eps=1e-5, momentum=0.1, affine=False, track_running_stats=True)
+
+    def load_safetensors(self, path: str):
+        sd = safetensors.torch.load_file(path)
+        remapped = {}
+        for k, v in sd.items():
+            # Skip the decoder half of a full Flux2-VAE ckpt — we only need the encoder.
+            if k.startswith(("decoder.", "post_quant_conv.")):
+                continue
+            # Comfy / diffusers-style naming → our flattened naming.
+            m = re.match(r"encoder\.down_blocks\.(\d+)\.resnets\.(\d+)\.(.+)", k)
+            if m:
+                remapped[f"encoder.down_{m.group(1)}_resnets.{m.group(2)}.{m.group(3)}"] = v
+                continue
+            m = re.match(r"encoder\.down_blocks\.(\d+)\.downsamplers\.0\.(.+)", k)
+            if m:
+                remapped[f"encoder.down_{m.group(1)}_sampler.{m.group(2)}"] = v
+                continue
+            m = re.match(r"encoder\.mid_block\.resnets\.(\d+)\.(.+)", k)
+            if m:
+                remapped[f"encoder.mid_resnets.{m.group(1)}.{m.group(2)}"] = v
+                continue
+            m = re.match(r"encoder\.mid_block\.attentions\.0\.(.+)", k)
+            if m:
+                remapped[f"encoder.mid_attn.{m.group(1)}"] = v
+                continue
+            remapped[k] = v
+        missing, unexpected = self.load_state_dict(remapped, strict=False)
+        if missing:
+            raise KeyError(f"[VAE] Missing keys: {missing}")
+        if unexpected:
+            raise KeyError(f"[VAE] Unexpected keys: {unexpected}")
+
+    def encode(self, images, deterministic: bool = True, generator: torch.Generator = None):
+        moments = self.quant_conv(self.encoder(images))
+        mean, logvar = moments.chunk(2, dim=1)
+        if deterministic:
+            latents = mean
+        else:
+            noise = torch.randn(mean.shape, dtype=mean.dtype, device=mean.device, generator=generator)
+            latents = mean + torch.exp(0.5 * logvar) * noise
+        B, C, H, W = latents.shape
+        latents = latents.view(B, C, H // 2, 2, W // 2, 2).permute(0, 1, 3, 5, 2, 4)
+        latents = latents.reshape(B, C * 4, H // 2, W // 2)
+        bn_mean = self.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+        bn_std = torch.sqrt(self.bn.running_var.view(1, -1, 1, 1) + self.bn.eps).to(latents.device, latents.dtype)
+        return ((latents - bn_mean) / bn_std).to(torch.float32).flatten(2).transpose(1, 2).contiguous()
+
+        return rgba
+
+
+# ---------------------------------------------------------------------------
+# BiRefNet background removal (Swin-L + ASPP-deformable decoder)
+# ---------------------------------------------------------------------------
+
+# -- timm-style helpers, inlined to avoid a timm dependency --------------------
+
+def _trunc_normal_(tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0,
+                   a: float = -2.0, b: float = 2.0) -> torch.Tensor:
+    # Initialization helper — only used at __init__ time, the released ckpt
+    # overwrites everything in load_safetensors so the exact distribution here
+    # is unimportant.
+    with torch.no_grad():
+        tensor.normal_(mean, std).clamp_(mean + a * std, mean + b * std)
+    return tensor
+
+
+# -- Swin Transformer (Swin-Large preset) -------------------------------------
+
+class _SwinMlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None):
+        super().__init__()
+        hidden_features = hidden_features or in_features
+        out_features = out_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+
+    def forward(self, x):
+        return self.fc2(self.act(self.fc1(x)))
+
+
+def _window_partition(x, window_size):
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+
+
+def _window_reverse(windows, window_size, H, W):
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    return x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+
+
+class _WindowAttention(nn.Module):
+    def __init__(self, dim, window_size, num_heads):
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # (Wh, Ww)
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))
+        coords_h = torch.arange(window_size[0])
+        coords_w = torch.arange(window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))
+        coords_flatten = torch.flatten(coords, 1)
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+        relative_coords[:, :, 0] += window_size[0] - 1
+        relative_coords[:, :, 1] += window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+        self.register_buffer("relative_position_index", relative_coords.sum(-1))
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+        _trunc_normal_(self.relative_position_bias_table, std=0.02)
+
+    def forward(self, x, mask=None):
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
+        bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+            self.window_size[0] * self.window_size[1],
+            self.window_size[0] * self.window_size[1], -1)
+        attn = attn + bias.permute(2, 0, 1).contiguous().unsqueeze(0)
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+        attn = attn.softmax(dim=-1)
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        return self.proj(x)
+
+
+class _SwinBlock(nn.Module):
+    def __init__(self, dim, num_heads, window_size, shift_size, mlp_ratio=4.0):
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = _WindowAttention(dim, (window_size, window_size), num_heads)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = _SwinMlp(dim, int(dim * mlp_ratio))
+        self.H = None
+        self.W = None
+
+    def forward(self, x, mask_matrix):
+        B, L, C = x.shape
+        H, W = self.H, self.W
+        shortcut = x
+        x = self.norm1(x).view(B, H, W, C)
+        pad_r = (self.window_size - W % self.window_size) % self.window_size
+        pad_b = (self.window_size - H % self.window_size) % self.window_size
+        x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
+        _, Hp, Wp, _ = x.shape
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            attn_mask = mask_matrix
+        else:
+            shifted_x = x
+            attn_mask = None
+        x_windows = _window_partition(shifted_x, self.window_size).view(
+            -1, self.window_size * self.window_size, C)
+        attn_windows = self.attn(x_windows, mask=attn_mask).view(
+            -1, self.window_size, self.window_size, C)
+        shifted_x = _window_reverse(attn_windows, self.window_size, Hp, Wp)
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+        if pad_r > 0 or pad_b > 0:
+            x = x[:, :H, :W, :].contiguous()
+        x = x.view(B, H * W, C)
+        x = shortcut + x
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class _PatchMerging(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = nn.LayerNorm(4 * dim)
+
+    def forward(self, x, H, W):
+        B, L, C = x.shape
+        x = x.view(B, H, W, C)
+        if H % 2 == 1 or W % 2 == 1:
+            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
+        x0 = x[:, 0::2, 0::2, :]
+        x1 = x[:, 1::2, 0::2, :]
+        x2 = x[:, 0::2, 1::2, :]
+        x3 = x[:, 1::2, 1::2, :]
+        x = torch.cat([x0, x1, x2, x3], -1).view(B, -1, 4 * C)
+        return self.reduction(self.norm(x))
+
+
+class _SwinBasicLayer(nn.Module):
+    def __init__(self, dim, depth, num_heads, window_size, mlp_ratio=4.0, downsample=True):
+        super().__init__()
+        self.window_size = window_size
+        self.shift_size = window_size // 2
+        self.depth = depth
+        self.blocks = nn.ModuleList([
+            _SwinBlock(dim=dim, num_heads=num_heads, window_size=window_size,
+                       shift_size=0 if (i % 2 == 0) else window_size // 2,
+                       mlp_ratio=mlp_ratio)
+            for i in range(depth)
+        ])
+        self.downsample = _PatchMerging(dim) if downsample else None
+
+    def forward(self, x, H, W):
+        Hp = int(math.ceil(H / self.window_size)) * self.window_size
+        Wp = int(math.ceil(W / self.window_size)) * self.window_size
+        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)
+        h_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        w_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+        mask_windows = _window_partition(img_mask, self.window_size).view(-1, self.window_size ** 2)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)) \
+                             .masked_fill(attn_mask == 0, float(0.0)).to(x.dtype)
+        for blk in self.blocks:
+            blk.H, blk.W = H, W
+            x = blk(x, attn_mask)
+        if self.downsample is not None:
+            x_down = self.downsample(x, H, W)
+            Wh, Ww = (H + 1) // 2, (W + 1) // 2
+            return x, H, W, x_down, Wh, Ww
+        return x, H, W, x, H, W
+
+
+class _SwinPatchEmbed(nn.Module):
+    def __init__(self, patch_size=4, in_channels=3, embed_dim=192):
+        super().__init__()
+        self.patch_size = (patch_size, patch_size)
+        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
+        self.norm = nn.LayerNorm(embed_dim)
+        self.embed_dim = embed_dim
+
+    def forward(self, x):
+        _, _, H, W = x.shape
+        if W % self.patch_size[1] != 0:
+            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
+        if H % self.patch_size[0] != 0:
+            x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
+        x = self.proj(x)
+        Wh, Ww = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)
+        x = self.norm(x)
+        return x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
+
+
+class _SwinLarge(nn.Module):
+    """Swin-Large backbone matching the BiRefNet HF release.
+
+    embed_dim=192, depths=[2,2,18,2], num_heads=[6,12,24,48], window_size=12.
+    """
+    def __init__(self):
+        super().__init__()
+        embed_dim = 192
+        depths = [2, 2, 18, 2]
+        num_heads = [6, 12, 24, 48]
+        window_size = 12
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.patch_embed = _SwinPatchEmbed(patch_size=4, in_channels=3, embed_dim=embed_dim)
+        self.layers = nn.ModuleList([
+            _SwinBasicLayer(
+                dim=int(embed_dim * 2 ** i),
+                depth=depths[i],
+                num_heads=num_heads[i],
+                window_size=window_size,
+                downsample=(i < self.num_layers - 1),
+            ) for i in range(self.num_layers)
+        ])
+        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
+        self.num_features = num_features
+        for i in range(self.num_layers):
+            self.add_module(f"norm{i}", nn.LayerNorm(num_features[i]))
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+        Wh, Ww = x.size(2), x.size(3)
+        x = x.flatten(2).transpose(1, 2)
+        outs = []
+        for i in range(self.num_layers):
+            x_out, H, W, x, Wh, Ww = self.layers[i](x, Wh, Ww)
+            norm_layer = getattr(self, f"norm{i}")
+            x_out = norm_layer(x_out)
+            out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
+            outs.append(out)
+        return tuple(outs)
+
+
+# -- ASPP-Deformable -----------------------------------------------------------
+
+class _DeformableConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False):
+        super().__init__()
+        if isinstance(kernel_size, int):
+            kernel_size = (kernel_size, kernel_size)
+        self.stride = (stride, stride) if isinstance(stride, int) else stride
+        self.padding = padding
+        self.offset_conv = nn.Conv2d(in_channels, 2 * kernel_size[0] * kernel_size[1],
+                                     kernel_size=kernel_size, stride=stride, padding=padding, bias=True)
+        self.modulator_conv = nn.Conv2d(in_channels, 1 * kernel_size[0] * kernel_size[1],
+                                        kernel_size=kernel_size, stride=stride, padding=padding, bias=True)
+        self.regular_conv = nn.Conv2d(in_channels, out_channels,
+                                      kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)
+
+    def forward(self, x):
+        offset = self.offset_conv(x)
+        modulator = 2.0 * torch.sigmoid(self.modulator_conv(x))
+        return deform_conv2d(
+            input=x, offset=offset,
+            weight=self.regular_conv.weight, bias=self.regular_conv.bias,
+            padding=self.padding, mask=modulator, stride=self.stride,
+        )
+
+
+class _ASPPModuleDeformable(nn.Module):
+    def __init__(self, in_channels, planes, kernel_size, padding):
+        super().__init__()
+        self.atrous_conv = _DeformableConv2d(in_channels, planes, kernel_size=kernel_size,
+                                             stride=1, padding=padding, bias=False)
+        self.bn = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        return self.relu(self.bn(self.atrous_conv(x)))
+
+
+class _ASPPDeformable(nn.Module):
+    def __init__(self, in_channels, out_channels=None, parallel_block_sizes=(1, 3, 7)):
+        super().__init__()
+        if out_channels is None:
+            out_channels = in_channels
+        inter = 256
+        self.aspp1 = _ASPPModuleDeformable(in_channels, inter, 1, padding=0)
+        self.aspp_deforms = nn.ModuleList([
+            _ASPPModuleDeformable(in_channels, inter, k, padding=k // 2)
+            for k in parallel_block_sizes
+        ])
+        self.global_avg_pool = nn.Sequential(
+            nn.AdaptiveAvgPool2d((1, 1)),
+            nn.Conv2d(in_channels, inter, 1, stride=1, bias=False),
+            nn.BatchNorm2d(inter),
+            nn.ReLU(inplace=True),
+        )
+        self.conv1 = nn.Conv2d(inter * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x1 = self.aspp1(x)
+        x_aspp_deforms = [m(x) for m in self.aspp_deforms]
+        x5 = self.global_avg_pool(x)
+        x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True)
+        y = torch.cat((x1, *x_aspp_deforms, x5), dim=1)
+        return self.relu(self.bn1(self.conv1(y)))
+
+
+# -- Decoder blocks ------------------------------------------------------------
+
+class _BasicDecBlk(nn.Module):
+    def __init__(self, in_channels, out_channels, inter_channels=64):
+        super().__init__()
+        self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1)
+        self.bn_in = nn.BatchNorm2d(inter_channels)
+        self.relu_in = nn.ReLU(inplace=True)
+        self.dec_att = _ASPPDeformable(in_channels=inter_channels)
+        self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1)
+        self.bn_out = nn.BatchNorm2d(out_channels)
+
+    def forward(self, x):
+        x = self.relu_in(self.bn_in(self.conv_in(x)))
+        x = self.dec_att(x)
+        x = self.bn_out(self.conv_out(x))
+        return x
+
+
+class _BasicLatBlk(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class _SimpleConvs(nn.Module):
+    def __init__(self, in_channels, out_channels, inter_channels=64):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1)
+        self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1)
+
+    def forward(self, x):
+        return self.conv_out(self.conv1(x))
+
+
+# -- Image → patch-stack helper -----------------------------------------------
+
+def _image2patches(image, patch_ref):
+    """`einops` rearrange 'b c (hg h) (wg w) -> b (c hg wg) h w' replacement.
+
+    Splits `image` into hg×wg non-overlapping patches and stacks them along
+    the channel axis. `hg`/`wg` are inferred from image and patch_ref sizes.
+    """
+    b, c, h_full, w_full = image.shape
+    hg, wg = h_full // patch_ref.shape[-2], w_full // patch_ref.shape[-1]
+    h, w = h_full // hg, w_full // wg
+    # (b, c, hg*h, wg*w) -> (b, c, hg, h, wg, w) -> (b, c, hg, wg, h, w) -> (b, c*hg*wg, h, w)
+    return image.view(b, c, hg, h, wg, w).permute(0, 1, 2, 4, 3, 5).reshape(b, c * hg * wg, h, w)
+
+
+# -- Decoder + top-level BiRefNet ---------------------------------------------
+
+class _BiRefNetDecoder(nn.Module):
+    def __init__(self, channels=(3072, 1536, 768, 384)):
+        super().__init__()
+        c = channels  # high-to-low resolution channel counts
+        # input-modulator blocks (one per resolution; channels are
+        # `3 * patch_grid**2`, see _image2patches docstring).
+        self.ipt_blk5 = _SimpleConvs(2 ** 10 * 3, c[0] // 8, inter_channels=64)
+        self.ipt_blk4 = _SimpleConvs(2 ** 8 * 3,  c[0] // 8, inter_channels=64)
+        self.ipt_blk3 = _SimpleConvs(2 ** 6 * 3,  c[1] // 8, inter_channels=64)
+        self.ipt_blk2 = _SimpleConvs(2 ** 4 * 3,  c[2] // 8, inter_channels=64)
+        self.ipt_blk1 = _SimpleConvs(2 ** 0 * 3,  c[3] // 8, inter_channels=64)
+
+        self.decoder_block4 = _BasicDecBlk(c[0] + c[0] // 8, c[1])
+        self.decoder_block3 = _BasicDecBlk(c[1] + c[0] // 8, c[2])
+        self.decoder_block2 = _BasicDecBlk(c[2] + c[1] // 8, c[3])
+        self.decoder_block1 = _BasicDecBlk(c[3] + c[2] // 8, c[3] // 2)
+        self.conv_out1 = nn.Sequential(nn.Conv2d(c[3] // 2 + c[3] // 8, 1, 1, 1, 0))
+
+        self.lateral_block4 = _BasicLatBlk(c[1], c[1])
+        self.lateral_block3 = _BasicLatBlk(c[2], c[2])
+        self.lateral_block2 = _BasicLatBlk(c[3], c[3])
+
+        # multi-scale supervision heads (training only — kept for state_dict
+        # parity with the released checkpoint; not consumed at inference).
+        self.conv_ms_spvn_4 = nn.Conv2d(c[1], 1, 1, 1, 0)
+        self.conv_ms_spvn_3 = nn.Conv2d(c[2], 1, 1, 1, 0)
+        self.conv_ms_spvn_2 = nn.Conv2d(c[3], 1, 1, 1, 0)
+
+        # gradient-decoder-triggering (gdt) attention: used at inference to
+        # gate p4/p3/p2.
+        _N = 16
+        def _gdt_branch(in_c):
+            return nn.Sequential(nn.Conv2d(in_c, _N, 3, 1, 1), nn.BatchNorm2d(_N), nn.ReLU(inplace=True))
+        self.gdt_convs_4 = _gdt_branch(c[1])
+        self.gdt_convs_3 = _gdt_branch(c[2])
+        self.gdt_convs_2 = _gdt_branch(c[3])
+
+        def _head_1x1():
+            return nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+        # multi-scale supervision heads on the gdt branch (training only)
+        self.gdt_convs_pred_4 = _head_1x1()
+        self.gdt_convs_pred_3 = _head_1x1()
+        self.gdt_convs_pred_2 = _head_1x1()
+        # attention heads
+        self.gdt_convs_attn_4 = _head_1x1()
+        self.gdt_convs_attn_3 = _head_1x1()
+        self.gdt_convs_attn_2 = _head_1x1()
+
+    def forward(self, x, x1, x2, x3, x4):
+        x4 = torch.cat((x4, self.ipt_blk5(_image2patches(x, x4))), 1)
+        p4 = self.decoder_block4(x4)
+        p4 = p4 * self.gdt_convs_attn_4(self.gdt_convs_4(p4)).sigmoid()
+        _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True)
+        _p3 = _p4 + self.lateral_block4(x3)
+
+        _p3 = torch.cat((_p3, self.ipt_blk4(_image2patches(x, _p3))), 1)
+        p3 = self.decoder_block3(_p3)
+        p3 = p3 * self.gdt_convs_attn_3(self.gdt_convs_3(p3)).sigmoid()
+        _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True)
+        _p2 = _p3 + self.lateral_block3(x2)
+
+        _p2 = torch.cat((_p2, self.ipt_blk3(_image2patches(x, _p2))), 1)
+        p2 = self.decoder_block2(_p2)
+        p2 = p2 * self.gdt_convs_attn_2(self.gdt_convs_2(p2)).sigmoid()
+        _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True)
+        _p1 = _p2 + self.lateral_block2(x1)
+
+        _p1 = torch.cat((_p1, self.ipt_blk2(_image2patches(x, _p1))), 1)
+        _p1 = self.decoder_block1(_p1)
+        _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True)
+
+        _p1 = torch.cat((_p1, self.ipt_blk1(_image2patches(x, _p1))), 1)
+        return self.conv_out1(_p1)
+
+
+class BiRefNet(nn.Module):
+    """BiRefNet (ZhengPeng7/BiRefNet) with Swin-L backbone, multi-scale input
+    concatenation, ASPP-deformable squeeze block, and the 4-level
+    input-modulating decoder used in the v1 release.
+
+    `forward(x)` returns a single 1-channel alpha map in `[0, 1]` (post-sigmoid).
+    `remove_background(pil_img)` is the PIL helper used by the pipeline —
+    accepts a PIL RGB image and returns an RGBA copy with the predicted matte
+    in the alpha channel.
+    """
+
+    INPUT_SIZE = (1024, 1024)
+    # backbone channel counts post mul_scl_ipt='cat' (doubled from raw Swin-L)
+    _CHANNELS = (3072, 1536, 768, 384)
+    # ImageNet normalization used by the BiRefNet recipe
+    _NORM_MEAN = (0.485, 0.456, 0.406)
+    _NORM_STD  = (0.229, 0.224, 0.225)
+
+    def __init__(self):
+        super().__init__()
+        self.bb = _SwinLarge()
+        cxt = list(self._CHANNELS[1:][::-1][-3:])  # = [384, 768, 1536]
+        self.squeeze_module = nn.Sequential(
+            _BasicDecBlk(self._CHANNELS[0] + sum(cxt), self._CHANNELS[0])
+        )
+        self.decoder = _BiRefNetDecoder(channels=self._CHANNELS)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def _forward_enc(self, x):
+        x1, x2, x3, x4 = self.bb(x)
+        # mul_scl_ipt='cat': re-run backbone at half resolution, concat features
+        B, C, H, W = x.shape
+        x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H // 2, W // 2),
+                                                   mode='bilinear', align_corners=True))
+        x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], 1)
+        x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], 1)
+        x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], 1)
+        x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], 1)
+        # cxt: upsample x1/x2/x3 to x4 spatial and concat for the squeeze input
+        x4 = torch.cat([
+            F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True),
+            F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True),
+            F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True),
+            x4,
+        ], 1)
+        return x1, x2, x3, x4
+
+    def forward(self, x):
+        x1, x2, x3, x4 = self._forward_enc(x)
+        x4 = self.squeeze_module(x4)
+        logits = self.decoder(x, x1, x2, x3, x4)
+        return torch.sigmoid(logits)
+
+    def load_safetensors(self, path: str) -> None:
+        sd = safetensors.torch.load_file(path)
+        # The decoder's gdt_convs_pred_* / conv_ms_spvn_* heads are training-only
+        # but are kept as submodules for state_dict parity. strict=True works.
+        missing, unexpected = self.load_state_dict(sd, strict=False)
+        if unexpected:
+            raise KeyError(f"[birefnet] unexpected keys (e.g. {unexpected[:3]})")
+        if missing:
+            raise KeyError(f"[birefnet] missing keys (e.g. {missing[:3]})")
+
+    @torch.no_grad()
+    def remove_background(self, image) -> "Image.Image":
+        from PIL import Image
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+        W, H = image.size
+        arr = np.array(image, dtype=np.float32) / 255.0
+        t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
+        t = F.interpolate(t, size=self.INPUT_SIZE, mode='bilinear', align_corners=True)
+        mean = torch.tensor(self._NORM_MEAN).view(1, 3, 1, 1)
+        std  = torch.tensor(self._NORM_STD).view(1, 3, 1, 1)
+        t = ((t - mean) / std).to(device=self.device, dtype=self.dtype)
+        alpha = self.forward(t)
+        alpha = F.interpolate(alpha.float(), size=(H, W), mode='bilinear', align_corners=True)[0, 0]
+        a = (alpha.clamp(0, 1) * 255).to(torch.uint8).cpu().numpy()
+        rgba = image.copy()
+        rgba.putalpha(Image.fromarray(a, mode="L"))
+        return rgba
+
+
+# ---------------------------------------------------------------------------
+# Shared transformer helpers
+# ---------------------------------------------------------------------------
+
+class LayerNorm32(nn.LayerNorm):
+    def forward(self, x):
+        origin_dtype = x.dtype
+        return F.layer_norm(
+            x.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        ).to(origin_dtype)
+
+
+class MultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+    def forward(self, x):
+        origin_dtype = x.dtype
+        return (F.normalize(x.float(), dim=-1) * self.gamma.float() * self.scale).to(origin_dtype)
+
+
+def apply_rotary_emb(hidden_states, freqs):
+    x_rotated = torch.view_as_complex(hidden_states.float().reshape(*hidden_states.shape[:-1], -1, 2))
+    x_rotated = x_rotated * freqs
+    x_out = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1)
+    return x_out.type_as(hidden_states)
+
+
+def clamp_mul(x, f):
+    f_t = f.tanh()
+    return x * f_t + x.detach() * (f - f_t)
+
+
+def scaled_dot_product_attention(qkv=None, q=None, k=None, v=None, kv=None):
+    if qkv is not None:
+        q, k, v = qkv.unbind(dim=2)
+    elif kv is not None:
+        k, v = kv.unbind(dim=2)
+    q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
+    return F.scaled_dot_product_attention(q, k, v).permute(0, 2, 1, 3)
+
+
+# ---------------------------------------------------------------------------
+# Positional embeddings
+# ---------------------------------------------------------------------------
+
+class RePo3DRotaryEmbedding(nn.Module):
+    def __init__(self, model_channels, num_heads, head_dim, repo_hidden_ratio=0.125, max_freq=16.0):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        repo_hidden_size = int(model_channels * repo_hidden_ratio)
+        self.norm = LayerNorm32(model_channels)
+        self.gate_map = nn.Linear(model_channels, repo_hidden_size, bias=False)
+        self.content_map = nn.Linear(model_channels, repo_hidden_size, bias=False)
+        self.act = nn.SiLU()
+        self.final_map = nn.Linear(repo_hidden_size, 3 * num_heads, bias=False)
+        self.dim_0 = 2 * (head_dim // 6)
+        self.dim_1 = 2 * (head_dim // 6)
+        self.dim_2 = head_dim - self.dim_0 - self.dim_1
+        dims = [self.dim_0, self.dim_1, self.dim_2]
+        freqs_list = []
+        for d in dims:
+            freq_dim = d // 2
+            freqs_list.append(torch.linspace(1.0, float(max_freq), steps=freq_dim, dtype=torch.float32))
+        self.freqs_0 = nn.Parameter(freqs_list[0])
+        self.freqs_1 = nn.Parameter(freqs_list[1])
+        self.freqs_2 = nn.Parameter(freqs_list[2])
+
+    def forward(self, hidden_states):
+        h = self.norm(hidden_states)
+        feat = self.act(self.gate_map(h)) * self.content_map(h)
+        out = self.final_map(feat)
+        B, L, _ = out.shape
+        delta_pos = out.reshape(B, L, self.num_heads, 3)
+        ang_0 = clamp_mul(delta_pos[..., 0].unsqueeze(-1), self.freqs_0) * torch.pi
+        ang_1 = clamp_mul(delta_pos[..., 1].unsqueeze(-1), self.freqs_1) * torch.pi
+        ang_2 = clamp_mul(delta_pos[..., 2].unsqueeze(-1), self.freqs_2) * torch.pi
+        ang = torch.cat([ang_0, ang_1, ang_2], dim=-1).float()  # fp32 needed for torch.polar → complex64
+        return torch.polar(torch.ones_like(ang), ang).type(torch.complex64)
+
+
+class PcdAbsolutePositionEmbedder(nn.Module):
+    def __init__(self, channels: int, in_channels: int = 3, max_res: int = 16):
+        super().__init__()
+        self.channels = channels
+        self.in_channels = in_channels
+        self.max_res = max_res
+        self.freq_dim = channels // in_channels // 2
+
+    def _freqs(self, device):
+        freqs_2exp = torch.arange(self.max_res, dtype=torch.float32, device=device)
+        res_dim = max(0, self.freq_dim - self.max_res)
+        freqs_res = (torch.arange(res_dim, dtype=torch.float32, device=device) / max(res_dim, 1) * self.max_res
+                     if res_dim > 0 else torch.empty(0, device=device))
+        freqs = torch.cat([freqs_2exp, freqs_res], dim=0)[:self.freq_dim]
+        return torch.pow(2.0, freqs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        orig_dtype = x.dtype
+        x = x.float()
+        *dims, D = x.shape
+        out = torch.outer(x.reshape(-1), self._freqs(x.device)) * 2 * torch.pi
+        out = torch.cat([out.sin(), out.cos()], dim=-1).reshape(*dims, -1)
+        if out.shape[-1] < self.channels:
+            out = torch.cat([out, torch.zeros(*dims, self.channels - out.shape[-1],
+                                              device=out.device, dtype=out.dtype)], dim=-1)
+        return out.to(orig_dtype)
+
+
+class PcdAbsolutePositionEmbedderV2(nn.Module):
+    def __init__(self, channels: int, in_channels: int = 3, max_res: int = 10):
+        super().__init__()
+        self.channels = channels
+        self.in_channels = in_channels
+        self.max_res = max_res
+        self.freq_dim = channels // in_channels // 2
+
+    def _freqs(self, device):
+        logs = torch.linspace(0.0, float(self.max_res), steps=self.freq_dim, dtype=torch.float32, device=device)
+        return torch.pow(2.0, logs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        orig_dtype = x.dtype
+        x = x.float()
+        N, D = x.shape
+        ang = x.unsqueeze(-1) * self._freqs(x.device) * torch.pi
+        embed = torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1).reshape(N, -1)
+        if embed.shape[1] < self.channels:
+            embed = torch.cat([embed, torch.zeros(N, self.channels - embed.shape[1],
+                                                   device=embed.device, dtype=embed.dtype)], dim=-1)
+        return embed.to(orig_dtype)
+
+
+# ---------------------------------------------------------------------------
+# Transformer building blocks
+# ---------------------------------------------------------------------------
+
+class FeedForwardNet(nn.Module):
+    def __init__(self, channels, mlp_ratio=4.0, channels_out=None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, int(channels * mlp_ratio)),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(int(channels * mlp_ratio), channels if channels_out is None else channels_out),
+        )
+
+    def forward(self, x):
+        return self.mlp(x)
+
+
+class MLP(nn.Module):
+    def __init__(self, channels: int, inner_channels: int, channels_out: Optional[int] = None,
+                 mlp_layer_num: int = 2):
+        super().__init__()
+        layers = []
+        for i in range(mlp_layer_num - 1):
+            layers.append(nn.Linear(channels if i == 0 else inner_channels, inner_channels))
+            layers.append(nn.GELU(approximate="tanh"))
+        layers.append(nn.Linear(inner_channels, channels if channels_out is None else channels_out))
+        self.mlp = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+
+class RopeMultiHeadAttention(nn.Module):
+    def __init__(self, channels, num_heads, ctx_channels=None, type="self",
+                 attn_mode="full", qkv_bias=True, qk_rms_norm=False, use_rope=False):
+        super().__init__()
+        self.channels = channels
+        self.num_heads = num_heads
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self._type = type
+        self.qk_rms_norm = qk_rms_norm
+        self.use_rope = use_rope
+        if self._type == "self":
+            self.qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+        if self.qk_rms_norm:
+            self.q_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            self.k_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+        self.out = nn.Linear(channels, channels)
+
+    def forward(self, x, context=None, rope_emb=None):
+        B, L, C = x.shape
+        if self._type == "self":
+            qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim)
+            q, k, v = qkv.unbind(2)
+            if self.use_rope:
+                q = apply_rotary_emb(q, rope_emb)
+                k = apply_rotary_emb(k, rope_emb)
+        else:
+            q = self.q(x).reshape(B, L, self.num_heads, self.head_dim)
+            if context is None:
+                raise ValueError("Context must be provided for cross attention")
+            kv = self.kv(context).reshape(B, context.shape[1], 2, self.num_heads, self.head_dim)
+            k, v = kv.unbind(2)
+        if self.qk_rms_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        h = scaled_dot_product_attention(q=q, k=k, v=v)
+        return self.out(h.reshape(B, L, C))
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, channels, num_heads, ctx_channels=None, type="self",
+                 attn_mode="full", qkv_bias=True, qk_rms_norm=False):
+        super().__init__()
+        assert channels % num_heads == 0
+        assert type in ["self", "cross"]
+        assert attn_mode == "full"
+        self.channels = channels
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.qk_rms_norm = qk_rms_norm
+        if self._type == "self":
+            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+        if self.qk_rms_norm:
+            self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+        self.to_out = nn.Linear(channels, channels)
+
+    def forward(self, x, context=None):
+        B, L, C = x.shape
+        if self._type == "self":
+            qkv = self.to_qkv(x).reshape(B, L, 3, self.num_heads, -1)
+            if self.qk_rms_norm:
+                q, k, v = qkv.unbind(dim=2)
+                q = self.q_rms_norm(q)
+                k = self.k_rms_norm(k)
+                qkv = torch.stack([q, k, v], dim=2)
+            h = scaled_dot_product_attention(qkv=qkv)
+        else:
+            Lkv = context.shape[1]
+            q = self.to_q(x).reshape(B, L, self.num_heads, -1)
+            kv = self.to_kv(context).reshape(B, Lkv, 2, self.num_heads, -1)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=2)
+                k = self.k_rms_norm(k)
+                h = scaled_dot_product_attention(q=q, k=k, v=v)
+            else:
+                h = scaled_dot_product_attention(q=q, kv=kv)
+        return self.to_out(h.reshape(B, L, -1))
+
+
+class UnifiedTransformerBlock(nn.Module):
+    def __init__(self, channels, num_heads, mlp_ratio=4.0, attn_mode="full",
+                 use_checkpoint=False, use_rope=False, qk_rms_norm=False, qkv_bias=True,
+                 modulation=True, share_mod=False, use_shift_table=False):
+        super().__init__()
+        self.modulation = modulation
+        self.share_mod = share_mod
+        self.norm1 = LayerNorm32(channels, elementwise_affine=not modulation, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=not modulation, eps=1e-6)
+        self.attn = RopeMultiHeadAttention(channels, num_heads=num_heads, type="self",
+                                            attn_mode=attn_mode, qkv_bias=qkv_bias,
+                                            use_rope=use_rope, qk_rms_norm=qk_rms_norm)
+        self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+        if modulation:
+            if not share_mod:
+                self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(channels, 6 * channels, bias=True))
+            self.shift_table = nn.Parameter(torch.randn(1, 6 * channels) / channels ** 0.5) if use_shift_table else None
+
+    def forward(self, x, mod=None, rotary_emb=None):
+        if self.modulation:
+            if not self.share_mod:
+                mod = self.adaLN_modulation(mod)
+            if hasattr(self, 'shift_table') and self.shift_table is not None:
+                mod = mod + self.shift_table.type(mod.dtype)
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
+            h = self.norm1(x)
+            h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+            h = self.attn(h, rope_emb=rotary_emb)
+            x = x + h * gate_msa.unsqueeze(1)
+            h = self.norm2(x)
+            h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+            x = x + self.mlp(h) * gate_mlp.unsqueeze(1)
+        else:
+            x = x + self.attn(self.norm1(x), rope_emb=rotary_emb)
+            x = x + self.mlp(self.norm2(x))
+        return x
+
+
+# ---------------------------------------------------------------------------
+# Quasi-random sampling utilities
+# ---------------------------------------------------------------------------
+
+PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
+
+
+def radical_inverse(base, n):
+    val = 0
+    inv_base = 1.0 / base
+    inv_base_n = inv_base
+    while n > 0:
+        digit = n % base
+        val += digit * inv_base_n
+        n //= base
+        inv_base_n *= inv_base
+    return val
+
+
+def halton_sequence(dim, n):
+    return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]
+
+
+def hammersley_sequence(dim, n, num_samples):
+    return [n / num_samples] + halton_sequence(dim - 1, n)
+
+
+@torch.no_grad()
+def sample_probs(probs, counts, algo="systematic"):
+    batch_shape = counts.shape
+    B = counts.numel()
+    P = probs.size(-1)
+    device = probs.device
+    probs = probs.view(B, P)
+    counts = counts.view(B)
+
+    probs = probs.to(torch.float32).clamp_min_(0)
+    row_sums = probs.sum(1, keepdim=True)
+    zero_mask = row_sums.eq(0)
+    probs = probs / row_sums.clamp_min_(1)
+    if zero_mask.any():
+        probs = probs.clone()
+        probs[zero_mask.expand_as(probs)] = 1.0 / P
+
+    counts = counts.to(device=device, dtype=torch.long)
+    out = torch.zeros(B, P, dtype=torch.long, device=device)
+    cdf = probs.cumsum(dim=1).clamp(max=1.0 - 1e-12)
+    unique_n, inv = counts.unique(sorted=False, return_inverse=True)
+    for i, n in enumerate(unique_n.tolist()):
+        if n == 0:
+            continue
+        rows = (inv == i).nonzero(as_tuple=False).squeeze(1)
+        r = rows.numel()
+        U0 = torch.rand(r, 1, device=device) / float(n)
+        grid = torch.arange(n, device=device, dtype=torch.float32)[None, :] / float(n)
+        us = (U0 + grid).clamp(max=1.0 - 1e-12)
+        cdf_rows = cdf.index_select(0, rows)
+        idx = torch.searchsorted(cdf_rows, us).clamp_max(probs.size(1) - 1)
+        buf = torch.zeros(r, P, dtype=torch.float32, device=device)
+        buf.scatter_add_(1, idx, torch.ones_like(idx, dtype=buf.dtype))
+        out.index_copy_(0, rows, buf.to(torch.long))
+
+    return out.view(*batch_shape, P)
+
+
+# ---------------------------------------------------------------------------
+# VAE decoders
+# ---------------------------------------------------------------------------
+
+class LevelEmbedder(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256, max_period=1024):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+        self.max_period = max_period
+
+    @staticmethod
+    def level_embedding(t, dim, max_period=1024):
+        half = dim // 2
+        freqs = torch.exp(-np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device)
+        args = t[:, None].float() * freqs[None] * 2 * torch.pi
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        emb = self.level_embedding(t, self.frequency_embedding_size, self.max_period)
+        return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+class ModulatedTransformerCrossOnlyBlock(nn.Module):
+    def __init__(self, channels, ctx_channels, num_heads, mlp_ratio=4.0, share_mod=False,
+                 qk_rms_norm_cross=True, qkv_bias=True):
+        super().__init__()
+        self.share_mod = share_mod
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.cross_attn = MultiHeadAttention(channels, ctx_channels=ctx_channels, num_heads=num_heads,
+                                              type="cross", attn_mode="full", qkv_bias=qkv_bias,
+                                              qk_rms_norm=qk_rms_norm_cross)
+        self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+        if not share_mod:
+            self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(channels, 6 * channels, bias=True))
+
+    def forward(self, x, mod, context):
+        if self.share_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+        h = self.norm1(x) * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+        x = x + self.cross_attn(h, context) * gate_msa.unsqueeze(1)
+        h = self.norm2(x) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+        x = x + self.mlp(h) * gate_mlp.unsqueeze(1)
+        return x
+
+
+class ModulatedCrossOnlyTransformerBase(nn.Module):
+    def __init__(self, in_channels, model_channels, cond_channels, num_blocks, num_heads=None,
+                 num_head_channels=64, mlp_ratio=4.0, share_mod=False, additional_level_embed=False,
+                 qk_rms_norm_cross=True):
+        super().__init__()
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.share_mod = share_mod
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+
+        self.input_layer = nn.Linear(in_channels, model_channels)
+        self.l_embedder = LevelEmbedder(model_channels)
+        self.l_embedder2 = LevelEmbedder(model_channels, max_period=100) if additional_level_embed else None
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(), nn.Linear(model_channels, 6 * model_channels, bias=True))
+        if cond_channels is not None:
+            self.blocks = nn.ModuleList([
+                ModulatedTransformerCrossOnlyBlock(
+                    model_channels, ctx_channels=cond_channels, num_heads=self.num_heads,
+                    mlp_ratio=self.mlp_ratio, qk_rms_norm_cross=self.qk_rms_norm_cross,
+                    share_mod=self.share_mod)
+                for _ in range(num_blocks)
+            ])
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def forward(self, x, l, cond, l2=None):
+        h = self.input_layer(x)
+        l_emb = self.l_embedder(l)
+        if self.l_embedder2 is not None and l2 is not None:
+            l_emb = l_emb + self.l_embedder2(l2)
+        if self.share_mod:
+            l_emb = self.adaLN_modulation(l_emb)
+        for block in self.blocks:
+            h = block(h, l_emb, cond)
+        return h
+
+
+class OctreeProbabilityFixedlenDecoder(ModulatedCrossOnlyTransformerBase):
+    def __init__(self, model_channels, cond_channels, num_blocks, num_heads=None,
+                 num_head_channels=64, mlp_ratio=4.0, share_mod=False,
+                 additional_level_embed=False, qk_rms_norm_cross=True, *,
+                 no_norm=False):
+        super().__init__(
+            in_channels=model_channels, model_channels=model_channels,
+            cond_channels=cond_channels, num_blocks=num_blocks,
+            num_heads=num_heads, num_head_channels=num_head_channels,
+            mlp_ratio=mlp_ratio, share_mod=share_mod,
+            additional_level_embed=additional_level_embed,
+            qk_rms_norm_cross=qk_rms_norm_cross,
+        )
+        self.out_proj = nn.Linear(self.model_channels, 8)
+        self.no_norm = no_norm
+        self.in_proj = nn.Linear(3, self.model_channels)
+        self.pos_embedder = PcdAbsolutePositionEmbedderV2(channels=model_channels, in_channels=3)
+
+    def forward(self, x, l, cond, l2=None):
+        d = self.dtype
+        B, L, C = x.shape
+        h = self.in_proj(x.to(d)) + self.pos_embedder(x.reshape(-1, 3)).reshape(B, L, -1).to(d)
+        if l2 is not None:
+            l2 = torch.log2(l2)
+        h = super().forward(h, l, cond.to(d), l2)
+        h = F.layer_norm(h.float(), h.shape[-1:]).to(d) if not self.no_norm else h / (1 + 2 * self.num_blocks) ** 0.5
+        logits = self.out_proj(h)
+        return {"logits": logits, "probs": torch.softmax(logits, dim=-1)}
+
+    @staticmethod
+    def sample(model, cond, num_points, level, temperature=1.0, algo="systematic"):
+        B = cond.shape[0]
+        device = cond.device
+        child_offset = torch.tensor([[i, j, k] for k in [0, 1] for j in [0, 1] for i in [0, 1]],
+                                    dtype=torch.long, device=device)
+        prev_coords_int = torch.zeros(B, 1, 3, dtype=torch.long, device=device)
+        prev_counts = torch.full((B, 1), num_points, dtype=torch.long, device=device)
+        prev_log_probs = torch.zeros(B, 1, dtype=torch.float32, device=device)
+        batch_indices_range = torch.arange(B, device=device).unsqueeze(1)
+        num_tensor = torch.full((B,), num_points, dtype=torch.long, device=device)
+
+        for lv in range(1, level + 1):
+            res_p = 1 << (lv - 1)
+            res = 1 << lv
+            parent_coords_norm = (prev_coords_int.to(torch.float32) + 0.5) / res_p
+            res_tensor = torch.full((B,), res, dtype=torch.long, device=device)
+            pred_logits = model(parent_coords_norm, res_tensor, cond, num_tensor)["logits"] / temperature
+            pred_probs = torch.softmax(pred_logits, dim=-1)
+            pred_log_probs = torch.log_softmax(pred_logits, dim=-1)
+            sampled = sample_probs(pred_probs, prev_counts, algo=algo).flatten(1, 2)
+            pred_log_probs = pred_log_probs.flatten(1, 2)
+            prev_log_probs_expanded = prev_log_probs.repeat_interleave(8, dim=1)
+            child_coords_int = (prev_coords_int[:, :, None, :] * 2 + child_offset[None, None, :, :]).flatten(1, 2)
+            mask = sampled > 0
+            max_valid = mask.sum(dim=1).max().item()
+            scatter_indices = mask.cumsum(dim=1) - 1
+            valid_scatter_indices = scatter_indices[mask]
+            valid_batch_indices = batch_indices_range.expand_as(mask)[mask]
+            next_prev_coords_int = torch.zeros(B, max_valid, 3, dtype=child_coords_int.dtype, device=device)
+            next_prev_coords_int[valid_batch_indices, valid_scatter_indices] = child_coords_int[mask]
+            next_prev_counts = torch.zeros(B, max_valid, dtype=sampled.dtype, device=device)
+            next_prev_counts[valid_batch_indices, valid_scatter_indices] = sampled[mask]
+            next_prev_log_probs = torch.zeros(B, max_valid, dtype=prev_log_probs.dtype, device=device)
+            next_prev_log_probs[valid_batch_indices, valid_scatter_indices] = (prev_log_probs_expanded + pred_log_probs)[mask]
+            prev_coords_int = next_prev_coords_int
+            prev_counts = next_prev_counts
+            prev_log_probs = next_prev_log_probs
+
+        res = 1 << level
+        prev_log_probs = torch.repeat_interleave(prev_log_probs.flatten(0, 1), prev_counts.flatten(0, 1), dim=0).reshape(B, num_points)
+        coords_int = torch.repeat_interleave(prev_coords_int.flatten(0, 1), prev_counts.flatten(0, 1), dim=0).reshape(B, num_points, -1)
+        coords_norm = (coords_int.to(torch.float32) + torch.rand_like(coords_int, dtype=torch.float32)) / res
+        return {"points": coords_norm, "log_probs": prev_log_probs}
+
+
+class TransformerCrossBlock(nn.Module):
+    def __init__(self, channels, ctx_channels, num_heads, mlp_ratio=4.0, attn_mode="full",
+                 qk_rms_norm=True, qk_rms_norm_cross=True, qkv_bias=True):
+        super().__init__()
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.self_attn = MultiHeadAttention(channels, num_heads=num_heads, type="self",
+                                             attn_mode=attn_mode, qkv_bias=qkv_bias,
+                                             qk_rms_norm=qk_rms_norm)
+        self.cross_attn = MultiHeadAttention(channels, ctx_channels=ctx_channels, num_heads=num_heads,
+                                              type="cross", attn_mode="full", qkv_bias=qkv_bias,
+                                              qk_rms_norm=qk_rms_norm_cross)
+        self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+
+    def forward(self, x, context):
+        x = x + self.self_attn(self.norm1(x))
+        x = x + self.cross_attn(self.norm2(x), context)
+        x = x + self.mlp(self.norm3(x))
+        return x
+
+
+class TransformerBase(nn.Module):
+    def __init__(self, in_channels, model_channels, cond_channels, num_blocks, num_heads=None,
+                 num_head_channels=64, mlp_ratio=4.0, attn_mode="full", window_num=None,
+                 qk_rms_norm=True, qk_rms_norm_cross=True):
+        super().__init__()
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.input_layer = nn.Linear(in_channels, model_channels)
+        if cond_channels is not None:
+            self.blocks = nn.ModuleList([
+                TransformerCrossBlock(model_channels, ctx_channels=cond_channels,
+                                     num_heads=self.num_heads, mlp_ratio=self.mlp_ratio,
+                                     attn_mode="full", qk_rms_norm=qk_rms_norm,
+                                     qk_rms_norm_cross=qk_rms_norm_cross)
+                for _ in range(num_blocks)
+            ])
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    def forward(self, x, cond=None, l=None, cond2=None):
+        h = self.input_layer(x)
+        for block in self.blocks:
+            h = block(h, cond)
+        return h
+
+
+class FixedlenDecoder(TransformerBase):
+    def __init__(self, in_channels, model_channels, cond_channels, num_blocks, num_heads=None,
+                 num_head_channels=64, mlp_ratio=4.0, attn_mode="full", window_num=None,
+                 qk_rms_norm=True, qk_rms_norm_cross=True):
+        super().__init__(in_channels=model_channels, model_channels=model_channels,
+                         cond_channels=cond_channels, num_blocks=num_blocks,
+                         num_heads=num_heads, num_head_channels=num_head_channels,
+                         mlp_ratio=mlp_ratio, attn_mode=attn_mode, window_num=window_num,
+                         qk_rms_norm=qk_rms_norm, qk_rms_norm_cross=qk_rms_norm_cross)
+        self.in_proj = nn.Linear(in_channels, model_channels)
+        self.pos_embedder = PcdAbsolutePositionEmbedderV2(channels=model_channels, in_channels=3)
+
+    def forward(self, x=None, cond=None):
+        pcd = x["points"]
+        d = self.dtype
+        B, L, C = pcd.shape
+        h = self.in_proj(pcd.to(d)) + self.pos_embedder(pcd.reshape(-1, 3)).reshape(B, L, -1).to(d)
+        return super().forward(h, cond.to(d))
+
+
+class ElasticGaussianFixedlenDecoder(FixedlenDecoder):
+    def __init__(self, in_channels, model_channels, cond_channels, num_blocks, num_heads=None,
+                 num_head_channels=64, mlp_ratio=4.0, attn_mode="full", window_num=None,
+                 *, no_norm=False, representation_config=None,
+                 use_learned_offset_scale=True, use_per_offset=True,
+                 qk_rms_norm=True, qk_rms_norm_cross=True):
+        self.rep_config = representation_config
+        self.use_learned_offset_scale = use_learned_offset_scale
+        self.use_per_offset = use_per_offset
+        self.out_channels = self._calc_layout()
+        super().__init__(in_channels=in_channels, model_channels=model_channels,
+                         cond_channels=cond_channels, num_blocks=num_blocks,
+                         num_heads=num_heads, num_head_channels=num_head_channels,
+                         mlp_ratio=mlp_ratio, attn_mode=attn_mode, window_num=window_num,
+                         qk_rms_norm=qk_rms_norm, qk_rms_norm_cross=qk_rms_norm_cross)
+        self.out_proj = nn.Linear(model_channels, self.out_channels)
+        self.no_norm = no_norm
+        self._build_perturbation()
+
+    def _calc_layout(self):
+        ng = self.rep_config['num_gaussians']
+        self.layout = {
+            '_xyz':         {'shape': (ng, 3),    'size': ng * 3},
+            '_features_dc': {'shape': (ng, 1, 3), 'size': ng * 3},
+            '_scaling':     {'shape': (ng, 3),    'size': ng * 3},
+            '_rotation':    {'shape': (ng, 4),    'size': ng * 4},
+            '_opacity':     {'shape': (ng, 1),    'size': ng},
+        }
+        if self.use_learned_offset_scale and self.use_per_offset:
+            self.layout['_offset_scale'] = {'shape': (ng, 1), 'size': ng}
+        start = 0
+        for k, v in self.layout.items():
+            v['range'] = (start, start + v['size'])
+            start += v['size']
+        return start
+
+    def _build_perturbation(self):
+        ng = self.rep_config['num_gaussians']
+        perturbation = torch.tensor([hammersley_sequence(3, i, ng) for i in range(ng)]).float()
+        perturbation = torch.atanh((perturbation * 2 - 1) / self.rep_config['perturbe_size'])
+        self.register_buffer('points_offset_perturbation', perturbation)
+        if self.use_learned_offset_scale:
+            base = torch.tensor(self.rep_config['offset_scale'])
+            self.register_buffer('base_offset_scale', torch.log(torch.exp(base) - 1.0))
+
+    def _get_offset(self, h):
+        B = h.shape[0]
+        if self.use_learned_offset_scale:
+            r = self.layout['_offset_scale']['range']
+            _offset_scale = F.softplus(
+                h[:, :, r[0]:r[1]].reshape(B, -1, *self.layout['_offset_scale']['shape'])
+                + self.base_offset_scale)
+
+        r = self.layout['_xyz']['range']
+        offset = h[:, :, r[0]:r[1]].reshape(B, -1, *self.layout['_xyz']['shape'])
+        offset = offset * self.rep_config['lr']['_xyz']
+        if self.rep_config['perturb_offset']:
+            offset = offset + self.points_offset_perturbation
+        offset = torch.tanh(offset) * 0.5 * self.rep_config['perturbe_size']
+        offset = offset * (_offset_scale if self.use_learned_offset_scale else self.rep_config['offset_scale'])
+        return offset
+
+    def forward(self, x=None, cond=None):
+        h = super().forward(x, cond)
+        h = F.layer_norm(h.float(), h.shape[-1:]).to(h.dtype) if not self.no_norm else h / (1 + 3 * self.num_blocks) ** 0.5
+        return {"features": self.out_proj(h)}
+
+
+# ---------------------------------------------------------------------------
+# Flow matching denoiser
+# ---------------------------------------------------------------------------
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(-np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        emb = self.timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+class LatentSeqMMFlowModel(nn.Module):
+    def __init__(self, q_token_length, in_channels, model_channels, cond_channels,
+                 out_channels, num_blocks, num_refiner_blocks=2, num_heads=None,
+                 num_head_channels=64, cam_channels=None, cond2_channels=None,
+                 mlp_ratio=4, share_mod=True, qk_rms_norm=False, use_shift_table=False):
+        super().__init__()
+        self.q_token_length = q_token_length
+        self.in_channels = in_channels
+        self.cam_channels = cam_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.cond2_channels = cond2_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_refiner_blocks = num_refiner_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.share_mod = share_mod
+        self.qk_rms_norm = qk_rms_norm
+        self.use_shift_table = use_shift_table
+
+        self.t_embedder = TimestepEmbedder(model_channels)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(), nn.Linear(model_channels, 6 * model_channels, bias=True))
+
+        self.input_layer = nn.Linear(in_channels, model_channels)
+        self.cond_embedder = nn.Linear(cond_channels, model_channels)
+        self.cond_embedder2 = nn.Linear(cond2_channels, model_channels) if cond2_channels is not None else None
+
+        sobol_seq = torch.quasirandom.SobolEngine(dimension=3, scramble=True, seed=123).draw(q_token_length)
+        self.pos_pe = sobol_seq.unsqueeze(0)
+        self.pos_embedder = PcdAbsolutePositionEmbedder(model_channels)
+        self.noise_repo_layers = nn.ModuleList([
+            RePo3DRotaryEmbedding(model_channels, num_heads=self.num_heads, head_dim=num_head_channels)
+            for _ in range(num_refiner_blocks)])
+        self.context_repo_layers = nn.ModuleList([
+            RePo3DRotaryEmbedding(model_channels, num_heads=self.num_heads, head_dim=num_head_channels)
+            for _ in range(num_refiner_blocks)])
+        self.repo_layers = nn.ModuleList([
+            RePo3DRotaryEmbedding(model_channels, num_heads=self.num_heads, head_dim=num_head_channels)
+            for _ in range(num_blocks)])
+
+        block_kwargs = dict(num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, attn_mode='full',
+                            use_rope=True, qk_rms_norm=self.qk_rms_norm,
+                            use_shift_table=self.use_shift_table)
+        self.noise_refiner = nn.ModuleList([
+            UnifiedTransformerBlock(model_channels, modulation=True, share_mod=self.share_mod, **block_kwargs)
+            for _ in range(num_refiner_blocks)])
+        self.context_refiner = nn.ModuleList([
+            UnifiedTransformerBlock(model_channels, modulation=False, **block_kwargs)
+            for _ in range(num_refiner_blocks)])
+        if self.cam_channels is not None:
+            self.cam_refiner = MLP(self.cam_channels, model_channels, model_channels,
+                                   mlp_layer_num=num_refiner_blocks)
+        self.blocks = nn.ModuleList([
+            UnifiedTransformerBlock(model_channels, modulation=True, share_mod=self.share_mod, **block_kwargs)
+            for _ in range(num_blocks)])
+        self.shift_table = nn.Parameter(torch.randn(1, 2, model_channels) / model_channels**0.5) if use_shift_table else None
+        self.out_layer = nn.Linear(model_channels, out_channels)
+        if cam_channels is not None:
+            self.cam_out_layer = nn.Linear(model_channels, cam_channels)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def load_safetensors(self, path: str) -> None:
+        self.load_state_dict(safetensors.torch.load_file(path), strict=True)
+
+    def forward(self, x_t, t, cond):
+        d = self.dtype
+        z = x_t['latent'].to(d)
+        feat1 = cond['feature1'].to(d)
+        feat2 = cond['feature2'].to(d) if self.cond_embedder2 is not None else None
+        self.pos_pe = self.pos_pe.to(z.device)
+
+        h_x = self.input_layer(z)
+        h_cond = self.cond_embedder(feat1)
+        if feat2 is not None:
+            h_cond = h_cond + self.cond_embedder2(feat2)
+        t_emb = self.t_embedder(t)
+        t_mod = self.adaLN_modulation(t_emb) if self.share_mod else t_emb
+
+        h_x = h_x + self.pos_embedder(self.pos_pe).to(d)
+
+        for i, block in enumerate(self.noise_refiner):
+            h_x = block(h_x, mod=t_mod, rotary_emb=self.noise_repo_layers[i](h_x))
+
+        for i, block in enumerate(self.context_refiner):
+            h_cond = block(h_cond, mod=None, rotary_emb=self.context_repo_layers[i](h_cond))
+
+        if self.cam_channels is not None:
+            cam = x_t.get('camera').to(d)
+            h_cam = self.cam_refiner(cam)
+
+        h = torch.cat([h_x, h_cond], dim=1)
+        if self.cam_channels is not None:
+            h = torch.cat([h, h_cam], dim=1)
+
+        for i, block in enumerate(self.blocks):
+            h = block(h, mod=t_mod, rotary_emb=self.repo_layers[i](h))
+
+        h_x = F.layer_norm(h[:, :z.shape[1]].float(), h.shape[-1:]).type(d)
+        if self.cam_channels is not None:
+            h_cam = F.layer_norm(h[:, -cam.shape[1]:].float(), h.shape[-1:]).type(d)
+
+        if self.use_shift_table:
+            shift, scale = (self.shift_table + t_emb.unsqueeze(1)).chunk(2, dim=1)
+            h_x = h_x * (1 + scale) + shift
+            if self.cam_channels is not None:
+                h_cam = h_cam * (1 + scale) + shift
+
+        out = {'latent': self.out_layer(h_x)}
+        if self.cam_channels is not None:
+            out['camera'] = self.cam_out_layer(h_cam)
+        return out
+
+
+# ---------------------------------------------------------------------------
+# OctreeGaussianDecoder
+# ---------------------------------------------------------------------------
+
+class OctreeGaussianDecoder(nn.Module):
+    _MAX_VOXEL_LEVEL = 8
+
+    def __init__(self, octree_args: dict, gs_args: dict):
+        super().__init__()
+        self.octree = OctreeProbabilityFixedlenDecoder(**octree_args)
+        self.gs     = ElasticGaussianFixedlenDecoder(**gs_args)
+
+    def load_safetensors(self, path: str) -> None:
+        self.load_state_dict(safetensors.torch.load_file(path), strict=True)
+
+    @property
+    def gaussians_per_point(self) -> int:
+        return self.gs.rep_config['num_gaussians']
+
+    @torch.no_grad()
+    def decode(self, latent: torch.Tensor, num_gaussians: int):
+        from triposplat import _build_gaussians  # local import: avoid model.py ↔ triposplat.py cycle
+        num_decoder_tokens = max(1, num_gaussians // self.gaussians_per_point)
+        points_pred = OctreeProbabilityFixedlenDecoder.sample(
+            self.octree, latent,
+            num_points=num_decoder_tokens, level=self._MAX_VOXEL_LEVEL,
+            temperature=1.0, algo='systematic',
+        )
+        pred = self.gs(x=points_pred, cond=latent)
+        return _build_gaussians(self.gs, points_pred, pred)[0]

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+numpy
+safetensors
+pillow
+tqdm

static/viewer/viewer.html

@@ -0,0 +1,167 @@
+<!doctype html>
+<html lang="en">
+<head>
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width,initial-scale=1">
+<title>3DGS Viewer</title>
+<style>
+  *, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
+  html, body { width: 100%; height: 100%; overflow: hidden;
+    background: radial-gradient(ellipse at 50% 60%, #1a2035 0%, #080b12 100%); }
+  canvas { display: block; }
+  #hud {
+    position: fixed; top: 12px; left: 16px;
+    color: #94a3b8; font: 12px/1.5 system-ui, sans-serif;
+    pointer-events: none; user-select: none;
+  }
+  #loading {
+    position: fixed; inset: 0;
+    display: flex; flex-direction: column;
+    align-items: center; justify-content: center;
+    color: #94a3b8; font: 14px system-ui, sans-serif;
+    gap: 12px;
+  }
+  #spinner {
+    width: 36px; height: 36px;
+    border: 3px solid #334155;
+    border-top-color: #60a5fa;
+    border-radius: 50%;
+    animation: spin 0.9s linear infinite;
+  }
+  @keyframes spin { to { transform: rotate(360deg); } }
+  #error {
+    display: none;
+    position: fixed; inset: 0;
+    align-items: center; justify-content: center;
+    color: #f87171; font: 13px system-ui, sans-serif;
+    padding: 32px; text-align: center; white-space: pre-wrap;
+  }
+</style>
+</head>
+<body>
+<div id="loading">
+  <div id="spinner"></div>
+  <div id="loading-label">Loading splat…</div>
+</div>
+<div id="error"></div>
+<div id="hud">drag to orbit &nbsp;·&nbsp; scroll to zoom &nbsp;·&nbsp; right-drag to pan</div>
+
+<script type="importmap">
+{
+  "imports": {
+    "three": "https://cdnjs.cloudflare.com/ajax/libs/three.js/0.180.0/three.module.js",
+    "three/addons/": "https://unpkg.com/three@0.180.0/examples/jsm/",
+    "@sparkjsdev/spark": "https://unpkg.com/@sparkjsdev/spark@2.0.0/dist/spark.module.js"
+  }
+}
+</script>
+
+<script type="module">
+import * as THREE from "three";
+import { OrbitControls } from "three/addons/controls/OrbitControls.js";
+import { SparkRenderer, SplatMesh } from "@sparkjsdev/spark";
+
+const params = new URLSearchParams(location.search);
+const plyURL = params.get("ply");
+
+const loadingEl = document.getElementById("loading");
+const loadLabel = document.getElementById("loading-label");
+const errorEl   = document.getElementById("error");
+
+function showError(msg) {
+  loadingEl.style.display = "none";
+  errorEl.style.display = "flex";
+  errorEl.textContent = msg;
+}
+
+if (!plyURL) {
+  showError("No ?ply= parameter provided.");
+} else {
+  init(plyURL);
+}
+
+function init(url) {
+  const scene = new THREE.Scene();
+
+  const camera = new THREE.PerspectiveCamera(45, window.innerWidth / window.innerHeight, 0.01, 1000);
+  camera.position.set(0, 0.3, 1.8);
+
+  const renderer = new THREE.WebGLRenderer({ antialias: false });
+  renderer.setPixelRatio(window.devicePixelRatio);
+  renderer.setSize(window.innerWidth, window.innerHeight);
+  document.body.appendChild(renderer.domElement);
+
+  const spark = new SparkRenderer({ renderer });
+  scene.add(spark);
+
+  const controls = new OrbitControls(camera, renderer.domElement);
+  controls.enableDamping = true;
+  controls.dampingFactor = 0.07;
+  controls.minDistance = 0.2;
+  controls.maxDistance = 50;
+  controls.target.set(0, 0, 0);
+
+  const splat = new SplatMesh({ url });
+  // DEG saves PLYs with -Y up and the front of the object facing roughly
+  // along the horizontal axis.  Re-orient for Three.js (+Y up, camera on
+  // +Z looking toward -Z): a Group lets us apply a yaw first (to bring the
+  // front-facing side to +Z) and then flip 180° around X to put up on +Y.
+  const splatRoot = new THREE.Group();
+  splatRoot.add(splat);
+  splat.rotation.y = Math.PI / 2;    // yaw the model so its front faces the camera
+  splatRoot.rotation.x = Math.PI;    // flip Y/Z so it stands upright
+  scene.add(splatRoot);
+
+  let framed = false;
+  function tryFrame() {
+    if (framed) return;
+    const box = new THREE.Box3().setFromObject(splatRoot);
+    if (box.isEmpty() || !isFinite(box.min.x)) return;
+    framed = true;
+    // const center = new THREE.Vector3();
+    // const size   = new THREE.Vector3();
+    // box.getCenter(center);
+    // box.getSize(size);
+    // console.log(size);
+    // const maxDim = Math.max(size.x, size.y, size.z);
+    // // With a 45° FOV, half-fov tan ≈ 0.414, so the minimum distance to fit
+    // // a sphere of diameter `maxDim` in view is maxDim / (2*0.414) ≈ 1.21*maxDim.
+    // // Use a small margin (1.15) so the object nearly fills the viewport.
+    // const dist = maxDim * 1.15;
+    // camera.position.copy(center).add(new THREE.Vector3(0, maxDim * 0.15, dist));
+    // controls.target.copy(center);
+    // controls.update();
+    loadingEl.style.display = "none";
+  }
+
+  fetch(url, { method: "HEAD" }).then(r => {
+    if (!r.ok) throw new Error(`HTTP ${r.status} ${r.statusText}`);
+    const len = r.headers.get("content-length");
+    if (len) loadLabel.textContent = `Loading splat (${(len / 1024 / 1024).toFixed(1)} MB)…`;
+  }).catch(e => showError("Fetch failed: " + e.message));
+
+  let checkFrames = 0;
+  function checkLoaded() {
+    checkFrames++;
+    if (checkFrames < 5) return;
+    tryFrame();
+  }
+
+  window.addEventListener("resize", () => {
+    camera.aspect = window.innerWidth / window.innerHeight;
+    camera.updateProjectionMatrix();
+    renderer.setSize(window.innerWidth, window.innerHeight);
+  });
+
+  renderer.setAnimationLoop(() => {
+    controls.update();
+    renderer.render(scene, camera);
+    checkLoaded();
+  });
+
+  // Fallback: hide loading after 4s regardless of bbox-readiness.
+  setTimeout(() => { loadingEl.style.display = "none"; }, 4000);
+}
+</script>
+</body>
+</html>

triposplat.py

@@ -0,0 +1,598 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import safetensors.torch
+from PIL import Image, ImageFilter
+from torchvision import transforms
+from tqdm.auto import tqdm
+
+from model import (
+    DinoV3ViT, Flux2VAEEncoder, BiRefNet,
+    OctreeProbabilityFixedlenDecoder, ElasticGaussianFixedlenDecoder,
+    LatentSeqMMFlowModel, OctreeGaussianDecoder,
+)
+
+
+# ---------------------------------------------------------------------------
+# Gaussian
+# ---------------------------------------------------------------------------
+
+class Gaussian:
+    def __init__(self, aabb: list, sh_degree: int = 0, mininum_kernel_size: float = 0.0,
+                 scaling_bias: float = 0.01, opacity_bias: float = 0.1,
+                 scaling_activation: str = "exp", device='cuda'):
+        self.sh_degree = sh_degree
+        self.mininum_kernel_size = mininum_kernel_size
+        self.scaling_bias = scaling_bias
+        self.opacity_bias = opacity_bias
+        self.device = device
+        self.aabb = torch.tensor(aabb, dtype=torch.float32, device=device)
+
+        if scaling_activation == "exp":
+            self._scaling_activation = torch.exp
+            self._inverse_scaling_activation = torch.log
+        elif scaling_activation == "softplus":
+            self._scaling_activation = F.softplus
+            self._inverse_scaling_activation = lambda x: x + torch.log(-torch.expm1(-x))
+
+        self._opacity_activation = torch.sigmoid
+        self._inverse_opacity_activation = lambda x: torch.log(x / (1 - x))
+
+        self.scale_bias = self._inverse_scaling_activation(torch.tensor(self.scaling_bias)).to(self.device)
+        self.rots_bias = torch.zeros(4, device=self.device)
+        self.rots_bias[0] = 1
+        self.opacity_bias_val = self._inverse_opacity_activation(torch.tensor(self.opacity_bias)).to(self.device)
+
+        self._storage = {}
+
+    def _get_store(self, name):
+        return self._storage.get(name)
+
+    def _set_store(self, name, value):
+        self._storage[name] = value
+
+    @property
+    def _xyz(self):
+        return self._get_store("_xyz")
+    @_xyz.setter
+    def _xyz(self, value):
+        if value is None:
+            self._set_store("_xyz", None); self._set_store("xyz", None); return
+        self._set_store("_xyz", value)
+        self._set_store("xyz", value * self.aabb[None, 3:] + self.aabb[None, :3])
+
+    @property
+    def get_xyz(self):
+        return self._get_store("xyz")
+
+    @property
+    def _features_dc(self):
+        return self._get_store("_features_dc")
+    @_features_dc.setter
+    def _features_dc(self, value):
+        self._set_store("_features_dc", value)
+
+    @property
+    def _opacity(self):
+        return self._get_store("_opacity")
+    @_opacity.setter
+    def _opacity(self, value):
+        if value is None:
+            self._set_store("_opacity", None); self._set_store("opacity", None); return
+        self._set_store("_opacity", value)
+        self._set_store("opacity", self._opacity_activation(value + self.opacity_bias_val))
+
+    @property
+    def get_opacity(self):
+        return self._get_store("opacity")
+
+    @property
+    def _scaling(self):
+        return self._get_store("_scaling")
+    @_scaling.setter
+    def _scaling(self, value):
+        if value is None:
+            self._set_store("_scaling", None); self._set_store("scaling", None); return
+        self._set_store("_scaling", value)
+        s = self._scaling_activation(value + self.scale_bias)
+        s = torch.square(s) + self.mininum_kernel_size ** 2
+        self._set_store("scaling", torch.sqrt(s))
+
+    @property
+    def get_scaling(self):
+        return self._get_store("scaling")
+
+    @property
+    def _rotation(self):
+        return self._get_store("_rotation")
+    @_rotation.setter
+    def _rotation(self, value):
+        self._set_store("_rotation", value)
+
+    def construct_list_of_attributes(self):
+        l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
+        dc = self._features_dc
+        for i in range(dc.shape[1] * dc.shape[2]):
+            l.append(f'f_dc_{i}')
+        l.append('opacity')
+        for i in range(self._scaling.shape[1]):
+            l.append(f'scale_{i}')
+        for i in range(self._rotation.shape[1]):
+            l.append(f'rot_{i}')
+        return l
+
+    _DEFAULT_TRANSFORM = [[1, 0, 0], [0, 0, -1], [0, 1, 0]]
+
+    def _get_ply_data(self, transform=None):
+        xyz = self.get_xyz.detach().cpu().numpy()
+        normals = np.zeros_like(xyz)
+        f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+        opacities = self._inverse_opacity_activation(self.get_opacity).detach().cpu().numpy()
+        scale = torch.log(self.get_scaling).detach().cpu().numpy()
+        rotation = (self._rotation + self.rots_bias[None, :]).detach().cpu().numpy()
+        if transform is not None:
+            transform = np.array(transform)
+            xyz = np.matmul(xyz, transform.T)
+            R_mat = _quat_to_matrix(rotation)
+            R_mat = np.matmul(transform, R_mat)
+            rotation = _matrix_to_quat(R_mat)
+        return xyz, normals, f_dc, opacities, scale, rotation
+
+    def _transformed_xyz_rot(self, transform=None):
+        if transform is None:
+            transform = self._DEFAULT_TRANSFORM
+        transform = np.array(transform, dtype=np.float32)
+        xyz = self.get_xyz.detach().cpu().numpy().astype(np.float32)
+        rotation = (self._rotation + self.rots_bias[None, :]).detach().cpu().numpy()
+        xyz = np.matmul(xyz, transform.T)
+        R_mat = _quat_to_matrix(rotation)
+        R_mat = np.matmul(transform, R_mat)
+        rotation = _matrix_to_quat(R_mat)
+        return xyz, rotation
+
+    def to_ply_bytes(self, transform=None) -> bytes:
+        if transform is None:
+            transform = self._DEFAULT_TRANSFORM
+        xyz, normals, f_dc, opacities, scale, rotation = self._get_ply_data(transform=transform)
+        dtype_full = [(attr, 'f4') for attr in self.construct_list_of_attributes()]
+        elements = np.empty(xyz.shape[0], dtype=dtype_full)
+        elements[:] = list(map(tuple, np.concatenate((xyz, normals, f_dc, opacities, scale, rotation), axis=1)))
+        return _binary_ply_bytes(elements, dtype_full)
+
+    def to_splat_bytes(self, transform=None) -> bytes:
+        if transform is None:
+            transform = self._DEFAULT_TRANSFORM
+        xyz, rotation = self._transformed_xyz_rot(transform=transform)
+        scale = self.get_scaling.detach().cpu().numpy().astype(np.float32)
+        opacity = self.get_opacity.detach().cpu().numpy()
+        f_dc = self._features_dc.detach().cpu().numpy()
+        C0 = 0.28209479177387814
+        # .splat packs color as 4 bytes RGBA: RGB from the SH DC term, A from opacity.
+        rgb = np.clip((f_dc[:, 0, :] * C0 + 0.5) * 255, 0, 255).astype(np.uint8)
+        alpha = np.clip(opacity[:, 0:1] * 255, 0, 255).astype(np.uint8)
+        rgba = np.concatenate([rgb, alpha], axis=1)
+        rot = rotation / np.linalg.norm(rotation, axis=-1, keepdims=True)
+        rot_u8 = np.clip(rot * 128 + 128, 0, 255).astype(np.uint8)
+        order = np.argsort(-opacity[:, 0] * np.prod(scale, axis=-1))
+        xyz, scale, rgba, rot_u8 = xyz[order], scale[order], rgba[order], rot_u8[order]
+        # Per-splat record is exactly 32 bytes: xyz(12) + scale(12) + rgba(4) + rot(4).
+        data = np.concatenate([
+            xyz.astype(np.float32).view(np.uint8).reshape(-1, 12),
+            scale.astype(np.float32).view(np.uint8).reshape(-1, 12),
+            rgba.reshape(-1, 4),
+            rot_u8.reshape(-1, 4),
+        ], axis=1).reshape(-1)
+        return data.tobytes()
+
+    def save_ply(self, path, transform=None):
+        with open(path, 'wb') as f:
+            f.write(self.to_ply_bytes(transform=transform))
+
+    def save_splat(self, path, transform=None):
+        with open(path, 'wb') as f:
+            f.write(self.to_splat_bytes(transform=transform))
+
+
+def _binary_ply_bytes(elements, dtype_full) -> bytes:
+    num_vertices = len(elements)
+    header = "ply\nformat binary_little_endian 1.0\n"
+    header += f"element vertex {num_vertices}\n"
+    type_map = {'f4': 'float', 'u1': 'uchar', 'i4': 'int'}
+    for name, t in dtype_full:
+        header += f"property {type_map.get(t, t)} {name}\n"
+    header += "end_header\n"
+    return header.encode('ascii') + elements.tobytes()
+
+
+def _quat_to_matrix(q):
+    q = q / np.linalg.norm(q, axis=-1, keepdims=True)
+    w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3]
+    R = np.stack([
+        1 - 2*(y*y + z*z), 2*(x*y - w*z),     2*(x*z + w*y),
+        2*(x*y + w*z),     1 - 2*(x*x + z*z), 2*(y*z - w*x),
+        2*(x*z - w*y),     2*(y*z + w*x),     1 - 2*(x*x + y*y),
+    ], axis=-1).reshape(-1, 3, 3)
+    return R
+
+
+def _matrix_to_quat(R):
+    trace = R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2]
+    q = np.zeros((R.shape[0], 4), dtype=R.dtype)
+    s = np.sqrt(np.maximum(trace + 1, 0)) * 2
+    q[:, 0] = 0.25 * s
+    q[:, 1] = (R[:, 2, 1] - R[:, 1, 2]) / np.where(s != 0, s, 1)
+    q[:, 2] = (R[:, 0, 2] - R[:, 2, 0]) / np.where(s != 0, s, 1)
+    q[:, 3] = (R[:, 1, 0] - R[:, 0, 1]) / np.where(s != 0, s, 1)
+    m01 = (R[:, 0, 0] >= R[:, 1, 1]) & (R[:, 0, 0] >= R[:, 2, 2]) & (s == 0)
+    s1 = np.sqrt(np.maximum(1 + R[:, 0, 0] - R[:, 1, 1] - R[:, 2, 2], 0)) * 2
+    q[m01, 0] = (R[m01, 2, 1] - R[m01, 1, 2]) / s1[m01]
+    q[m01, 1] = 0.25 * s1[m01]
+    q[m01, 2] = (R[m01, 0, 1] + R[m01, 1, 0]) / s1[m01]
+    q[m01, 3] = (R[m01, 0, 2] + R[m01, 2, 0]) / s1[m01]
+    m11 = (R[:, 1, 1] > R[:, 0, 0]) & (R[:, 1, 1] >= R[:, 2, 2]) & (s == 0)
+    s2 = np.sqrt(np.maximum(1 + R[:, 1, 1] - R[:, 0, 0] - R[:, 2, 2], 0)) * 2
+    q[m11, 0] = (R[m11, 0, 2] - R[m11, 2, 0]) / s2[m11]
+    q[m11, 1] = (R[m11, 0, 1] + R[m11, 1, 0]) / s2[m11]
+    q[m11, 2] = 0.25 * s2[m11]
+    q[m11, 3] = (R[m11, 1, 2] + R[m11, 2, 1]) / s2[m11]
+    m21 = (R[:, 2, 2] > R[:, 0, 0]) & (R[:, 2, 2] > R[:, 1, 1]) & (s == 0)
+    s3 = np.sqrt(np.maximum(1 + R[:, 2, 2] - R[:, 0, 0] - R[:, 1, 1], 0)) * 2
+    q[m21, 0] = (R[m21, 1, 0] - R[m21, 0, 1]) / s3[m21]
+    q[m21, 1] = (R[m21, 0, 2] + R[m21, 2, 0]) / s3[m21]
+    q[m21, 2] = (R[m21, 1, 2] + R[m21, 2, 1]) / s3[m21]
+    q[m21, 3] = 0.25 * s3[m21]
+    return q / np.linalg.norm(q, axis=-1, keepdims=True)
+
+
+def _build_gaussians(decoder: ElasticGaussianFixedlenDecoder, points_pred: dict, pred: dict):
+    x = points_pred
+    offset = decoder._get_offset(pred['features'])
+    h = pred["features"]
+    ret = []
+    for i in range(h.shape[0]):
+        g = Gaussian(
+            sh_degree=0,
+            aabb=[-0.5, -0.5, -0.5, 1.0, 1.0, 1.0],
+            mininum_kernel_size=decoder.rep_config['filter_kernel_size_3d'],
+            scaling_bias=decoder.rep_config['scaling_bias'],
+            opacity_bias=decoder.rep_config['opacity_bias'],
+            scaling_activation=decoder.rep_config['scaling_activation'],
+        )
+        _x = x["points"][i, :, None, :]
+        for k, v in decoder.layout.items():
+            if k == '_xyz':
+                setattr(g, k, (offset[i] + _x).flatten(0, 1))
+            elif k in ('_xyz_center', '_offset_scale'):
+                continue
+            else:
+                feats = h[i][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape']).flatten(0, 1)
+                setattr(g, k, feats * decoder.rep_config['lr'][k])
+        ret.append(g)
+    return ret
+
+
+# ---------------------------------------------------------------------------
+# Euler flow sampler
+# ---------------------------------------------------------------------------
+
+class FlowEulerCfgSampler:
+    def __init__(self, sigma_min: float = 1e-5):
+        self.sigma_min = sigma_min
+
+    def _get_batch_size(self, x_t):
+        return next(iter(x_t.values())).shape[0] if isinstance(x_t, dict) else x_t.shape[0]
+
+    def _get_device(self, x_t):
+        return next(iter(x_t.values())).device if isinstance(x_t, dict) else x_t.device
+
+    def _inference_model(self, model, x_t, t, cond=None):
+        batch = self._get_batch_size(x_t)
+        device = self._get_device(x_t)
+        t_scaled = torch.tensor([1000 * t] * batch, device=device, dtype=torch.float32)
+        if isinstance(cond, dict):
+            for k, v in cond.items():
+                if isinstance(v, torch.Tensor) and v.shape[0] == 1 and batch > 1:
+                    cond[k] = v.repeat(batch, *([1] * (len(v.shape) - 1)))
+        elif cond is not None and cond.shape[0] == 1 and batch > 1:
+            cond = cond.repeat(batch, *([1] * (len(cond.shape) - 1)))
+        return model(x_t, t_scaled, cond)
+
+    def _cfg_prediction(self, model, x_t, t, cond, neg_cond, guidance_scale):
+        # Diffusers-style convention: guidance_scale == 1 (or <= 1, or None) means no CFG —
+        # only the conditional pass runs, halving the per-step cost. > 1 enables CFG and
+        # blends as `pred = s * cond + (1 - s) * uncond = s * cond - (s - 1) * uncond`.
+        pred_v = self._inference_model(model, x_t, t, cond)
+        if isinstance(guidance_scale, dict):
+            if not any(s > 1 for s in guidance_scale.values()):
+                return pred_v
+            neg_pred_v = self._inference_model(model, x_t, t, neg_cond)
+            for key in pred_v:
+                s = guidance_scale.get(key, 1.0)
+                if s > 1:
+                    pred_v[key] = s * pred_v[key] - (s - 1) * neg_pred_v[key]
+            return pred_v
+        if guidance_scale is None or guidance_scale <= 1:
+            return pred_v
+        neg_pred_v = self._inference_model(model, x_t, t, neg_cond)
+        for key in pred_v:
+            pred_v[key] = guidance_scale * pred_v[key] - (guidance_scale - 1) * neg_pred_v[key]
+        return pred_v
+
+    @torch.no_grad()
+    def sample(self, model, noise, cond, neg_cond, steps=50, shift=1.0,
+               guidance_scale=None, show_progress=False, callback=None):
+        sample = noise
+        t_seq = shift * np.linspace(1, 0, steps + 1) / (1 + (shift - 1) * np.linspace(1, 0, steps + 1))
+        t_pairs = list(zip(t_seq[:-1], t_seq[1:]))
+        iterator = tqdm(t_pairs, desc="Sampling", total=steps) if show_progress else t_pairs
+        for i, (t, t_prev) in enumerate(iterator):
+            x_t = {k: v.clone() for k, v in sample.items()} if isinstance(sample, dict) else sample.clone()
+            pred_v = self._cfg_prediction(model, x_t, t, cond, neg_cond, guidance_scale)
+            dt = t - t_prev
+            if isinstance(sample, dict):
+                for key in sample:
+                    sample[key] = sample[key] - pred_v[key] * dt
+            else:
+                sample = sample - pred_v * dt
+            if callback is not None:
+                callback(i + 1, steps)
+        return sample
+
+
+# ---------------------------------------------------------------------------
+# Component loaders
+# ---------------------------------------------------------------------------
+
+def _place(m, device, dtype):
+    if device is not None or dtype is not None:
+        m = m.to(device=device, dtype=dtype)
+    return m.eval()
+
+
+def load_dinov3(path: str, device=None, dtype=None) -> DinoV3ViT:
+    m = DinoV3ViT()
+    m.load_safetensors(path)
+    return _place(m, device, dtype)
+
+
+def load_vae_encoder(path: str, device=None, dtype=None) -> Flux2VAEEncoder:
+    m = Flux2VAEEncoder()
+    m.load_safetensors(path)
+    return _place(m, device, dtype)
+
+
+def load_rmbg(path: str, device=None, dtype=None) -> BiRefNet:
+    m = BiRefNet()
+    m.load_safetensors(path)
+    return _place(m, device, dtype)
+
+
+FLOW_MODEL_ARGS = dict(
+    q_token_length=8192, in_channels=16, cam_channels=5, out_channels=16,
+    model_channels=1024, cond_channels=1280, cond2_channels=128,
+    num_refiner_blocks=2, num_blocks=24, num_heads=16, mlp_ratio=4,
+    qk_rms_norm=True, share_mod=True, use_shift_table=True,
+)
+
+
+def load_flow_model(path: str, device=None, dtype=None) -> LatentSeqMMFlowModel:
+    m = LatentSeqMMFlowModel(**FLOW_MODEL_ARGS)
+    m.load_safetensors(path)
+    return _place(m, device, dtype)
+
+
+OCTREE_DECODER_ARGS = dict(
+    model_channels=1024, cond_channels=16,
+    num_blocks=4, num_heads=16, mlp_ratio=4, share_mod=True,
+)
+
+GS_DECODER_ARGS = dict(
+    in_channels=3, model_channels=1024, cond_channels=16,
+    attn_mode="full", num_blocks=16, num_heads=16, mlp_ratio=4,
+    use_learned_offset_scale=True, use_per_offset=True,
+    representation_config=dict(
+        lr=dict(_xyz=1.0, _features_dc=1.0, _opacity=1.0, _scaling=1.0, _rotation=0.1),
+        perturb_offset=True, perturbe_size=1.5, offset_scale=0.05, num_gaussians=32,
+        filter_kernel_size_3d=0.0009, scaling_bias=0.004, opacity_bias=0.1,
+        scaling_activation="softplus",
+    ),
+)
+
+
+def load_decoder(path: str, device=None, dtype=None) -> OctreeGaussianDecoder:
+    m = OctreeGaussianDecoder(OCTREE_DECODER_ARGS, GS_DECODER_ARGS)
+    m.load_safetensors(path)
+    return _place(m, device, dtype)
+
+
+# ---------------------------------------------------------------------------
+# Pipeline stages
+# ---------------------------------------------------------------------------
+
+_CANVAS_SIZE = 1024
+
+
+def _image_to_pil(image) -> Image.Image:
+    if isinstance(image, Image.Image):
+        return image
+    if isinstance(image, (str, bytes)) or hasattr(image, "__fspath__"):
+        return Image.open(image)
+    if isinstance(image, torch.Tensor):
+        t = image.detach().cpu()
+        if t.ndim == 4:
+            assert t.shape[0] == 1, (
+                f"batched image input is not supported (got B={t.shape[0]}); "
+                "pass one image at a time"
+            )
+            t = t[0]
+        arr = (t.clamp(0, 1) * 255).to(torch.uint8).numpy()
+        mode = "RGBA" if arr.shape[-1] == 4 else "RGB"
+        return Image.fromarray(arr, mode=mode)
+    raise TypeError(f"unsupported image type: {type(image)}")
+
+
+def preprocess_image(image, rmbg: BiRefNet, erode_radius: int = 1) -> Image.Image:
+    image = _image_to_pil(image)
+    size = _CANVAS_SIZE
+    w, h = image.size
+    s = size / min(w, h)
+    image = image.resize((max(1, int(round(w * s))), max(1, int(round(h * s)))), Image.LANCZOS)
+    has_real_alpha = (image.mode == "RGBA"
+                      and np.array(image.getchannel(3), dtype=np.int32).min() < 255)
+    if not has_real_alpha:
+        image = rmbg.remove_background(image.convert("RGB"))
+    if erode_radius > 0:
+        image.putalpha(image.getchannel(3).filter(ImageFilter.MinFilter(2 * erode_radius + 1)))
+    alpha = np.array(image.getchannel(3))
+    ys, xs = np.nonzero(alpha)
+    bbox = [xs.min(), ys.min(), xs.max(), ys.max()]
+    cx, cy = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
+    half = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2 * 1.2
+    image = image.crop([int(cx - half), int(cy - half), int(cx + half), int(cy + half)])
+    image = image.resize((size, size), Image.LANCZOS)
+    bg = Image.new("RGB", (size, size), (0, 0, 0))
+    bg.paste(image, mask=image.split()[3])
+    return bg
+
+
+_DINOV3_NORMALIZE = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+
+@torch.no_grad()
+def encode_image(image: Image.Image, dinov3: DinoV3ViT, vae_encoder: Flux2VAEEncoder,
+                 generator: torch.Generator = None) -> dict:
+    device = next(dinov3.parameters()).device
+    img_tensor   = transforms.ToTensor()(image).unsqueeze(0).to(device=device, dtype=torch.float32)
+    img_normed   = _DINOV3_NORMALIZE(img_tensor)
+    dinov3_dtype = next(dinov3.parameters()).dtype
+    vae_dtype    = next(vae_encoder.parameters()).dtype
+    dinov3_feat = dinov3(pixel_values=img_normed.to(dinov3_dtype))
+    dinov3_feat = F.layer_norm(dinov3_feat.float(), dinov3_feat.shape[-1:])
+    vae_feat = vae_encoder.encode(img_tensor.to(vae_dtype) * 2 - 1,
+                                  deterministic=False, generator=generator)
+    # pad 5 zero tokens so feature2's token length matches feature1's (cls + 4 registers + patches)
+    zero_reg = torch.zeros(vae_feat.shape[0], 5, vae_feat.shape[2],
+                           dtype=vae_feat.dtype, device=vae_feat.device)
+    vae_feat = torch.cat([zero_reg, vae_feat], dim=1)
+    return {'feature1': dinov3_feat, 'feature2': vae_feat}
+
+
+@torch.no_grad()
+def sample_latent(flow_model: LatentSeqMMFlowModel, cond: dict,
+                  steps: int = 50, guidance_scale: float = 7.0, shift: float = 3.0,
+                  generator: torch.Generator = None,
+                  show_progress: bool = False, callback=None) -> dict:
+    device = flow_model.device
+    neg_cond = {k: torch.zeros_like(v) for k, v in cond.items()}
+    noise = {'latent': torch.randn(1, flow_model.q_token_length, flow_model.in_channels,
+                                   device=device, generator=generator)}
+    if flow_model.cam_channels is not None:
+        noise['camera'] = torch.randn(1, 1, flow_model.cam_channels,
+                                      device=device, generator=generator)
+    sampler = FlowEulerCfgSampler()
+    return sampler.sample(flow_model, noise, cond=cond, neg_cond=neg_cond,
+                          steps=steps, guidance_scale=guidance_scale, shift=shift,
+                          show_progress=show_progress, callback=callback)
+
+
+# ---------------------------------------------------------------------------
+# Pipeline
+# ---------------------------------------------------------------------------
+
+class TripoSplatPipeline:
+    def __init__(self, ckpt_path: str, decoder_path: str, dinov3_path: str,
+                 flux2_vae_encoder_path: str, rmbg_path: str, device: str = "cuda"):
+        self._device = torch.device(device)
+        self.dinov3      = load_dinov3      (dinov3_path,             device=self._device, dtype=torch.bfloat16)
+        self.vae_encoder = load_vae_encoder (flux2_vae_encoder_path,  device=self._device, dtype=torch.bfloat16)
+        self.rmbg        = load_rmbg        (rmbg_path,               device=self._device, dtype=torch.float16)
+        self.flow_model  = load_flow_model  (ckpt_path,               device=self._device, dtype=torch.float16)
+        self.decoder     = load_decoder     (decoder_path,            device=self._device, dtype=torch.float16)
+
+    def preprocess_image(self, image, erode_radius: int = 1) -> Image.Image:
+        return preprocess_image(image, self.rmbg, erode_radius=erode_radius)
+
+    def encode_image(self, image: Image.Image, generator: torch.Generator = None) -> dict:
+        return encode_image(image, self.dinov3, self.vae_encoder, generator=generator)
+
+    def sample_latent(self, cond: dict, steps: int = 50, guidance_scale: float = 7.0,
+                      shift: float = 3.0, generator: torch.Generator = None,
+                      show_progress: bool = False, callback=None) -> dict:
+        return sample_latent(self.flow_model, cond, steps=steps, guidance_scale=guidance_scale,
+                             shift=shift, generator=generator,
+                             show_progress=show_progress, callback=callback)
+
+    def decode_latent(self, latent: torch.Tensor, num_gaussians: int = 262144):
+        return self.decoder.decode(latent, num_gaussians=num_gaussians)
+
+    _NUM_GAUSSIANS_MIN = 32768
+    _NUM_GAUSSIANS_MAX = 262144
+
+    def _validate_num_gaussians(self, n: int) -> int:
+        assert self._NUM_GAUSSIANS_MIN <= n <= self._NUM_GAUSSIANS_MAX, (
+            f"num_gaussians must be in [{self._NUM_GAUSSIANS_MIN}, {self._NUM_GAUSSIANS_MAX}], got {n}"
+        )
+        gpp = self.decoder.gaussians_per_point
+        if n % gpp == 0:
+            return n
+        rounded = round(n / gpp) * gpp
+        print(f"[TripoSplatPipeline] num_gaussians={n} is not a multiple of {gpp}; rounding to {rounded}")
+        return rounded
+
+    @torch.no_grad()
+    def run(self, image, seed: int = 42, steps: int = 20, guidance_scale: float = 3.0,
+            shift: float = 3.0, num_gaussians=262144, erode_radius: int = 1,
+            show_progress: bool = False, callback=None):
+        """
+        Args:
+            image: Input image. Accepts a file path / PIL.Image / torch.Tensor
+                (`[1,H,W,C]` or `[H,W,C]`, float in `[0, 1]`, optional alpha
+                channel as the 4th channel).
+            seed: RNG seed for the VAE encoder's stochastic latent sampling and
+                the initial flow-matching noise. Same seed → same output.
+            steps: Number of Euler integrator steps in the flow-matching sampler.
+                More steps → better fidelity, linear runtime cost.
+                Recommend: 10~20.
+            guidance_scale: Classifier-free-guidance strength (diffusers
+                convention). `≤ 1.0` disables CFG. Higher → more detail,
+                stronger adherence to the input image; too high can cause color
+                oversaturation.
+                Recommend: 3.0.
+            shift: Flow-matching timestep schedule shift. `1.0` gives a uniform
+                schedule; `>1.0` allocates more steps to the early/high-noise end.
+                Recommend: 3.0.
+            num_gaussians: Target Gaussian-splat count. An `int` returns a
+                single `Gaussian`. A `list` / `tuple` of ints returns a
+                `list[Gaussian]`. Each count is rounded to the nearest multiple
+                of 32. More gaussians → more detail but higher rendering and
+                storage cost.
+                Recommend: 32768~262144.
+            erode_radius: Pixel radius used to erode the alpha matte after
+                background removal, to avoid segmentation-border bleed before
+                compositing on black. `0` disables; `1` is a 3×3 minimum filter.
+                Recommend: 1.
+            show_progress: Print a `tqdm` progress bar over sampler steps.
+            callback: Optional `fn(step, total)` invoked after each sampler step.
+                Useful for external progress UIs (e.g. ComfyUI's
+                `ProgressBar.update`).
+
+        Returns:
+            `(gaussian, prepared_image)` for an `int` `num_gaussians`, or
+            `(list_of_gaussians, prepared_image)` for a `list` / `tuple`. The
+            second element is the RGB composite the encoders actually saw —
+            useful for display / debugging.
+        """
+        if isinstance(num_gaussians, (list, tuple)):
+            counts = [self._validate_num_gaussians(n) for n in num_gaussians]
+        else:
+            counts = [self._validate_num_gaussians(num_gaussians)]
+
+        gen = torch.Generator(device=self._device).manual_seed(seed)
+        prepared = self.preprocess_image(image, erode_radius=erode_radius)
+        cond = self.encode_image(prepared, generator=gen)
+        out = self.sample_latent(cond, steps=steps, guidance_scale=guidance_scale, shift=shift,
+                                 generator=gen, show_progress=show_progress, callback=callback)
+        gaussians = [self.decode_latent(out['latent'], num_gaussians=n) for n in counts]
+        if isinstance(num_gaussians, (list, tuple)):
+            return gaussians, prepared
+        return gaussians[0], prepared