Delete triposplat.py

triposplat.py

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