Upload upstream/upstream_src.txt with huggingface_hub
ded5557 verified 8 days ago
125 kB
TripoSplat @ d8db9e0
========================================================================

========================================================================
===== .gitignore =====
========================================================================
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[codz]
*$py.class

# C extensions
*.so

# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST

# PyInstaller
#   Usually these files are written by a python script from a template
#   before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py.cover
.hypothesis/
.pytest_cache/
cover/

# Translations
*.mo
*.pot

# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal

# Flask stuff:
instance/
.webassets-cache

# Scrapy stuff:
.scrapy

# Sphinx documentation
docs/_build/

# PyBuilder
.pybuilder/
target/

# Jupyter Notebook
.ipynb_checkpoints

# IPython
profile_default/
ipython_config.py

# pyenv
#   For a library or package, you might want to ignore these files since the code is
#   intended to run in multiple environments; otherwise, check them in:
# .python-version

# pipenv
#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
#   However, in case of collaboration, if having platform-specific dependencies or dependencies
#   having no cross-platform support, pipenv may install dependencies that don't work, or not
#   install all needed dependencies.
# Pipfile.lock

# UV
#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
#   This is especially recommended for binary packages to ensure reproducibility, and is more
#   commonly ignored for libraries.
# uv.lock

# poetry
#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
#   This is especially recommended for binary packages to ensure reproducibility, and is more
#   commonly ignored for libraries.
#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
# poetry.lock
# poetry.toml

# pdm
#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
# pdm.lock
# pdm.toml
.pdm-python
.pdm-build/

# pixi
#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
# pixi.lock
#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
#   in the .venv directory. It is recommended not to include this directory in version control.
.pixi

# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/

# Celery stuff
celerybeat-schedule
celerybeat.pid

# Redis
*.rdb
*.aof
*.pid

# RabbitMQ
mnesia/
rabbitmq/
rabbitmq-data/

# ActiveMQ
activemq-data/

# SageMath parsed files
*.sage.py

# Environments
.env
.envrc
.venv
env/
venv/
ENV/
env.bak/
venv.bak/

# Spyder project settings
.spyderproject
.spyproject

# Rope project settings
.ropeproject

# mkdocs documentation
/site

# mypy
.mypy_cache/
.dmypy.json
dmypy.json

# Pyre type checker
.pyre/

# pytype static type analyzer
.pytype/

# Cython debug symbols
cython_debug/

# PyCharm
#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
#   and can be added to the global gitignore or merged into this file.  For a more nuclear
#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
# .idea/

# Abstra
#   Abstra is an AI-powered process automation framework.
#   Ignore directories containing user credentials, local state, and settings.
#   Learn more at https://abstra.io/docs
.abstra/

# Visual Studio Code
#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
#   and can be added to the global gitignore or merged into this file. However, if you prefer, 
#   you could uncomment the following to ignore the entire vscode folder
# .vscode/
# Temporary file for partial code execution
tempCodeRunnerFile.py

# Ruff stuff:
.ruff_cache/

# PyPI configuration file
.pypirc

# Marimo
marimo/_static/
marimo/_lsp/
__marimo__/

# Streamlit
.streamlit/secrets.toml


========================================================================
===== LICENSE =====
========================================================================
MIT License

Copyright (c) 2026 VAST

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.


========================================================================
===== README.md =====
========================================================================
# TripoSplat
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. This is the inference-only repo for TripoSplat. For the training code, see [TripoSplat-Training](https://github.com/runjie-yan/TripoSplat-Training).

<a href="https://arxiv.org/abs/2605.16355"><img src="https://img.shields.io/badge/Read%20Paper-B31B1B?style=for-the-badge&logo=arxiv" alt="Paper"></a>
<a href="https://www.tripo3d.ai/research/triposplat"><img src="https://img.shields.io/badge/Technical%20Blog-grey?style=for-the-badge&logo=data:image/svg%2bxml;base64,PHN2ZyB3aWR0aD0iNjUiIGhlaWdodD0iNjUiIHZpZXdCb3g9IjAgMCA2NSA2NSIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj4KPHBhdGggZD0iTTkuNDk5MSA5LjYzNDc3TDE2LjQzNzQgMjEuNDU1NkMxNi40MzkzIDIxLjQ1ODkgMTYuNDQxMiAyMS40NjIyIDE2LjQ0MzEgMjEuNDY1NUwzMC4yNjU4IDQ1LjA1NDhDMzEuNTMyNyA0Ny4yMTY3IDM0LjcwNDUgNDcuMjE2NyAzNS45NzE0IDQ1LjA1NDhMNDkuMzg2MiAyMi4xNjE2SDU5LjQ2MThMNDEuMjY2IDUzLjE2MkMzNy42NDQ5IDU5LjMzMTMgMjguNTkyMyA1OS4zMzEzIDI0Ljk3MTIgNTMuMTYyTDYuNjM5NjcgMjEuOTMwMkM0LjAyNCAxNy40NzM5IDUuNjU5NTYgMTIuMjEyNyA5LjQ5OTEgOS42MzQ3N1oiIGZpbGw9IndoaXRlIi8+CjxwYXRoIGQ9Ik0yMC4xMTIxIDE2LjYwODdIMzQuNjkyNkwyOC42MjIgMjcuMDQ0MkMyOC4yMDMzIDI3Ljc2NCAyOC4yMDgzIDI4LjY0OTIgMjguNjM1MSAyOS4zNjQ0TDMxLjA1MjcgMzMuNDE1MUMzMS45NjU0IDM0Ljk0NDUgMzQuMjE2MyAzNC45MzY1IDM1LjExNzggMzMuNDAwNkw0NC45NzM5IDE2LjYwODdINDYuOTQyTDQ2Ljk0NTUgMTYuNjA4N0g2MC44NDQ2QzYwLjQ4MzIgMTIuMDU4NyA1Ni42NzMxIDguMDQ4ODMgNTEuNDUwOSA4LjA0ODgzTDE1LjA4NzkgOC4wNDg4M0wyMC4xMTIxIDE2LjYwODdaIiBmaWxsPSIjRjhDRjAwIi8+Cjwvc3ZnPgo=" alt="Technical Blog"></a>
<a href="https://huggingface.co/spaces/VAST-AI/TripoSplat"><img src="https://img.shields.io/badge/Huggingface%20Demo-grey?style=for-the-badge&logo=huggingface" alt="HuggingFace Demo"></a>

| ![](static/doc/001.webp) | ![](static/doc/002.webp) |
|---|---|
| ![](static/doc/003.webp) | ![](static/doc/004.webp) |

## Highlights
- **High-quality, versatile generation** that handles a wide range of image styles.
- **Arbitrary Gaussian count** (up to 262,144) — trade off visual quality against rendering cost according to your need.
- **Minimal, readable code**: two files (`triposplat.py` and `model.py`), ~2,000 LOC total. Easy to customize and integrate into other ecosystems.
- **Near-zero dependencies**: no `transformers`, no `diffusers`, no version-conflict hell. Runs on any platform.
- **Official ComfyUI support**: drop the [official workflow template](https://github.com/Comfy-Org/workflow_templates/blob/main/templates/3d_triposplat_image_to_gaussian_splat.json) into ComfyUI and start playing with TripoSplat right away.

## Quickstart
Download model weights to `ckpts/` from [HuggingFace](https://huggingface.co/VAST-AI/TripoSplat). 
```bash
# Use one of the following ways to download model weights.

# 1. Use HuggingFace CLI
hf download VAST-AI/TripoSplat --local-dir ckpts/

# 2. Use huggingface_hub
pip install huggingface_hub
python -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='VAST-AI/TripoSplat', local_dir='ckpts/')"

# 3. Use ModelScope CLI
pip install modelscope
modelscope download VAST-AI-Research/TripoSplat --local_dir ckpts/

# 4. Use modelscope Python SDK
pip install modelscope
python -c "from modelscope import snapshot_download; snapshot_download('VAST-AI-Research/TripoSplat', local_dir='ckpts/')"

# 5. Manual download from HuggingFace / ModelScope.
```

Setup the environment and run the example inference script.
```bash
# install torch and torchvision according to your environment
pip install numpy safetensors pillow tqdm
python run_example.py
```

The exported `.ply` / `.splat` files can be visualized in any 3D Gaussian
viewer — e.g. [SparkJS](https://sparkjs.dev) or
[SuperSplat](https://superspl.at/editor).


## Gradio Demo

```bash
pip install gradio
python run_gradio.py
```

## License
TripoSplat code and weight models are released under the [MIT License](https://github.com/VAST-AI-Research/TripoSplat/blob/main/LICENSE).

## Citation
If you find TripoSplat useful, please cite:
```bibtex
@misc{yan2026generative3dgaussianslearned,
    title={Generative 3D Gaussians with Learned Density Control}, 
    author={Runjie Yan and Yan-Pei Cao and Peng Wang and Ding Liang and Yuan-Chen Guo},
    year={2026},
    eprint={2605.16355},
    archivePrefix={arXiv},
    primaryClass={cs.GR},
    url={https://arxiv.org/abs/2605.16355}, 
}
```


========================================================================
===== model.py =====
========================================================================
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]


========================================================================
===== run_example.py =====
========================================================================
"""TripoSplat minimal examples.
Usage: python run_example.py
"""
from triposplat import TripoSplatPipeline


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

INPUT = "static/example_inputs/building_stone_house.webp"


# ---------------------------------------------------------------------------
# Example 1 — one image → PLY + SPLAT
# ---------------------------------------------------------------------------

gaussian, prepared = pipe.run(INPUT, num_gaussians=262144, show_progress=True)

prepared.save("preprocessed_image.webp")
gaussian.save_ply("output.ply")
gaussian.save_splat("output.splat")


# ---------------------------------------------------------------------------
# Example 2 — one image → Gaussians at several densities (denoiser runs once,
# decoder is replayed per count) → one PLY per count
# ---------------------------------------------------------------------------

counts = [32768, 65536, 131072, 262144]
gaussians, _ = pipe.run(INPUT, num_gaussians=counts, show_progress=True)

for n, g in zip(counts, gaussians):
    g.save_ply(f"output_{n}.ply")


========================================================================
===== run_gradio.py =====
========================================================================
"""TripoSplat Gradio demo with Spark.js in-browser viewer.
Usage: python run_gradio.py
"""
import time
from pathlib import Path
from uuid import uuid4

import gradio as gr
import torch

from triposplat import TripoSplatPipeline


# ----------------------------------------------------------------------------
# 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()
EXAMPLES_DIR = Path("static/example_inputs").resolve()
EXAMPLES = [
    str(EXAMPLES_DIR / "creature_butterfly.webp"),
    str(EXAMPLES_DIR / "building_stone_house.webp"),
    str(EXAMPLES_DIR / "vehicle_pirate_ship.webp"),
    str(EXAMPLES_DIR / "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 generate(image, seed: int, steps: int, guidance_scale: float,
             num_gaussians: int, output_format: str,
             progress=gr.Progress(track_tqdm=True)):
    """Run the full pipeline (preprocess + encode + sample + decode)."""
    if image is None:
        raise gr.Error("Please upload an image first.")

    progress(0, desc="Generating...")
    t0 = time.time()
    prepared = PIPE.preprocess_image(image)
    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 prepared, _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)"
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_in = gr.Image(label="Input image", type="pil", image_mode="RGBA",
                                height=320)

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

    run_btn.click(
        fn=generate,
        inputs=[image_in, seed_in, steps_in, cfg_in, num_g_in, fmt_in],
        outputs=[prepared_out, viewer_out, file_out, info_out],
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        allowed_paths=[
            str(VIEWER_HTML.parent),
            str(OUT_ROOT),
            str(EXAMPLES_DIR),
        ],
    )


========================================================================
===== triposplat.py =====
========================================================================
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'],
            device=x['points'].device,
        )
        _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