Leaner idol package (inference code only)

idol/conditioning.py

@@ -13,12 +13,11 @@ DTYPES = {'fp32': torch.float32, 'bf16': torch.bfloat16, 'fp16': torch.float16}
 
 
 class ImageEncoder:
-    def __init__(self, dinov3_repo: str, flux2_vae_repo: str, resolution: int = 1024,
+    def __init__(self, dinov3_repo: str, flux2_vae_repo: str,
                  dinov3_dtype: str = 'bf16', device: str = 'cuda'):
         from diffusers.models import AutoencoderKLFlux2
         from transformers import AutoModel
 
-        self.resolution = resolution
         self.device = torch.device(device)
         self.dinov3 = AutoModel.from_pretrained(dinov3_repo, torch_dtype=DTYPES[dinov3_dtype])
         self.dinov3 = self.dinov3.eval().requires_grad_(False).to(self.device)
@@ -26,18 +25,6 @@ class ImageEncoder:
                                                             torch_dtype=torch.bfloat16)
         self.flux2_vae = self.flux2_vae.eval().requires_grad_(False).to(self.device)
 
-    def to(self, device):
-        self.device = torch.device(device)
-        self.dinov3.to(self.device)
-        self.flux2_vae.to(self.device)
-        return self
-
-    def _resize(self, images):
-        size = (self.resolution, self.resolution)
-        if images.shape[-2:] != size:
-            images = F.interpolate(images, size=size, mode='bilinear', align_corners=False)
-        return images
-
     def _dinov3_tokens(self, images):
         mean = torch.tensor(IMAGENET_MEAN, device=images.device).view(1, 3, 1, 1)
         std = torch.tensor(IMAGENET_STD, device=images.device).view(1, 3, 1, 1)
@@ -60,8 +47,8 @@ class ImageEncoder:
 
     @torch.no_grad()
     def __call__(self, images: torch.Tensor) -> TensorDict:
-        """images: [B, 3, H, W] in [0, 1] -> TensorDict(feature1=DINOv3 tokens, feature2=FLUX.2 VAE tokens)."""
-        images = self._resize(images.to(self.device))
+        """images: [B, 3, 1024, 1024] in [0, 1] -> TensorDict(feature1=DINOv3 tokens, feature2=FLUX.2 VAE tokens)."""
+        images = images.to(self.device)
         feature1 = self._dinov3_tokens(images)
         feature2 = self._flux2_tokens(images)
         return TensorDict({'feature1': feature1, 'feature2': feature2}, batch_size=feature1.shape[0])

idol/gaussian.py

@@ -1,5 +1,7 @@
-import torch
+"""The generated object: 3D Gaussians with position, colour, scale, rotation and opacity."""
 import numpy as np
+import torch
+import torch.nn.functional as F
 from plyfile import PlyData, PlyElement
 
 
@@ -7,377 +9,54 @@ def inverse_sigmoid(x):
     return torch.log(x / (1 - x))
 
 
-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.init_params = {
-            'aabb': aabb,
-            'sh_degree': sh_degree,
-            'mininum_kernel_size': mininum_kernel_size,
-            'scaling_bias': scaling_bias,
-            'opacity_bias': opacity_bias,
-            'scaling_activation': scaling_activation,
-        }
-        
-        self.sh_degree = sh_degree
-        self.active_sh_degree = sh_degree
-        self.mininum_kernel_size = mininum_kernel_size 
-        self.scaling_bias = scaling_bias
-        self.opacity_bias = opacity_bias
-        self.scaling_activation_type = scaling_activation
-        self.device = device
-        self.aabb = torch.tensor(aabb, dtype=torch.float32, device=device)
-        self.setup_functions()
-
-        self._storage = {}
-        # Pre-initialize keys
-        for key in ['_xyz', 'xyz', '_features_dc', 'features_dc', 
-                    '_features_rest', 'features_rest', 
-                    '_scaling', 'scaling', 
-                    '_rotation', 'rotation', 
-                    '_opacity', 'opacity',
-                    'features']:
-            self._storage[key] = None
-
-    def setup_functions(self):
-        if self.scaling_activation_type == "exp":
-            self.scaling_activation = torch.exp
-            self.inverse_scaling_activation = torch.log
-        elif self.scaling_activation_type == "softplus":
-            self.scaling_activation = torch.nn.functional.softplus
-            self.inverse_scaling_activation = lambda x: x + torch.log(-torch.expm1(-x))
-
-        self.opacity_activation = torch.sigmoid
-        self.inverse_opacity_activation = inverse_sigmoid
-
-        self.rotation_activation = torch.nn.functional.normalize
-        
-        self.scale_bias = self.inverse_scaling_activation(torch.tensor(self.scaling_bias)).to(self.device)
-        self.rots_bias = torch.zeros((4)).to(self.device)
-        self.rots_bias[0] = 1
-        self.opacity_bias = self.inverse_opacity_activation(torch.tensor(self.opacity_bias)).to(self.device)
-
-    def _get_store(self, name):
-        return self._storage.get(name)
-
-    def _set_store(self, name, value):
-        self._storage[name] = value
-
-    # --- Opacity ---
-    @property
-    def opacity(self):
-        return self._get_store("opacity")
-
-    @opacity.setter
-    def opacity(self, value):
-        self.set_opacity(value)
-
-    @property
-    def _opacity(self):
-        return self._get_store("_opacity")
-
-    @_opacity.setter
-    def _opacity(self, value):
-        self.set_raw_opacity(value)
-
-    def set_opacity(self, value):
-        if value is None:
-            self._set_store("opacity", None)
-            self._set_store("_opacity", None)
-            return
-        self._set_store("opacity", value)
-        # Update raw: inverse activation
-        raw = self.inverse_opacity_activation(value) - self.opacity_bias
-        self._set_store("_opacity", raw)
-
-    def set_raw_opacity(self, value):
-        if value is None:
-            self._set_store("_opacity", None)
-            self._set_store("opacity", None)
-            return
-        self._set_store("_opacity", value)
-        # Update activated
-        activated = self.opacity_activation(value + self.opacity_bias)
-        self._set_store("opacity", activated)
-    
-    # --- Scaling ---
-    @property
-    def scaling(self):
-        return self._get_store("scaling")
-
-    @scaling.setter
-    def scaling(self, value):
-        self.set_scaling(value)
-
-    @property
-    def _scaling(self):
-        return self._get_store("_scaling")
-
-    @_scaling.setter
-    def _scaling(self, value):
-        self.set_raw_scaling(value)
-
-    def set_scaling(self, value):
-        if value is None:
-            self._set_store("scaling", None)
-            self._set_store("_scaling", None)
-            return
-        self._set_store("scaling", value)
-        # Inverse activation
-        s = torch.sqrt(torch.square(value) - self.mininum_kernel_size ** 2)
-        raw = self.inverse_scaling_activation(s) - self.scale_bias
-        self._set_store("_scaling", raw)
-
-    def set_raw_scaling(self, value):
-        if value is None:
-            self._set_store("_scaling", None)
-            self._set_store("scaling", None)
-            return
-        self._set_store("_scaling", value)
-        # Forward activation
-        s = self.scaling_activation(value + self.scale_bias)
-        s = torch.square(s) + self.mininum_kernel_size ** 2
-        activated = torch.sqrt(s)
-        self._set_store("scaling", activated)
-
-    # --- Rotation ---
-    @property
-    def rotation(self):
-        return self._get_store("rotation")
-
-    @rotation.setter
-    def rotation(self, value):
-        self.set_rotation(value)
-
-    @property
-    def _rotation(self):
-        return self._get_store("_rotation")
-
-    @_rotation.setter
-    def _rotation(self, value):
-        self.set_raw_rotation(value)
-
-    def set_rotation(self, value):
-        if value is None:
-            self._set_store("rotation", None)
-            self._set_store("_rotation", None)
-            return
-        self._set_store("rotation", value)
-        # Inverse: just subtract bias
-        raw = value - self.rots_bias[None, :]
-        self._set_store("_rotation", raw)
+def inverse_softplus(x):
+    return x + torch.log(-torch.expm1(-x))
 
-    def set_raw_rotation(self, value):
-        if value is None:
-            self._set_store("_rotation", None)
-            self._set_store("rotation", None)
-            return
-        self._set_store("_rotation", value)
-        # Forward activation
-        activated = self.rotation_activation(value + self.rots_bias[None, :])
-        self._set_store("rotation", activated)
 
-    # --- XYZ ---
-    @property
-    def xyz(self):
-        return self._get_store("xyz")
-
-    @xyz.setter
-    def xyz(self, value):
-        self.set_xyz(value)
-
-    @property
-    def _xyz(self):
-        return self._get_store("_xyz")
-
-    @_xyz.setter
-    def _xyz(self, value):
-        self.set_raw_xyz(value)
-
-    def set_xyz(self, value):
-        if value is None:
-            self._set_store("xyz", None)
-            self._set_store("_xyz", None)
-            return
-        self._set_store("xyz", value)
-        # Inverse
-        raw = (value - self.aabb[None, :3]) / self.aabb[None, 3:]
-        self._set_store("_xyz", raw)
-
-    def set_raw_xyz(self, value):
-        if value is None:
-            self._set_store("_xyz", None)
-            self._set_store("xyz", None)
-            return
-        self._set_store("_xyz", value)
-        # Forward activation
-        activated = value * self.aabb[None, 3:] + self.aabb[None, :3]
-        self._set_store("xyz", activated)
-
-    # --- Features (Combined) ---
-    @property
-    def features(self):
-        return self._get_store("features")
-
-    @features.setter
-    def features(self, value):
-        self.set_features(value)
-
-    def set_features(self, value):
-        self._set_store("features", value)
-        if value is None:
-            self._set_store("_features_dc", None)
-            self._set_store("features_dc", None)
-            self._set_store("_features_rest", None)
-            self._set_store("features_rest", None)
-            return
-
-        # Split features into DC and Rest
-        # Assuming shape (N, 1+K, 3)
-        dc = value[:, :1, :]
-        self._set_store("_features_dc", dc)
-        self._set_store("features_dc", dc)
-        
-        if value.shape[1] > 1:
-            rest = value[:, 1:, :]
-            self._set_store("_features_rest", rest)
-            self._set_store("features_rest", rest)
-        else:
-            self._set_store("_features_rest", None)
-            self._set_store("features_rest", None)
-
-    def _update_features_from_components(self):
-        dc = self._get_store("_features_dc")
-        rest = self._get_store("_features_rest")
-        if dc is None:
-            self._set_store("features", None)
-            return
-            
-        if rest is not None:
-            # Concatenate on dim 1 (N, 1+K, 3)
-            f = torch.cat((dc, rest), dim=1)
-            self._set_store("features", f)
-        else:
-            self._set_store("features", dc)
-
-    # --- Features DC ---
-    @property
-    def features_dc(self):
-        return self._get_store("features_dc")
-
-    @features_dc.setter
-    def features_dc(self, value):
-        self.set_features_dc(value)
-
-    @property
-    def _features_dc(self):
-        return self._get_store("_features_dc")
-
-    @_features_dc.setter
-    def _features_dc(self, value):
-        self.set_raw_features_dc(value)
-
-    def set_features_dc(self, value):
-        self._set_store("features_dc", value)
-        self._set_store("_features_dc", value) # Identity
-        self._update_features_from_components()
-
-    def set_raw_features_dc(self, value):
-        self._set_store("_features_dc", value)
-        self._set_store("features_dc", value) # Identity
-        self._update_features_from_components()
-
-    # --- Features Rest ---
-    @property
-    def features_rest(self):
-        return self._get_store("features_rest")
-
-    @features_rest.setter
-    def features_rest(self, value):
-        self.set_features_rest(value)
-
-    @property
-    def _features_rest(self):
-        return self._get_store("_features_rest")
-
-    @_features_rest.setter
-    def _features_rest(self, value):
-        self.set_raw_features_rest(value)
-
-    def set_features_rest(self, value):
-        self._set_store("features_rest", value)
-        self._set_store("_features_rest", value)
-        self._update_features_from_components()
+class Gaussian:
+    """Built from the Gaussian decoder's raw outputs; the activated values are what renderers and PLY files use.
+
+    xyz [N, 3] positions, features_dc [N, 1, 3] colour (degree-0 spherical harmonics), scaling [N, 3],
+    rotation [N, 4] unit quaternions (w, x, y, z), opacity [N, 1]. The object is z-up, inside [-0.5, 0.5]^3."""
+
+    def __init__(self, raw_xyz, features_dc, raw_scaling, raw_rotation, raw_opacity,
+                 scaling_bias: float, opacity_bias: float, minimum_kernel_size: float):
+        device = raw_xyz.device
+        aabb = torch.tensor([-0.5, -0.5, -0.5, 1.0, 1.0, 1.0], dtype=torch.float32, device=device)
+        scaling_offset = inverse_softplus(torch.tensor(scaling_bias)).to(device)
+        opacity_offset = inverse_sigmoid(torch.tensor(opacity_bias)).to(device)
+        rotation_offset = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device)
+
+        self.xyz = raw_xyz * aabb[None, 3:] + aabb[None, :3]
+        self.features_dc = features_dc
+        scaling = F.softplus(raw_scaling + scaling_offset)
+        self.scaling = torch.sqrt(torch.square(scaling) + minimum_kernel_size ** 2)
+        self.rotation = F.normalize(raw_rotation + rotation_offset[None, :])
+        self.opacity = torch.sigmoid(raw_opacity + opacity_offset)
 
-    def set_raw_features_rest(self, value):
-        self._set_store("_features_rest", value)
-        self._set_store("features_rest", value)
-        self._update_features_from_components()
-    
-    # --- Old API ---
-    @property
-    def get_scaling(self):
-        return self._get_store("scaling")
-    
-    @property
-    def get_rotation(self):
-        return self._get_store("rotation")
-    
-    @property
-    def get_xyz(self):
-        return self._get_store("xyz")
-    
-    @property
-    def get_features(self):
-        return self.features
-    
-    @property
-    def get_opacity(self):
-        return self._get_store("opacity")
-    
-    # --- Saver and Loader ---
-    def construct_list_of_attributes(self):
-        l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
-        # All channels except the 3 DC
-        for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
-            l.append('f_dc_{}'.format(i))
-        l.append('opacity')
-        for i in range(self._scaling.shape[1]):
-            l.append('scale_{}'.format(i))
-        for i in range(self._rotation.shape[1]):
-            l.append('rot_{}'.format(i))
-        return l
-        
     def save_ply(self, path):
         """Write the standard 3D Gaussian Splatting PLY (z-up object turned to the y-up convention viewers use)."""
-        write_ply_attributes(path, self.construct_list_of_attributes(), self.ply_attributes())
-
-    def ply_attributes(self):
-        """The (N, attributes) float array save_ply writes, in construct_list_of_attributes() order."""
-        xyz = self.get_xyz.detach().cpu().numpy()
+        xyz = self.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()
+        f_dc = self.features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
         # Clamp before the inverse activations: a saturated opacity (exactly 1.0 after the sigmoid) or a zero
         # scale would otherwise be written as +/-inf, which some PLY viewers reject.
         eps = 1e-6
-        opacities = inverse_sigmoid(self.get_opacity.clamp(eps, 1 - eps)).detach().cpu().numpy()
-        scale = torch.log(self.get_scaling.clamp_min(eps)).detach().cpu().numpy()
-        rotation = (self._rotation + self.rots_bias[None, :]).detach().cpu().numpy()
+        opacities = inverse_sigmoid(self.opacity.clamp(eps, 1 - eps)).detach().cpu().numpy()
+        scale = torch.log(self.scaling.clamp_min(eps)).detach().cpu().numpy()
+        rotation = self.rotation.detach().cpu().numpy()
 
         # z-up -> y-up: rotate +90 degrees about x, i.e. (x, y, z) -> (x, -z, y)
         xyz = np.stack([xyz[:, 0], -xyz[:, 2], xyz[:, 1]], axis=1)
-        rotation = rotation / np.linalg.norm(rotation, axis=1, keepdims=True)
         rotation = quaternion_multiply(np.array([np.sqrt(0.5), np.sqrt(0.5), 0.0, 0.0]), rotation)
 
-        return np.concatenate((xyz, normals, f_dc, opacities, scale, rotation), axis=1).astype(np.float32)
+        attributes = np.concatenate((xyz, normals, f_dc, opacities, scale, rotation), axis=1).astype(np.float32)
+        names = ['x', 'y', 'z', 'nx', 'ny', 'nz', 'f_dc_0', 'f_dc_1', 'f_dc_2', 'opacity',
+                 'scale_0', 'scale_1', 'scale_2', 'rot_0', 'rot_1', 'rot_2', 'rot_3']
+        elements = np.empty(attributes.shape[0], dtype=[(name, 'f4') for name in names])
+        for column, name in enumerate(names):
+            elements[name] = attributes[:, column]
+        PlyData([PlyElement.describe(elements, 'vertex')]).write(path)
 
 
 def quaternion_multiply(a, b):
@@ -390,14 +69,3 @@ def quaternion_multiply(a, b):
         aw * by - ax * bz + ay * bw + az * bx,
         aw * bz + ax * by - ay * bx + az * bw,
     ], axis=1)
-
-
-def write_ply_attributes(path, attribute_names, attributes):
-    """Write a Gaussian PLY from attribute names and an (N, len(names)) array (see Gaussian.ply_attributes)."""
-    dtype_full = [(attribute, 'f4') for attribute in attribute_names]
-    elements = np.empty(attributes.shape[0], dtype=dtype_full)
-    # Assign whole columns: constructing 262,144 Python tuples per PLY
-    # dominated validation export time and allocated millions of scalars.
-    for column_index, (attribute_name, _) in enumerate(dtype_full):
-        elements[attribute_name] = attributes[:, column_index]
-    PlyData([PlyElement.describe(elements, 'vertex')]).write(path)

idol/models/birefnet.py

@@ -4,7 +4,6 @@ image encoders see them."""
 import math
 
 import numpy as np
-import safetensors.torch
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -14,18 +13,6 @@ from torchvision.ops import deform_conv2d
 # 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):
@@ -77,7 +64,6 @@ class _WindowAttention(nn.Module):
 
         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
@@ -153,8 +139,6 @@ class _PatchMerging(nn.Module):
     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, :]
@@ -215,11 +199,6 @@ class _SwinPatchEmbed(nn.Module):
         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)
@@ -414,12 +393,6 @@ class _BiRefNetDecoder(nn.Module):
         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
@@ -431,10 +404,6 @@ class _BiRefNetDecoder(nn.Module):
 
         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()
@@ -527,21 +496,9 @@ class BiRefNet(nn.Module):
         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)

idol/models/decoder.py

@@ -1,424 +1,192 @@
-from typing import *
+"""The two decoders that turn a latent [B, 8192, 16] into Gaussians.
+
+OctreeProbabilityFixedlenDecoder places the anchor points: starting from the unit cube, it predicts how the points
+split among the 8 children of every occupied cell, level by level, then scatters each point inside its leaf cell.
+GaussianFixedlenDecoder turns every anchor point into 32 Gaussians (offsets, colour, scale, rotation, opacity)."""
+from typing import List
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from ..sampling_utils import hammersley_sequence, sample_probs
+
 from ..gaussian import Gaussian
-from .transformer_base import TransformerBase, ModulatedCrossOnlyTransformerBase, PcdAbsolutePositionEmbedder, PcdAbsolutePositionEmbedderv2
-import tensordict
-from tensordict import TensorDict
+from .layers import (ModulatedTransformerCrossOnlyBlock, PcdAbsolutePositionEmbedderv2, SinusoidalEmbedder,
+                     TransformerCrossBlock, convert_module_to_f16)
 
-class FixedlenDecoder(TransformerBase):
-    def __init__(
-        self,
-        in_channels: int,
-        model_channels: int,
-        cond_channels: int,
-        num_blocks: int,
-        num_heads: Optional[int] = None,
-        num_head_channels: Optional[int] = 64,
-        mlp_ratio: float = 4.0,
-        attn_mode: Literal["full"] = "full",
-        window_num: Optional[int] = None,
-        pcd_pe_mode: Literal["pcd_ape", "pcd_ape_v2"] = "pcd_ape_v2",
-        use_fp16: bool = False,
-        use_checkpoint: bool = False,
-        qk_rms_norm: bool = True,
-        qk_rms_norm_cross: bool = True,
-        share_mod: bool = False,
-    ):
-        """
-        Gaussian Feature Fixed Length Latent Decoder.
-        Decode fixed length latent to 3D Gaussian representation, encode condition with cross attention.
-        Conditions:
-            - latent: [B, Lq, C]
-        Input:
-            - x: [B, L, 3] (points to decode)
-        """
-        if attn_mode != "full":
-            raise ValueError(f"Only attn_mode='full' is supported in the public decoder, got {attn_mode!r}.")
-        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,
-            use_fp16=use_fp16,
-            use_checkpoint=use_checkpoint,
-            qk_rms_norm=qk_rms_norm,
-            qk_rms_norm_cross=qk_rms_norm_cross,
-            share_mod=share_mod,
-        )
-        self.in_channels = in_channels
-        self.in_proj = nn.Linear(in_channels, model_channels)
-        self.pcd_pe_mode = pcd_pe_mode
-        if pcd_pe_mode == "pcd_ape":
-            self.pos_embedder = PcdAbsolutePositionEmbedder(
-                channels=model_channels,
-                in_channels=3,
-            )
-        elif pcd_pe_mode == "pcd_ape_v2":
-            self.pos_embedder = PcdAbsolutePositionEmbedderv2(
-                channels=model_channels,
-                in_channels=3,
-            )
-        else:
-            self.pos_embedder = None
-            
-    def forward(self, x: tensordict.TensorDict = None, cond: Optional[torch.Tensor]= None) -> torch.Tensor:
-        """
-        Args:
-            x: a dict containing:
-                - points: [B, L, 3] the query points to decode
-            cond: [B, Lq, C] the latent features as condition
-        Returns:
-            list of Gaussian representations
-        """
-        pcd = x["points"]
-        B, L, C = pcd.shape
-        h = pcd
-        
-        # token num log scaling
-        l = torch.full((B,), L).to(h)
-        l = torch.log2(l).to(h)
-        
-        h = self.in_proj(h)
-        h = h + self.pos_embedder(pcd.view(-1,C)).view(B, L, -1) if hasattr(self, 'pos_embedder') else h
-        indices = None
-        # prepare spatial cache
 
-        h = super().forward(h, cond, l=l, indices=indices)
-        h = h.type(cond.dtype)
-        return h
+class OctreeProbabilityFixedlenDecoder(nn.Module):
+    def __init__(self, model_channels: int, cond_channels: int, num_blocks: int, num_heads: int,
+                 mlp_ratio: float = 4.0):
+        super().__init__()
+        self.dtype = torch.float16
+        self.in_proj = nn.Linear(3, model_channels)
+        self.pos_embedder = PcdAbsolutePositionEmbedderv2(channels=model_channels, in_channels=3)
+        self.input_layer = nn.Linear(model_channels, model_channels)
+        self.l_embedder = SinusoidalEmbedder(model_channels, max_period=1024, two_pi=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(model_channels, 6 * model_channels, bias=True))
+        self.blocks = nn.ModuleList([
+            ModulatedTransformerCrossOnlyBlock(model_channels, ctx_channels=cond_channels, num_heads=num_heads,
+                                               mlp_ratio=mlp_ratio)
+            for _ in range(num_blocks)
+        ])
+        self.out_proj = nn.Linear(model_channels, 2 ** 3)
+        self.blocks.apply(convert_module_to_f16)
 
-class GaussianFixedlenDecoder(FixedlenDecoder):
-    def __init__(
-        self,
-        in_channels: int,
-        model_channels: int,
-        cond_channels: int,
-        num_blocks: int,
-        num_heads: Optional[int] = None,
-        num_head_channels: Optional[int] = 64,
-        mlp_ratio: float = 4.0,
-        attn_mode: Literal["full"] = "full",
-        window_num: Optional[int] = None,
-        pcd_pe_mode: Literal["pcd_ape", "pcd_ape_v2"] = "pcd_ape_v2",
-        use_fp16: bool = False,
-        use_checkpoint: bool = False,
-        qk_rms_norm: bool = True,
-        qk_rms_norm_cross: bool = True,
-        share_mod: bool = False,
-        *,
-        representation_config: dict = None,
-        use_learned_offset_scale: bool = False,
-        gaussian_center_pertube_scale: float = 1.0,
-        use_per_offset=False,
-    ):
-        self.rep_config = representation_config
-        self.gaussian_center_pertube_scale = gaussian_center_pertube_scale
-        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,
-            pcd_pe_mode = pcd_pe_mode,
-            use_fp16 = use_fp16,
-            use_checkpoint = use_checkpoint,
-            qk_rms_norm = qk_rms_norm,
-            qk_rms_norm_cross = qk_rms_norm_cross,
-            share_mod = share_mod,
-        )
-        self.out_proj = nn.Linear(model_channels, self.out_channels)
-        self._build_perturbation()
-        
-        # weight init
-        self.initialize_weights()
-        if use_fp16:
-            self.convert_to_fp16()
+    def forward(self, x: torch.Tensor, level: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        """x: cell centres [B, L, 3] in [0, 1]; level: the cells' octree resolution [B]; cond: latent [B, Lc, C].
+        Returns the logits [B, L, 8] of how each cell's points split among its 8 children."""
+        B, L, _ = x.shape
+        h = self.in_proj(x) + self.pos_embedder(x.view(-1, 3)).view(B, L, -1)
+        h = self.input_layer(h)
+        level_emb = self.adaLN_modulation(self.l_embedder(level.type(h.dtype))).type(self.dtype)
+        out = h.type(self.dtype)
+        context = cond.type(self.dtype)
+        for block in self.blocks:
+            out = block(out, level_emb, context)
+        out = out.type(h.dtype).type(cond.dtype)
+        return self.out_proj(F.layer_norm(out, out.shape[-1:]))
 
-    def initialize_weights(self) -> None:
-        super().initialize_weights()
-        # zero-out output layers:
-        nn.init.constant_(self.out_proj.weight, 0)
-        nn.init.constant_(self.out_proj.bias, 0)
-        
-    def _calc_layout(self) -> None:
-        self.layout = {
-            '_xyz' : {'shape': (self.rep_config['num_gaussians'], 3), 'size': self.rep_config['num_gaussians'] * 3},
-            '_features_dc' : {'shape': (self.rep_config['num_gaussians'], 1, 3), 'size': self.rep_config['num_gaussians'] * 3},
-            '_scaling' : {'shape': (self.rep_config['num_gaussians'], 3), 'size': self.rep_config['num_gaussians'] * 3},
-            '_rotation' : {'shape': (self.rep_config['num_gaussians'], 4), 'size': self.rep_config['num_gaussians'] * 4},
-            '_opacity' : {'shape': (self.rep_config['num_gaussians'], 1), 'size': self.rep_config['num_gaussians']},
-        }
-        if self.use_learned_offset_scale:
-            if self.use_per_offset:
-                self.layout['_offset_scale'] = {'shape': (self.rep_config['num_gaussians'], 1), 'size': self.rep_config['num_gaussians']}
-            else:
-                raise ValueError("use_per_offset must be True when use_learned_offset_scale is True")
-        start = 0
-        for k, v in self.layout.items():
-            v['range'] = (start, start + v['size'])
-            start += v['size']
-        out_channels = start
-        return out_channels
-    
-    
-    def _build_perturbation(self) -> None:
-        perturbation = [hammersley_sequence(3, i, self.rep_config['num_gaussians']) for i in range(self.rep_config['num_gaussians'])]
-        perturbation = torch.tensor(perturbation).float() * 2 - 1
-        perturbation = perturbation / self.rep_config['perturbe_size']
-        perturbation = torch.atanh(perturbation).to(self.device)
-        self.register_buffer('points_offset_perturbation', perturbation)
-        
-        if self.use_learned_offset_scale:
-             base_offset_scale = torch.tensor(self.rep_config['offset_scale'])
-             base_offset_scale = torch.log(torch.exp(base_offset_scale) - 1.0) # inverse softplus
-             self.register_buffer('base_offset_scale', base_offset_scale)
-    
-    def _get_offset(self, h):
-        B = h.shape[0]
-        if self.use_learned_offset_scale:
-            _offset_scale = h[:,:, self.layout['_offset_scale']['range'][0]:self.layout['_offset_scale']['range'][1]].reshape(B, -1, *self.layout['_offset_scale']['shape'])
-            _offset_scale = F.softplus(_offset_scale+self.base_offset_scale)
-            
-        offset = h[:,:, self.layout['_xyz']['range'][0]:self.layout['_xyz']['range'][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']
-        if self.use_learned_offset_scale:
-            offset = offset * _offset_scale
-        else:
-            offset = offset * self.rep_config['offset_scale']
-        return offset
-        
+    @staticmethod
+    @torch.no_grad()
+    def sample(model, cond: torch.Tensor, num_points: int, level: int) -> torch.Tensor:
+        """Place num_points anchor points [B, num_points, 3] in [0, 1]^3, descending the octree to ``level``."""
+        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)
+        batch_indices_range = torch.arange(B, device=device).unsqueeze(1)
 
-    def to_representation(self, x: torch.Tensor, h: torch.Tensor) -> List[Gaussian]:
-        """
-        Convert a batch of network outputs to 3D representations.
+        for l in range(1, level + 1):
+            res_p = 1 << (l - 1)
+            parent_coords_norm = (prev_coords_int.to(torch.float32) + 0.5) / res_p
+            res_tensor = torch.full((B,), 1 << l, dtype=torch.long, device=device)
+            pred_probs = torch.softmax(model(parent_coords_norm, res_tensor, cond), dim=-1)    # [B, Np, 8]
+            sampled = systematic_sample(pred_probs, prev_counts).flatten(1, 2)                   # [B, Np*8]
+            child_coords_int = (prev_coords_int[:, :, None, :] * 2 + child_offset[None, None, :, :]).flatten(1, 2)
 
-        Args:
-            x: the input tensor of shape [B, L, C].
-            h: the output tensor of shape [B, L, C'] from the MLP.
+            # keep the children that received points, packed to the front
+            mask = sampled > 0
+            max_valid = mask.sum(dim=1).max().item()
+            valid_scatter_indices = (mask.cumsum(dim=1) - 1)[mask]
+            valid_batch_indices = batch_indices_range.expand_as(mask)[mask]
+            prev_coords_int = torch.zeros(B, max_valid, 3, dtype=child_coords_int.dtype, device=device)
+            prev_coords_int[valid_batch_indices, valid_scatter_indices] = child_coords_int[mask]
+            prev_counts = torch.zeros(B, max_valid, dtype=sampled.dtype, device=device)
+            prev_counts[valid_batch_indices, valid_scatter_indices] = sampled[mask]
 
-        Returns:
-            list of representations
-        """
-        x = h["points"] if "points" in h else x["points"]
-        offset = self._get_offset(h['features'])
-        h = h["features"]
-        ret = []
-        for i in range(h.shape[0]):
-            representation = Gaussian(
-                sh_degree=0,
-                aabb=[-0.5, -0.5, -0.5, 1.0, 1.0, 1.0],
-                mininum_kernel_size = self.rep_config['3d_filter_kernel_size'],
-                scaling_bias = self.rep_config['scaling_bias'],
-                opacity_bias = self.rep_config['opacity_bias'],
-                scaling_activation = self.rep_config['scaling_activation'],
-                device = h.device,
-            )
-            _x = x[i,:,None,:]
+        # each point lands at a uniformly random position inside its leaf cell
+        coords_int = torch.repeat_interleave(prev_coords_int.flatten(0, 1), prev_counts.flatten(0, 1),
+                                             dim=0).reshape(B, num_points, -1)
+        return (coords_int.to(torch.float32) + torch.rand_like(coords_int, dtype=torch.float32)) / (1 << level)
 
-            for k, v in self.layout.items():
-                if k == '_xyz':
-                    _xyz = offset[i] + _x # NOTE: based on x data range == [0, 1]
-                    setattr(representation, k, _xyz.flatten(0, 1))
-                elif k == '_xyz_center':
-                    continue
-                elif k == '_offset_scale':
-                    continue
-                else:
-                    feats = h[i][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape']).flatten(0, 1)
-                    feats = feats * self.rep_config['lr'][k]
-                    setattr(representation, k, feats)
-            ret.append(representation)
-        return ret
-    
-    def forward(self, x = None, cond = None) -> TensorDict:
-        h = super().forward(x, cond)
-        
-        h = F.layer_norm(h, h.shape[-1:])
-        h_out = self.out_proj(h)
-        return TensorDict({
-            'features': h_out
-        }, batch_size=h_out.shape[:1], device=h.device)
 
-class OctreeProbabilityFixedlenDecoder(ModulatedCrossOnlyTransformerBase):
-    def __init__(
-        self,
-        model_channels: int,
-        cond_channels: Optional[int],
-        num_blocks: int,
-        num_heads: Optional[int] = None,
-        num_head_channels: Optional[int] = 64,
-        mlp_ratio: float = 4.0,
-        use_fp16: bool = False,
-        use_checkpoint: bool = False,
-        share_mod: bool = False,
-        qk_rms_norm_cross: bool = True,
-        *,
-        pcd_pe_mode: Literal["pcd_ape", "pcd_ape_v2"] = "pcd_ape_v2",
-        **kwargs,
-    ):
-        super().__init__(
-            in_channels = model_channels, # pe embedded
-            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,
-            use_fp16 = use_fp16,
-            use_checkpoint = use_checkpoint,
-            share_mod = share_mod,
-            qk_rms_norm_cross = qk_rms_norm_cross,
-        )
-        self.out_proj = nn.Linear(self.model_channels, 2 ** 3) # to 8 logits
-        
-        self.in_proj = nn.Linear(3, self.model_channels)
-        self.pcd_pe_mode = pcd_pe_mode
-        if pcd_pe_mode == "pcd_ape":
-            self.pos_embedder = PcdAbsolutePositionEmbedder(
-                channels=model_channels,
-                in_channels=3,
-            )
-        elif pcd_pe_mode == "pcd_ape_v2":
-            self.pos_embedder = PcdAbsolutePositionEmbedderv2(
-                channels=model_channels,
-                in_channels=3,
-            )
-        else:
-            self.pos_embedder = None
-            
-        # weight init
-        self.initialize_weights()
-        if use_fp16:
-            self.convert_to_fp16()
-            
-    def initialize_weights(self) -> None:
-        super().initialize_weights()
-        # zero-out output layers:
-        nn.init.constant_(self.out_proj.weight, 0)
-        nn.init.constant_(self.out_proj.bias, 0)
-        
-    def forward(self, x: torch.Tensor, l: torch.Tensor, cond: Optional[torch.Tensor]= None, l2: Optional[torch.Tensor] = None) -> torch.Tensor:
-        """
-        Args:
-            x: query points voxel centers, [B(=1), L, 3]
-            l: query points level, [B(=1)]
-            cond: [B(=1), Lq, C] the latent features as condition
-        Returns:
-            list of Gaussian representations
-        """
-        B, L, C = x.shape
-        h = self.in_proj(x)
-        h = h + self.pos_embedder(x.view(-1,3)).view(B, L, -1) if hasattr(self, 'pos_embedder') else h
-        # FIXME: elegant l2 log scaling
-        if l2 is not None:
-            l2 = torch.log2(l2)
-        h = super().forward(h, l, cond, l2)
-        h = h.type(cond.dtype)
-        h = F.layer_norm(h, h.shape[-1:])
-        logits = self.out_proj(h)
-        probs = torch.softmax(logits, dim=-1)
-        return TensorDict({'logits': logits, 'probs': probs}, batch_size=logits.shape[:1], device=logits.device)
+@torch.no_grad()
+def systematic_sample(probs: torch.Tensor, counts: torch.Tensor) -> torch.Tensor:
+    """Split counts[...] items among the P bins of probs[..., P] by systematic resampling -> [..., P] counts."""
+    batch_shape = counts.shape
+    B = counts.numel()
+    P = probs.size(-1)
+    device = probs.device
+    probs = probs.view(B, P)
+    counts = counts.view(B)
 
-    @staticmethod
-    def sample(model, cond, num_points, level, temperature=1.0, algo: Literal["iid", "residual", "systematic"] = "systematic") -> TensorDict:
-        """
-        Sample points from the model.
-        Args:
-            cond: object latent condition, [B, Lc, C]
-            num_points: number of points to sample
-            level: level of the octree
-            temperature: temperature of the sampling
-        Returns:
-            points: sampled points, [B, num_points, 3]
-            log_probs: log probabilities of the sampled points, [B, num_points]
-        """
-        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) # [B, 1, 3]
-        prev_counts = torch.full((B, 1), num_points, dtype=torch.long, device=device) # [B, 1]
-        prev_log_probs = torch.zeros(B, 1, dtype=torch.float32, device=device) # [B, 1]
-        
-        batch_indices_range = torch.arange(B, device=device).unsqueeze(1)
-        num_tensor = torch.full((B,), num_points, dtype=torch.long, device=device)
+    # normalise, with a uniform fallback for zero-sum rows
+    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
 
-        for l in range(1, level + 1):
-            res_p = 1 << (l-1)
-            res = 1 << l
-                
-            parent_coords_int = prev_coords_int # [B, Np, 3]
-            parent_coords_norm = (parent_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'] # [B, Np, 8]
-            pred_logits = pred_logits / temperature
-            pred_probs = torch.softmax(pred_logits, dim=-1) # [B, Np, 8]
-            pred_log_probs = torch.log_softmax(pred_logits, dim=-1) # [B, Np, 8]
-            sampled = sample_probs(pred_probs, prev_counts, algo=algo,
-                                   group_rows=getattr(model, 'group_sampling_rows', False)).flatten(1,2)
-            pred_log_probs = pred_log_probs.flatten(1,2) # [B, Np*8]
-            prev_log_probs_expanded = prev_log_probs.repeat_interleave(8, dim=1) # [B, Np*8]
+    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 index, n in enumerate(unique_n.tolist()):
+        if n == 0:
+            continue
+        rows = (inv == index).nonzero(as_tuple=False).squeeze(1)
+        r = rows.numel()
+        u0 = torch.rand(r, 1, device=device) / float(n)                                   # random offset in [0, 1/n)
+        grid = torch.arange(n, device=device, dtype=torch.float32)[None, :] / float(n)    # stratified points
+        us = (u0 + grid).clamp(max=1.0 - 1e-12)
+        idx = torch.searchsorted(cdf.index_select(0, rows), us).clamp_max(P - 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)
+
+
+class GaussianFixedlenDecoder(nn.Module):
+    def __init__(self, in_channels: int, model_channels: int, cond_channels: int, num_blocks: int, num_heads: int,
+                 representation_config: dict, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.dtype = torch.float16
+        self.rep_config = representation_config
+        per_anchor = representation_config['num_gaussians']
+        # output channels per anchor: offset (3), colour (3), scale (3), rotation (4), opacity (1), offset scale (1)
+        self.layout = {
+            '_xyz': (3,),
+            '_features_dc': (1, 3),
+            '_scaling': (3,),
+            '_rotation': (4,),
+            '_opacity': (1,),
+            '_offset_scale': (1,),
+        }
+        start = 0
+        self.ranges = {}
+        for name, shape in self.layout.items():
+            size = per_anchor * int(torch.tensor(shape).prod())
+            self.ranges[name] = (start, start + size)
+            start += size
+
+        self.in_proj = nn.Linear(in_channels, model_channels)
+        self.pos_embedder = PcdAbsolutePositionEmbedderv2(channels=model_channels, in_channels=3)
+        self.input_layer = nn.Linear(model_channels, model_channels)
+        self.blocks = nn.ModuleList([
+            TransformerCrossBlock(model_channels, ctx_channels=cond_channels, num_heads=num_heads, mlp_ratio=mlp_ratio)
+            for _ in range(num_blocks)
+        ])
+        self.out_proj = nn.Linear(model_channels, start)
+        # fixed per-Gaussian offset pattern and the base offset scale, stored with the weights
+        self.register_buffer('points_offset_perturbation', torch.zeros(per_anchor, 3))
+        self.register_buffer('base_offset_scale', torch.zeros(()))
+        self.blocks.apply(convert_module_to_f16)
+
+    def forward(self, points: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        """points: anchors [B, L, 3]; cond: latent [B, Lc, C] -> per-anchor outputs [B, L, 32 x 15]"""
+        B, L, C = points.shape
+        h = self.in_proj(points) + self.pos_embedder(points.view(-1, C)).view(B, L, -1)
+        h = self.input_layer(h)
+        out = h.type(self.dtype)
+        context = cond.type(self.dtype)
+        for block in self.blocks:
+            out = block(out, context)
+        out = out.type(h.dtype).type(cond.dtype)
+        return self.out_proj(F.layer_norm(out, out.shape[-1:]))
 
-            # prepare coords
-            child_coords_int = parent_coords_int[:, :, None, :] * 2 + child_offset[None, None, :, :] # [B,. Np, 8, 3]
-            child_coords_int = child_coords_int.flatten(1,2) # [B, Np*8, 3]
-                
-            # mask out and padding
-            # Vectorized implementation
-            mask = sampled > 0 # [B, K]
-            valid_counts = mask.sum(dim=1) # [B]
-            max_valid = valid_counts.max().item()
-            
-            # Create scatter indices
-            scatter_indices = mask.cumsum(dim=1) - 1
-            valid_mask = mask # just mask
-            
-            valid_scatter_indices = scatter_indices[valid_mask] # [Total_Valid]
-            
-            # Batch indices
-            valid_batch_indices = batch_indices_range.expand_as(mask)[valid_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[valid_mask]
-            
-            next_prev_counts = torch.zeros(B, max_valid, dtype=sampled.dtype, device=device)
-            next_prev_counts[valid_batch_indices, valid_scatter_indices] = sampled[valid_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)[valid_mask]
-            
-            prev_coords_int = next_prev_coords_int
-            prev_counts = next_prev_counts
-            prev_log_probs = next_prev_log_probs
+    def _slice(self, features, name):
+        start, end = self.ranges[name]
+        return features[:, :, start:end].reshape(features.shape[0], -1, self.rep_config['num_gaussians'],
+                                                 *self.layout[name])
 
-            
-        # finially randomly put points in each voxel
-        res = 1 << level
-        # unflatten the coords according to counts
-        prev_log_probs = torch.repeat_interleave(prev_log_probs.flatten(0,1), prev_counts.flatten(0,1), dim=0).reshape(B, num_points,) # [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) # [B, num_points, 3]
-        coords_norm = (coords_int.to(torch.float32) + torch.rand_like(coords_int, dtype=torch.float32)) / res
+    def to_representation(self, points: torch.Tensor, features: torch.Tensor) -> List[Gaussian]:
+        """Anchors [B, L, 3] and per-anchor outputs [B, L, D] -> one Gaussian object per batch element."""
+        config = self.rep_config
+        offset_scale = F.softplus(self._slice(features, '_offset_scale') + self.base_offset_scale)
+        offset = self._slice(features, '_xyz') * config['lr']['_xyz'] + self.points_offset_perturbation
+        offset = torch.tanh(offset) * 0.5 * config['perturbe_size'] * offset_scale
 
-        points = coords_norm
-        log_probs = prev_log_probs
-        return TensorDict({
-            'points': points,
-            'log_probs': log_probs,
-        }, batch_size=points.shape[:1], device=points.device)
+        objects = []
+        for i in range(features.shape[0]):
+            raw = {name: (self._slice(features[i:i + 1], name)[0].flatten(0, 1) * config['lr'][name])
+                   for name in ('_features_dc', '_scaling', '_rotation', '_opacity')}
+            objects.append(Gaussian(
+                raw_xyz=(offset[i] + points[i, :, None, :]).flatten(0, 1),
+                features_dc=raw['_features_dc'], raw_scaling=raw['_scaling'], raw_rotation=raw['_rotation'],
+                raw_opacity=raw['_opacity'], scaling_bias=config['scaling_bias'],
+                opacity_bias=config['opacity_bias'], minimum_kernel_size=config['3d_filter_kernel_size']))
+        return objects

idol/models/dit.py

@@ -1,6 +1,6 @@
 # Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py
 # The original license is as follows:
-# 
+#
 # Copyright 2025 The Wan Team and The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,415 +14,109 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-from typing import *
+"""The denoiser: a rectified-flow transformer over the object latent (8192 tokens) and the input camera (1 token),
+conditioned on the image tokens. Joint attention over [latent, image, camera] with learned rotary positions."""
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import numpy as np
 from tensordict import TensorDict
 
-from ..modules.mlp import MLP
-from ..modules.transformer.modulated import UnifiedTransformerBlock
-from ..modules.attention.modules import RePo3DRotaryEmbedding, Rotary3DPoseEmbedding
-from ..modules.utils import convert_module_to_f16, convert_module_to_f32
-from .transformer_base import PcdAbsolutePositionEmbedder
-
-class TimestepEmbedder(nn.Module):
-    """
-    Embeds scalar timesteps into vector representations.
-    """
-    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):
-        """
-        Create sinusoidal timestep embeddings.
-
-        Args:
-            t: a 1-D Tensor of N indices, one per batch element.
-                These may be fractional.
-            dim: the dimension of the output.
-            max_period: controls the minimum frequency of the embeddings.
-
-        Returns:
-            an (N, D) Tensor of positional embeddings.
-        """
-        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
-        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
+from .layers import (MLP, PcdAbsolutePositionEmbedder, RePo3DRotaryEmbedding, SinusoidalEmbedder,
+                     UnifiedTransformerBlock, convert_module_to_f16)
 
-    def forward(self, t):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
-        t_emb = self.mlp(t_freq)
-        return t_emb
 
 class LatentSeqMMFlowModel(nn.Module):
-    """
-    Latent Sequence Multimodal Flow Model (Unified MMDiT).
-    Adapts ZImageTransformer2DModel architecture for 3D-Image unification.
-    MMDiT adapted from diffusers/models/transformers/transformer_z_image.py
-    REPO adapted from https://github.com/SakanaAI/repo
-    """
     def __init__(
         self,
         q_token_length: int,
         in_channels: int,
+        cam_channels: int,
         model_channels: int,
         cond_channels: int,
-        out_channels: int,
+        cond2_channels: int,
         num_blocks: int,
         num_refiner_blocks: int = 2,
-        num_heads: Optional[int] = None,
-        num_head_channels: Optional[int] = 64,
-        cam_channels: Optional[int] = None,
-        cond2_channels: Optional[int] = None,
+        num_heads: int = 16,
+        num_head_channels: int = 64,
         mlp_ratio: float = 4,
-        pe_mode: Literal["3d_ot_rope", "3d_ot_repo", "3d_ot_ape", "repo", "learnable", None] = "3d_ot_ape",
-        use_fp16: bool = False,
-        use_checkpoint: bool = False,
-        share_mod: bool = True,
-        qk_rms_norm: bool = False,
-        use_shift_table: bool = 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.pe_mode = pe_mode
-        self.use_fp16 = use_fp16
-        self.use_checkpoint = use_checkpoint
-        self.share_mod = share_mod
-        self.qk_rms_norm = qk_rms_norm
-        self.dtype = torch.float16 if use_fp16 else torch.float32
-        self.use_shift_table = use_shift_table
+        self.dtype = torch.float16
+
+        def blocks(count, modulation=True):
+            return nn.ModuleList([UnifiedTransformerBlock(model_channels, num_heads=num_heads, mlp_ratio=mlp_ratio,
+                                                          modulation=modulation) for _ in range(count)])
 
-        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)
-            )
+        def rope_layers(count):
+            return nn.ModuleList([RePo3DRotaryEmbedding(model_channels, num_heads=num_heads,
+                                                        head_dim=num_head_channels) for _ in range(count)])
 
+        self.t_embedder = SinusoidalEmbedder(model_channels)
+        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)
-        if cond2_channels is not None:
-            self.cond_embedder2 = nn.Linear(cond2_channels, model_channels)
-        else:
-            self.cond_embedder2 = None
-        
-        # Positional Embeddings
-        self.use_rope = (self.pe_mode in ["3d_ot_rope", "3d_ot_repo", "repo"])
-        if pe_mode == "3d_ot_rope":
-            sobol_seq = torch.quasirandom.SobolEngine(dimension=3, scramble=True, seed=123).draw(q_token_length)
-            self.register_buffer("pos_pe", sobol_seq[None])
-            self.pos_embedder = Rotary3DPoseEmbedding(num_head_channels)
-            self.cond_pos_embedder = None
-        elif pe_mode == "3d_ot_ape":
-            sobol_seq = torch.quasirandom.SobolEngine(dimension=3, scramble=True, seed=123).draw(q_token_length)
-            self.register_buffer("pos_pe", sobol_seq[None])
-            self.pos_embedder = PcdAbsolutePositionEmbedder(model_channels)
-            self.cond_pos_embedder = None
-        elif pe_mode == "repo":
-            self.pos_embedder = None
-            self.cond_pos_embedder = None
-            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)
-            ])
-        elif pe_mode == "3d_ot_repo":
-            sobol_seq = torch.quasirandom.SobolEngine(dimension=3, scramble=True, seed=123).draw(q_token_length)
-            self.register_buffer("pos_pe", sobol_seq[None])
-            self.pos_embedder = PcdAbsolutePositionEmbedder(model_channels)
-            self.cond_pos_embedder = None
-            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)
-            ])
-        elif pe_mode == "learnable":
-            # Not fully implemented adaptation for learnable
-            self.pos_pe = None
-            self.pos_embedder = None
-        else:
-            self.pos_pe = None
-            self.pos_embedder = None
+        self.cond_embedder2 = nn.Linear(cond2_channels, model_channels)
 
-        # Refiners
-        self.noise_refiner = nn.ModuleList([
-            UnifiedTransformerBlock(
-                model_channels,
-                num_heads=self.num_heads,
-                mlp_ratio=self.mlp_ratio,
-                use_checkpoint=self.use_checkpoint,
-                use_rope=self.use_rope,
-                qk_rms_norm=self.qk_rms_norm,
-                modulation=True,
-                share_mod=self.share_mod,
-                use_shift_table=self.use_shift_table,
-            )
-            for _ in range(num_refiner_blocks)
-        ])
-        
-        self.context_refiner = nn.ModuleList([
-            UnifiedTransformerBlock(
-                model_channels,
-                num_heads=self.num_heads,
-                mlp_ratio=self.mlp_ratio,
-                use_checkpoint=self.use_checkpoint,
-                use_rope=self.use_rope,
-                qk_rms_norm=self.qk_rms_norm,
-                modulation=False,
-            )
-            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,
-                num_heads=self.num_heads,
-                mlp_ratio=self.mlp_ratio,
-                use_checkpoint=self.use_checkpoint,
-                use_rope=self.use_rope,
-                qk_rms_norm=self.qk_rms_norm,
-                modulation=True,
-                share_mod=self.share_mod,
-                use_shift_table=self.use_shift_table,
-            )
-            for _ in range(num_blocks)
-        ])
+        # positions of the latent tokens (a fixed scrambled Sobol sequence, stored with the weights)
+        self.register_buffer("pos_pe", torch.zeros(1, q_token_length, 3))
+        self.pos_embedder = PcdAbsolutePositionEmbedder(model_channels)
+        self.noise_repo_layers = rope_layers(num_refiner_blocks)
+        self.context_repo_layers = rope_layers(num_refiner_blocks)
+        self.repo_layers = rope_layers(num_blocks)
 
-        if self.use_shift_table:
-            self.shift_table = nn.Parameter(torch.randn(1, 2, model_channels) / model_channels**0.5)
-        else:
-            self.shift_table = None
+        # each stream is refined on its own, then all tokens go through the joint blocks
+        self.noise_refiner = blocks(num_refiner_blocks)
+        self.context_refiner = blocks(num_refiner_blocks, modulation=False)
+        self.cam_refiner = MLP(cam_channels, model_channels, model_channels, mlp_layer_num=num_refiner_blocks)
+        self.blocks = blocks(num_blocks)
 
-        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)
-        
-        self.initialize_weights()
-        if use_fp16:
-            self.convert_to_fp16()
+        self.shift_table = nn.Parameter(torch.zeros(1, 2, model_channels))
+        self.out_layer = nn.Linear(model_channels, in_channels)
+        self.cam_out_layer = nn.Linear(model_channels, cam_channels)
 
-    @property
-    def device(self) -> torch.device:
-        return next(self.parameters()).device
-
-    def convert_to_fp16(self) -> None:
         self.blocks.apply(convert_module_to_f16)
         self.noise_refiner.apply(convert_module_to_f16)
         self.context_refiner.apply(convert_module_to_f16)
 
-    def convert_to_fp32(self) -> None:
-        self.blocks.apply(convert_module_to_f32)
-        self.noise_refiner.apply(convert_module_to_f32)
-        self.context_refiner.apply(convert_module_to_f32)
-
-    def initialize_weights(self) -> None:
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-        self.apply(_basic_init)
-
-        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
-
-        if self.share_mod:
-            nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
-            nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
-        else:
-            for block in self.blocks:
-                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
-                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
-            for block in self.noise_refiner:
-                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
-                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
-        
-        if self.use_shift_table:
-            for block in self.blocks:
-                nn.init.constant_(block.shift_table, 0)
-            for block in self.noise_refiner:
-                nn.init.constant_(block.shift_table, 0)
-            nn.init.constant_(self.shift_table, 0)
-            
-        if hasattr(self, 'repo_layers'):
-            for block in self.noise_repo_layers:
-                block.initialize_weights()
-            for block in self.context_repo_layers:
-                block.initialize_weights()
-            for block in self.repo_layers:
-                block.initialize_weights()
-                
-        if self.cond_embedder2 is not None:
-            nn.init.constant_(self.cond_embedder2.weight, 0)
-            nn.init.constant_(self.cond_embedder2.bias, 0)
-
-        nn.init.constant_(self.out_layer.weight, 0)
-        nn.init.constant_(self.out_layer.bias, 0)
-
-        if self.cam_channels is not None:
-            nn.init.constant_(self.cam_out_layer.weight, 0)
-            nn.init.constant_(self.cam_out_layer.bias, 0)
-
-    def forward(
-        self,
-        x: TensorDict,
-        t: torch.Tensor,
-        cond: TensorDict,
-        points: Optional[torch.Tensor] = None,
-        return_mid_features: bool = False,
-        **kwargs,
-    ) -> torch.Tensor:
+    def forward(self, x: TensorDict, t: torch.Tensor, cond: TensorDict) -> TensorDict:
         """
         Args:
-            x: TensorDict with 'latent' [B, L_x, C_in] key, optional 'cam' [B, 1, C_cam] key
-            cond: TensorDict with 'feature1' [B, L_c, C_cond1] key, optional 'feature2' [B, L_c, C_cond2] key
+            x: TensorDict with 'latent' [B, 8192, 16] and 'camera' [B, 1, 5]
+            t: [B] time in [0, 1000]
+            cond: TensorDict with 'feature1' [B, L_c, 1280] (DINOv3) and 'feature2' [B, L_c, 128] (FLUX.2 VAE)
+        Returns:
+            the velocity, a TensorDict shaped like x
         """
-        
-        z = x['latent']
+        z, cam = x['latent'], x['camera']
         h_x = self.input_layer(z).type(self.dtype)
-        if self.cond_embedder2 is not None:
-            h_cond = self.cond_embedder(cond['feature1']) + self.cond_embedder2(cond['feature2'])
-        else:
-            h_cond = self.cond_embedder(cond['feature1'])
+        h_cond = self.cond_embedder(cond['feature1']) + self.cond_embedder2(cond['feature2'])
         h_cond = h_cond.type(self.dtype)
-        
-        t_emb = self.t_embedder(t)
-        if self.share_mod:
-            t_mod = self.adaLN_modulation(t_emb)
-        else:
-            t_mod = t_emb
-        t_mod = t_mod.type(self.dtype)
 
-        rotary_emb = None
-        rope_x = None
-        rope_cond = None
+        t_emb = self.t_embedder(t)
+        t_mod = self.adaLN_modulation(t_emb).type(self.dtype)
 
-        if self.pos_embedder is not None:
-            if self.pe_mode == "3d_ot_rope":
-                rope_x = self.pos_embedder(self.pos_pe) # [1, L, 1, D/2]
-                rope_x = rope_x.expand(z.shape[0], -1, -1, -1)
-                
-                # cond RoPE (Identity)
-                rope_cond = torch.ones(z.shape[0], h_cond.shape[1], 1, rope_x.shape[-1], device=z.device, dtype=rope_x.real.dtype)
-                rope_cond = torch.polar(rope_cond, torch.zeros_like(rope_cond)) 
-                
-                parts = [rope_x, rope_cond]
-                if self.cam_channels is not None:
-                    rope_cam = torch.ones(z.shape[0], x['camera'].shape[1], 1, rope_x.shape[-1], device=z.device, dtype=rope_x.real.dtype)
-                    rope_cam = torch.polar(rope_cam, torch.zeros_like(rope_cam))
-                    parts.append(rope_cam)
-                
-                rotary_emb = torch.cat(parts, dim=1)
+        h_x = h_x + self.pos_embedder(self.pos_pe).type(self.dtype)
+        for repo, block in zip(self.noise_repo_layers, self.noise_refiner):
+            h_x = block(h_x, mod=t_mod, rotary_emb=repo(h_x))
+        for repo, block in zip(self.context_repo_layers, self.context_refiner):
+            h_cond = block(h_cond, mod=None, rotary_emb=repo(h_cond))
+        h_cam = self.cam_refiner(cam).type(self.dtype)
 
-            elif self.pe_mode in ["3d_ot_ape", "3d_ot_repo"]:
-                h_x = h_x + self.pos_embedder(self.pos_pe).type(self.dtype)
-        
-        # Refine noise (z)
-        for i, block in enumerate(self.noise_refiner):
-            if "repo" in self.pe_mode:
-                rope_x = self.noise_repo_layers[i](h_x)
-            h_x = block(h_x, mod=t_mod, rotary_emb=rope_x)
-            
-        # Refine context (cond)
-        # Note: We pass rope_cond (identity) if rope is used, though it might not affect much if it's identity
-        for i, block in enumerate(self.context_refiner):
-            if "repo" in self.pe_mode:
-                rope_cond = self.context_repo_layers[i](h_cond)
-            h_cond = block(h_cond, mod=None, rotary_emb=rope_cond)
-            
-        # Refine camera
-        if self.cam_channels is not None:
-            cam = x['camera']
-            h_cam = self.cam_refiner(cam)
-            h_cam = h_cam.type(self.dtype)
+        h = torch.cat([h_x, h_cond, h_cam], dim=1)
+        for repo, block in zip(self.repo_layers, self.blocks):
+            h = block(h, mod=t_mod, rotary_emb=repo(h))
 
-        # Concat
-        h = torch.cat([h_x, h_cond], dim=1)
-        if self.cam_channels is not None:
-            h = torch.cat([h, h_cam], dim=1)
-        
-        mid_features = None
-        mid_idx = len(self.blocks) // 2
-        for i, block in enumerate(self.blocks):
-            if "repo" in self.pe_mode:
-                rotary_emb = self.repo_layers[i](h)
-            h = block(h, mod=t_mod, rotary_emb=rotary_emb)
-            if i == mid_idx:
-                mid_features = h
-            
-        # Split
-        h_x = h[:, :z.shape[1]]
-        h_x = h_x.type(z.dtype)
+        h_x = h[:, :z.shape[1]].type(z.dtype)
         h_x = F.layer_norm(h_x, h_x.shape[-1:])
-        # camera out
-        if self.cam_channels is not None:
-            h_cam = h[:, -cam.shape[1]:]
-            h_cam = h_cam.type(cam.dtype)
-            h_cam = F.layer_norm(h_cam, h_cam.shape[-1:])
-            
-        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
-        
-        h_x = self.out_layer(h_x)
-        if self.cam_channels is not None:
-            h_cam = self.cam_out_layer(h_cam)
-        
-        out = TensorDict({
-            'latent': h_x,
-        }, batch_size=x.batch_size)
-        if self.cam_channels is not None:
-            out['camera'] = h_cam
-        if return_mid_features:
-            return out, mid_features
+        h_cam = h[:, -cam.shape[1]:].type(cam.dtype)
+        h_cam = F.layer_norm(h_cam, h_cam.shape[-1:])
+        shift, scale = (self.shift_table + t_emb.unsqueeze(1)).chunk(2, dim=1)
+        h_x = h_x * (1 + scale) + shift
+        h_cam = h_cam * (1 + scale) + shift
+
+        out = TensorDict({'latent': self.out_layer(h_x)}, batch_size=x.batch_size)
+        out['camera'] = self.cam_out_layer(h_cam)
         return out

idol/models/layers.py

@@ -0,0 +1,318 @@
+"""Building blocks shared by the denoiser and the decoders: attention, positional embeddings, transformer blocks."""
+import importlib.util
+import os
+from typing import Optional
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# flash-attn when it is installed, PyTorch SDPA otherwise (same result, flash-attn is faster).
+# Override with ATTN_BACKEND=flash_attn|sdpa.
+ATTN_BACKEND = os.environ.get('ATTN_BACKEND',
+                              'flash_attn' if importlib.util.find_spec('flash_attn') is not None else 'sdpa')
+if ATTN_BACKEND not in ('flash_attn', 'sdpa'):
+    raise ValueError(f'Unknown attention backend: {ATTN_BACKEND}')
+if ATTN_BACKEND == 'flash_attn':
+    import flash_attn
+
+
+def scaled_dot_product_attention(*args) -> torch.Tensor:
+    """(qkv [N, L, 3, H, C]) or (q [N, L, H, C], k [N, S, H, C], v [N, S, H, C]) -> [N, L, H, C]."""
+    if ATTN_BACKEND == 'flash_attn':
+        if len(args) == 1:
+            return flash_attn.flash_attn_qkvpacked_func(args[0])
+        return flash_attn.flash_attn_func(*args)
+    q, k, v = args[0].unbind(dim=2) if len(args) == 1 else args
+    out = F.scaled_dot_product_attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3))
+    return out.permute(0, 2, 1, 3)
+
+
+def convert_module_to_f16(module):
+    """Store a linear layer's weights in float16 (the transformer trunks run in half precision)."""
+    if isinstance(module, nn.Linear):
+        for parameter in module.parameters():
+            parameter.data = parameter.data.half()
+
+
+class LayerNorm32(nn.LayerNorm):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x.float()).type(x.dtype)
+
+
+class MultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (F.normalize(x.float(), dim=-1) * self.gamma * self.scale).to(x.dtype)
+
+
+class FeedForwardNet(nn.Module):
+    def __init__(self, channels: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, int(channels * mlp_ratio)),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(int(channels * mlp_ratio), channels),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+
+class MLP(nn.Module):
+    def __init__(self, channels: int, inner_channels: int, channels_out: int, 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_out))
+        self.mlp = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+
+# ---------------------------------------------------------------------------------------------------------------------
+# Positional and scalar embeddings
+# ---------------------------------------------------------------------------------------------------------------------
+
+class PcdAbsolutePositionEmbedder(nn.Module):
+    """Sinusoidal embedding of 3D positions in [0, 1] (the denoiser's latent token positions)."""
+    def __init__(self, channels: int, in_channels: int = 3, max_res=16):
+        super().__init__()
+        self.channels = channels
+        self.in_channels = in_channels
+        self.freq_dim = channels // in_channels // 2
+        freqs_2exp = torch.arange(max_res, dtype=torch.float32)
+        res_dim = max(0, self.freq_dim - max_res)
+        freqs_res = torch.arange(res_dim, dtype=torch.float32) / res_dim * max_res
+        self.freqs = 2 ** torch.cat([freqs_2exp, freqs_res], dim=0)[:self.freq_dim]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: (..., 3) positions -> (..., channels)"""
+        *dims, _ = x.shape
+        self.freqs = self.freqs.to(x.device)
+        out = torch.outer(x.reshape(-1), self.freqs) * 2 * torch.pi
+        embed = torch.cat([torch.sin(out), torch.cos(out)], dim=-1).reshape(*dims, -1)
+        if embed.shape[-1] < self.channels:
+            embed = torch.cat([embed, torch.zeros(*dims, self.channels - embed.shape[-1], device=embed.device)], dim=-1)
+        return embed
+
+
+class PcdAbsolutePositionEmbedderv2(nn.Module):
+    """Sinusoidal embedding of 3D points in [0, 1] (the decoders' query points)."""
+    def __init__(self, channels: int, in_channels: int = 3, max_res: int = 10):
+        super().__init__()
+        self.channels = channels
+        self.freq_dim = channels // in_channels // 2
+        logs = torch.linspace(0.0, float(max_res), steps=self.freq_dim, dtype=torch.float32)
+        self.register_buffer("freqs", torch.pow(2.0, logs), persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: (N, 3) -> (N, channels)"""
+        N = x.shape[0]
+        freqs = self.freqs.to(x.device, dtype=x.dtype)
+        ang = x.unsqueeze(-1) * freqs * 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
+
+
+class SinusoidalEmbedder(nn.Module):
+    """Embeds a scalar per batch element (diffusion time, octree level) into a vector."""
+    def __init__(self, hidden_size, frequency_embedding_size=256, max_period=10000, two_pi=False):
+        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
+        self.two_pi = two_pi
+
+    def forward(self, t):
+        half = self.frequency_embedding_size // 2
+        freqs = torch.exp(
+            -np.log(self.max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        if self.two_pi:
+            args = args * 2 * torch.pi
+        return self.mlp(torch.cat([torch.cos(args), torch.sin(args)], dim=-1))
+
+
+def apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+    """hidden_states [B, L, H, D] rotated by the complex frequencies freqs [B, L, H or 1, D/2]."""
+    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)
+
+
+class RePo3DRotaryEmbedding(nn.Module):
+    """Learned rotary position: each token's features predict a 3D offset per head, turned into RoPE angles."""
+    def __init__(self, model_channels: int, num_heads: int, head_dim: int, repo_hidden_ratio: float = 0.125,
+                 max_freq: float = 16.0):
+        super().__init__()
+        self.num_heads = num_heads
+        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)
+        # head_dim split into 3 parts for x, y, z
+        dim_0 = dim_1 = 2 * (head_dim // 6)
+        dims = [dim_0, dim_1, head_dim - dim_0 - dim_1]
+        freqs = [torch.linspace(1.0, float(max_freq), steps=d // 2, dtype=torch.float32) for d in dims]
+        self.freqs_0 = nn.Parameter(freqs[0])
+        self.freqs_1 = nn.Parameter(freqs[1])
+        self.freqs_2 = nn.Parameter(freqs[2])
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """hidden_states [B, L, C] -> complex rotations [B, L, H, D/2]"""
+        h = self.norm(hidden_states.float())
+        out = self.final_map(self.act(self.gate_map(h)) * self.content_map(h))
+        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)
+        return torch.polar(torch.ones_like(ang), ang).type(torch.complex64)
+
+
+# ---------------------------------------------------------------------------------------------------------------------
+# Attention
+# ---------------------------------------------------------------------------------------------------------------------
+
+class MultiHeadAttention(nn.Module):
+    """Self or cross attention with RMS-normalised queries and keys (the decoders)."""
+    def __init__(self, channels: int, num_heads: int, ctx_channels: Optional[int] = None, cross: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        self.cross = cross
+        head_dim = channels // num_heads
+        if cross:
+            self.to_q = nn.Linear(channels, channels)
+            self.to_kv = nn.Linear(ctx_channels, channels * 2)
+        else:
+            self.to_qkv = nn.Linear(channels, channels * 3)
+        self.q_rms_norm = MultiHeadRMSNorm(head_dim, num_heads)
+        self.k_rms_norm = MultiHeadRMSNorm(head_dim, num_heads)
+        self.to_out = nn.Linear(channels, channels)
+
+    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, L, C = x.shape
+        if self.cross:
+            q = self.q_rms_norm(self.to_q(x).reshape(B, L, self.num_heads, -1))
+            k, v = self.to_kv(context).reshape(B, context.shape[1], 2, self.num_heads, -1).unbind(dim=2)
+            h = scaled_dot_product_attention(q, self.k_rms_norm(k), v)
+        else:
+            q, k, v = self.to_qkv(x).reshape(B, L, 3, self.num_heads, -1).unbind(dim=2)
+            h = scaled_dot_product_attention(torch.stack([self.q_rms_norm(q), self.k_rms_norm(k), v], dim=2))
+        return self.to_out(h.reshape(B, L, -1))
+
+
+class RopeMultiHeadAttention(nn.Module):
+    """Self attention with rotary embeddings and RMS-normalised queries and keys (the denoiser)."""
+    def __init__(self, channels: int, num_heads: int):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = channels // num_heads
+        self.qkv = nn.Linear(channels, channels * 3)
+        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: torch.Tensor, rope_emb: torch.Tensor) -> torch.Tensor:
+        B, L, C = x.shape
+        q, k, v = self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim).unbind(2)
+        q = self.q_norm(apply_rotary_emb(q, rope_emb))
+        k = self.k_norm(apply_rotary_emb(k, rope_emb))
+        return self.out(scaled_dot_product_attention(q, k, v).reshape(B, L, C))
+
+
+# ---------------------------------------------------------------------------------------------------------------------
+# Transformer blocks
+# ---------------------------------------------------------------------------------------------------------------------
+
+class UnifiedTransformerBlock(nn.Module):
+    """Denoiser block: RoPE self attention + FFN, modulated by the diffusion time through a shared adaLN plus
+    this block's own shift table (modulation=False for the condition refiner, which has plain LayerNorms)."""
+    def __init__(self, channels: int, num_heads: int, mlp_ratio: float = 4.0, modulation: bool = True):
+        super().__init__()
+        self.modulation = modulation
+        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)
+        self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+        if modulation:
+            self.shift_table = nn.Parameter(torch.zeros(1, 6 * channels))
+
+    def forward(self, x: torch.Tensor, mod: Optional[torch.Tensor], rotary_emb: torch.Tensor) -> torch.Tensor:
+        if not self.modulation:
+            x = x + self.attn(self.norm1(x), rope_emb=rotary_emb)
+            return x + self.mlp(self.norm2(x))
+        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)
+        h = self.mlp(h)
+        return x + h * gate_mlp.unsqueeze(1)
+
+
+class TransformerCrossBlock(nn.Module):
+    """Gaussian decoder block: self attention + cross attention to the latent + FFN."""
+    def __init__(self, channels: int, ctx_channels: int, num_heads: int, mlp_ratio: float = 4.0):
+        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)
+        self.cross_attn = MultiHeadAttention(channels, num_heads=num_heads, ctx_channels=ctx_channels, cross=True)
+        self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+
+    def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
+        x = x + self.self_attn(self.norm1(x))
+        x = x + self.cross_attn(self.norm2(x), context)
+        return x + self.mlp(self.norm3(x))
+
+
+class ModulatedTransformerCrossOnlyBlock(nn.Module):
+    """Octree decoder block: cross attention to the latent + FFN, modulated by the octree level (shared adaLN)."""
+    def __init__(self, channels: int, ctx_channels: int, num_heads: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        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, num_heads=num_heads, ctx_channels=ctx_channels, cross=True)
+        self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+
+    def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
+        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.cross_attn(h, context)
+        x = x + h * gate_msa.unsqueeze(1)
+        h = self.norm2(x)
+        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+        h = self.mlp(h)
+        return x + h * gate_mlp.unsqueeze(1)

idol/models/transformer_base.py

@@ -1,351 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-from ..modules.utils import convert_module_to_f16, convert_module_to_f32
-from ..modules.transformer import AbsolutePositionEmbedder, TransformerCrossBlock, TransformerBlock
-from ..modules.transformer.blocks import FeedForwardNet
-from ..modules.norm import LayerNorm32
-from ..modules.attention import MultiHeadAttention
-from ..modules.transformer.modulated import ModulatedTransformerBlock, ModulatedTransformerCrossBlock, ModulatedTransformerCrossOnlyBlock
-import numpy as np
-
-class PcdAbsolutePositionEmbedder(nn.Module):
-    """
-    Embeds spatial positions in [0,1] into vector representations.
-    """
-    def __init__(self, channels: int, in_channels: int = 3, max_res=16):
-        super().__init__()
-        self.channels = channels
-        self.in_channels = in_channels
-        self.freq_dim = channels // in_channels // 2
-        self.max_res = max_res
-        freqs_2exp = torch.arange(max_res, dtype=torch.float32)
-        res_dim = max(0, self.freq_dim - max_res)
-        freqs_res = torch.arange(res_dim, dtype=torch.float32) / res_dim * max_res
-        self.freqs = torch.cat([freqs_2exp, freqs_res], dim=0)[:self.freq_dim]
-        self.freqs = 2 ** self.freqs # [1, 2 ^ max_res]
-        
-    def _sin_cos_embedding(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Create sinusoidal position embeddings.
-
-        Args:
-            x: a 1-D Tensor of N indices (in [0,1]^3 range)
-
-        Returns:
-            an (N, D) Tensor of positional embeddings.
-        """
-        self.freqs = self.freqs.to(x.device)
-        out = torch.outer(x, self.freqs) * 2 * torch.pi
-        out = torch.cat([torch.sin(out), torch.cos(out)], dim=-1)
-        return out
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            x (torch.Tensor): (N, D) or (B, N, D) tensor of spatial positions
-        """
-        *dims, D = x.shape
-        assert D == self.in_channels, "Input dimension must match number of input channels"
-        embed = self._sin_cos_embedding(x.reshape(-1))
-        embed = embed.reshape(*dims, -1)
-        if embed.shape[-1] < self.channels:
-            embed = torch.cat([embed, torch.zeros(*dims, self.channels - embed.shape[-1], device=embed.device)], dim=-1)
-        return embed
-
-
-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.freq_dim = channels // in_channels // 2
-        logs = torch.linspace(0.0, float(max_res), steps=self.freq_dim, dtype=torch.float32)
-        self.register_buffer("freqs", torch.pow(2.0, logs), persistent=False)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        N, D = x.shape
-        assert D == self.in_channels, "Input dimension must match number of input channels"
-        freqs = self.freqs.to(x.device, dtype=x.dtype)
-        ang = x.unsqueeze(-1) * freqs * 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
-
-
-class TransformerBase(nn.Module):
-    """
-    Transformer without output layers.
-    Serve as the base class for encoder and decoder.
-    """
-    def __init__(
-        self,
-        in_channels: int,
-        model_channels: int,
-        cond_channels: Optional[int],
-        num_blocks: int,
-        num_heads: Optional[int] = None,
-        num_head_channels: Optional[int] = 64,
-        mlp_ratio: float = 4.0,
-        use_fp16: bool = False,
-        use_checkpoint: bool = False,
-        qk_rms_norm: bool = True,
-        qk_rms_norm_cross: bool = True,
-        use_2_cross_block: bool = False,
-        share_mod: bool = False,
-    ):
-        super().__init__()
-        assert not use_2_cross_block or (cond_channels is not None), "use_2_cross_block requires cond_channels"
-        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.use_fp16 = use_fp16
-        self.use_checkpoint = use_checkpoint
-        self.qk_rms_norm = qk_rms_norm
-        self.qk_rms_norm_cross = qk_rms_norm_cross
-        self.share_mod = share_mod
-        self.use_2_cross_block = use_2_cross_block
-
-        self.dtype = torch.float16 if use_fp16 else torch.float32
-
-        self.input_layer = nn.Linear(in_channels, model_channels)
-        self.l_embedder = None
-        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,
-                        use_checkpoint=self.use_checkpoint,
-                        qk_rms_norm=self.qk_rms_norm,
-                        qk_rms_norm_cross=self.qk_rms_norm_cross,
-                        use_2_cross_block=use_2_cross_block,
-                    )
-                    for _ in range(num_blocks)
-                ])
-        else:
-            self.blocks = nn.ModuleList([
-                TransformerBlock(
-                    model_channels,
-                    num_heads=self.num_heads,
-                    mlp_ratio=self.mlp_ratio,
-                    use_checkpoint=self.use_checkpoint,
-                    qk_rms_norm=self.qk_rms_norm,
-                )
-                for _ in range(num_blocks)
-            ])
-
-    @property
-    def device(self) -> torch.device:
-        """
-        Return the device of the model.
-        """
-        return next(self.parameters()).device
-
-    def convert_to_fp16(self) -> None:
-        """
-        Convert the torso of the model to float16.
-        """
-        self.blocks.apply(convert_module_to_f16)
-
-    def convert_to_fp32(self) -> None:
-        """
-        Convert the torso of the model to float32.
-        """
-        self.blocks.apply(convert_module_to_f32)
-
-    def initialize_weights(self) -> None:
-        # Initialize transformer layers:
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-        self.apply(_basic_init)
-
-
-    def forward(self, x: torch.Tensor, cond: Optional[torch.Tensor]=None, l: Optional[torch.Tensor]=None, cond2: Optional[torch.Tensor]=None, indices=None, context_indices=None, context2_indices=None) -> torch.Tensor:
-        h = self.input_layer(x)
-        h = h.type(self.dtype)
-        cond = cond.type(self.dtype) if cond is not None else None
-        cond2 = cond2.type(self.dtype) if cond2 is not None else None
-        for block in self.blocks:
-            if cond is None:
-                h = block(h, indices=indices)
-            else:
-                h = block(h, cond, context2=cond2, indices=indices, context_indices=context_indices, context2_indices=context2_indices)
-        h = h.type(x.dtype)
-        return h
-
-class LevelEmbedder(nn.Module):
-    """
-    Embeds scalar timesteps into vector representations.
-    """
-    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):
-        """
-        Create sinusoidal timestep embeddings.
-
-        Args:
-            t: a 1-D Tensor of N indices, one per batch element.
-                These may be fractional.
-            dim: the dimension of the output.
-            max_period: controls the minimum frequency of the embeddings.
-
-        Returns:
-            an (N, D) Tensor of positional embeddings.
-        """
-        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
-        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):
-        t_freq = self.level_embedding(t, self.frequency_embedding_size, self.max_period)
-        t_emb = self.mlp(t_freq)
-        return t_emb
-
-class ModulatedCrossOnlyTransformerBase(nn.Module):
-    """
-    Transformer without output layers.
-    Serve as the base class for encoder and decoder.
-    """
-    def __init__(
-        self,
-        in_channels: int,
-        model_channels: int,
-        cond_channels: Optional[int],
-        num_blocks: int,
-        num_heads: Optional[int] = None,
-        num_head_channels: Optional[int] = 64,
-        mlp_ratio: float = 4.0,
-        use_fp16: bool = False,
-        use_checkpoint: bool = False,
-        share_mod: bool = False,
-        additional_level_embed: bool = False,
-        qk_rms_norm_cross: bool = True,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        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.use_fp16 = use_fp16
-        self.use_checkpoint = use_checkpoint
-        self.share_mod = share_mod
-        self.qk_rms_norm_cross = qk_rms_norm_cross
-        self.dtype = torch.float16 if use_fp16 else torch.float32
-
-        self.input_layer = nn.Linear(in_channels, model_channels)
-        self.l_embedder = LevelEmbedder(model_channels)
-        if additional_level_embed:
-            self.l_embedder2 = LevelEmbedder(model_channels, max_period=100)
-        else:
-            self.l_embedder2 = 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,
-                    use_checkpoint=self.use_checkpoint,
-                    qk_rms_norm_cross=self.qk_rms_norm_cross,
-                    share_mod=self.share_mod
-                )
-                for _ in range(num_blocks)
-            ])
-        else:
-            raise ValueError("ModulatedCrossOnlyTransformerBase requires cond_channels to be not None.")
-
-    @property
-    def device(self) -> torch.device:
-        """
-        Return the device of the model.
-        """
-        return next(self.parameters()).device
-
-    def convert_to_fp16(self) -> None:
-        """
-        Convert the torso of the model to float16.
-        """
-        self.blocks.apply(convert_module_to_f16)
-
-    def convert_to_fp32(self) -> None:
-        """
-        Convert the torso of the model to float32.
-        """
-        self.blocks.apply(convert_module_to_f32)
-
-    def initialize_weights(self) -> None:
-        # Initialize transformer layers:
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-        self.apply(_basic_init)
-        
-        # Initialize timestep embedding MLP:
-        nn.init.normal_(self.l_embedder.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.l_embedder.mlp[2].weight, std=0.02)
-        
-        if self.l_embedder2 is not None:
-            # zero projection init
-            nn.init.normal_(self.l_embedder2.mlp[0].weight, std=0.02)
-            nn.init.constant_(self.l_embedder2.mlp[2].weight, 0)
-            nn.init.constant_(self.l_embedder2.mlp[2].bias, 0)
-        
-        # Zero-out adaLN modulation layers in DiT blocks:
-        if self.share_mod:
-            nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
-            nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
-        else:
-            for block in self.blocks:
-                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
-                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
-        
-    def forward(self, x: torch.Tensor, l: torch.Tensor, cond: torch.Tensor, l2: Optional[torch.Tensor] = None) -> torch.Tensor:
-        h = self.input_layer(x)
-        l = l.type(h.dtype)
-        l_emb = self.l_embedder(l)
-        if self.l_embedder2 is not None and l2 is not None:
-            l2 = l2.type(h.dtype)
-            l_emb = l_emb + self.l_embedder2(l2)
-        if self.share_mod:
-            l_emb = self.adaLN_modulation(l_emb)
-        l_emb = l_emb.type(self.dtype)
-        h = h.type(self.dtype)
-        cond = cond.type(self.dtype)
-        for block in self.blocks:
-            h = block(h, l_emb, cond)
-        h = h.type(x.dtype)
-        return h

idol/modules/attention/__init__.py

@@ -1,14 +0,0 @@
-from typing import *
-import importlib.util
-import os
-
-# flash-attn when it is installed, PyTorch SDPA otherwise (same result, flash-attn is faster).
-# Override with ATTN_BACKEND=flash_attn|sdpa|naive.
-BACKEND = 'flash_attn' if importlib.util.find_spec('flash_attn') is not None else 'sdpa'
-BACKEND = os.environ.get('ATTN_BACKEND', BACKEND)
-if BACKEND not in ('flash_attn', 'sdpa', 'naive'):
-    raise ValueError(f'Unknown attention backend: {BACKEND}')
-
-
-from .full_attn import *
-from .modules import *

idol/modules/attention/full_attn.py

@@ -1,127 +0,0 @@
-from typing import *
-import torch
-import math
-from torch.nn.functional import scaled_dot_product_attention as sdpa
-from . import BACKEND
-
-if BACKEND == 'flash_attn':
-    import flash_attn
-
-
-__all__ = [
-    'scaled_dot_product_attention',
-]
-
-
-def _naive_sdpa(q, k, v):
-    """
-    Naive implementation of scaled dot product attention.
-    """
-    q = q.permute(0, 2, 1, 3)   # [N, H, L, C]
-    k = k.permute(0, 2, 1, 3)   # [N, H, L, C]
-    v = v.permute(0, 2, 1, 3)   # [N, H, L, C]
-    scale_factor = 1 / math.sqrt(q.size(-1))
-    attn_weight = q @ k.transpose(-2, -1) * scale_factor
-    attn_weight = torch.softmax(attn_weight, dim=-1)
-    out = attn_weight @ v
-    out = out.permute(0, 2, 1, 3)   # [N, L, H, C]
-    return out
-
-
-@overload
-def scaled_dot_product_attention(qkv: torch.Tensor) -> torch.Tensor:
-    """
-    Apply scaled dot product attention.
-
-    Args:
-        qkv (torch.Tensor): A [N, L, 3, H, C] tensor containing Qs, Ks, and Vs.
-    """
-    ...
-
-@overload
-def scaled_dot_product_attention(q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
-    """
-    Apply scaled dot product attention.
-
-    Args:
-        q (torch.Tensor): A [N, L, H, C] tensor containing Qs.
-        kv (torch.Tensor): A [N, L, 2, H, C] tensor containing Ks and Vs.
-    """
-    ...
-
-@overload
-def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
-    """
-    Apply scaled dot product attention.
-
-    Args:
-        q (torch.Tensor): A [N, L, H, Ci] tensor containing Qs.
-        k (torch.Tensor): A [N, L, H, Ci] tensor containing Ks.
-        v (torch.Tensor): A [N, L, H, Co] tensor containing Vs.
-
-    Note:
-        k and v are assumed to have the same coordinate map.
-    """
-    ...
-
-def scaled_dot_product_attention(*args, **kwargs):
-    arg_names_dict = {
-        1: ['qkv'],
-        2: ['q', 'kv'],
-        3: ['q', 'k', 'v']
-    }
-    num_all_args = len(args) + len(kwargs)
-    assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
-    for key in arg_names_dict[num_all_args][len(args):]:
-        assert key in kwargs, f"Missing argument {key}"
-
-    if num_all_args == 1:
-        qkv = args[0] if len(args) > 0 else kwargs['qkv']
-        assert len(qkv.shape) == 5 and qkv.shape[2] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, L, 3, H, C]"
-        device = qkv.device
-
-    elif num_all_args == 2:
-        q = args[0] if len(args) > 0 else kwargs['q']
-        kv = args[1] if len(args) > 1 else kwargs['kv']
-        assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
-        assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
-        assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
-        device = q.device
-
-    elif num_all_args == 3:
-        q = args[0] if len(args) > 0 else kwargs['q']
-        k = args[1] if len(args) > 1 else kwargs['k']
-        v = args[2] if len(args) > 2 else kwargs['v']
-        assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
-        assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
-        assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
-        assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
-        device = q.device    
-
-    if BACKEND == 'flash_attn':
-        if num_all_args == 1:
-            out = flash_attn.flash_attn_qkvpacked_func(qkv)
-        elif num_all_args == 2:
-            out = flash_attn.flash_attn_kvpacked_func(q, kv)
-        elif num_all_args == 3:
-            out = flash_attn.flash_attn_func(q, k, v)
-    elif BACKEND == 'sdpa':
-        if num_all_args == 1:
-            q, k, v = qkv.unbind(dim=2)
-        elif num_all_args == 2:
-            k, v = kv.unbind(dim=2)
-        q = q.permute(0, 2, 1, 3)   # [N, H, L, C]
-        k = k.permute(0, 2, 1, 3)   # [N, H, L, C]
-        v = v.permute(0, 2, 1, 3)   # [N, H, L, C]
-        out = sdpa(q, k, v)         # [N, H, L, C]
-        out = out.permute(0, 2, 1, 3)   # [N, L, H, C]
-    elif BACKEND == 'naive':
-        if num_all_args == 1:
-            q, k, v = qkv.unbind(dim=2)
-        elif num_all_args == 2:
-            k, v = kv.unbind(dim=2)
-        out = _naive_sdpa(q, k, v)
-    else:
-        raise ValueError(f"Unknown attention module: {BACKEND}")
-    
-    return out

idol/modules/attention/modules.py

@@ -1,338 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from .full_attn import scaled_dot_product_attention
-from ...modules.norm import LayerNorm32
-
-class MultiHeadRMSNorm(nn.Module):
-    def __init__(self, dim: int, heads: int):
-        super().__init__()
-        self.scale = dim ** 0.5
-        self.gamma = nn.Parameter(torch.ones(heads, dim))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return (F.normalize(x.float(), dim = -1) * self.gamma * self.scale).to(x.dtype)
-
-
-class RotaryPositionEmbedder(nn.Module):
-    def __init__(self, hidden_size: int, in_channels: int = 3):
-        super().__init__()
-        assert hidden_size % 2 == 0, "Hidden size must be divisible by 2"
-        self.hidden_size = hidden_size
-        self.in_channels = in_channels
-        self.freq_dim = hidden_size // in_channels // 2
-        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
-        self.freqs = 1.0 / (10000 ** self.freqs)
-        
-    def _get_phases(self, indices: torch.Tensor) -> torch.Tensor:
-        self.freqs = self.freqs.to(indices.device)
-        phases = torch.outer(indices, self.freqs)
-        phases = torch.polar(torch.ones_like(phases), phases)
-        return phases
-        
-    def _rotary_embedding(self, x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor:
-        x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
-        x_rotated = x_complex * phases
-        x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype)
-        return x_embed
-        
-    def forward(self, q: torch.Tensor, k: torch.Tensor, indices: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
-        if indices is None:
-            indices = torch.arange(q.shape[-2], device=q.device)
-            if len(q.shape) > 2:
-                indices = indices.unsqueeze(0).expand(q.shape[:-2] + (-1,))
-        
-        phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1)
-        if phases.shape[-1] < self.hidden_size // 2:
-            phases = torch.cat([phases, torch.polar(
-                torch.ones(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device),
-                torch.zeros(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device)
-            )], dim=-1)
-        q_embed = self._rotary_embedding(q, phases)
-        k_embed = self._rotary_embedding(k, phases)
-        return q_embed, k_embed
-    
-    def apply_rope(self, x: torch.Tensor, indices: Optional[torch.Tensor]) -> torch.Tensor:
-        if indices is None:
-            raise ValueError()
-        phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1)
-        if phases.shape[-1] < self.hidden_size // 2:
-            phases = torch.cat([phases, torch.polar(
-                torch.ones(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[-1], device=phases.device),
-                torch.zeros(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[-1], device=phases.device)
-            )], dim=-1)
-        return self._rotary_embedding(x, phases)
-
-
-class MultiHeadAttention(nn.Module):
-    def __init__(
-        self,
-        channels: int,
-        num_heads: int,
-        ctx_channels: Optional[int]=None,
-        type: Literal["self", "cross"] = "self",
-        qkv_bias: bool = True,
-        qk_rms_norm: bool = False,
-    ):
-        super().__init__()
-        assert channels % num_heads == 0
-        assert type in ["self", "cross"], f"Invalid attention type: {type}"
-
-        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: torch.Tensor, context: Optional[torch.Tensor] = None, indices: Optional[torch.Tensor] = None, context_indices: Optional[torch.Tensor] = None) -> torch.Tensor:
-        B, L, C = x.shape
-        if self._type == "self":
-            qkv = self.to_qkv(x)
-            qkv = qkv.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)
-        else:
-            Lkv = context.shape[1]
-            q = self.to_q(x)
-            kv = self.to_kv(context)
-            q = q.reshape(B, L, self.num_heads, -1)
-            kv = kv.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, k, v)
-            else:
-                h = scaled_dot_product_attention(q, kv)
-        h = h.reshape(B, L, -1)
-        h = self.to_out(h)
-        return h
-
-  
-
-def apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
-    """
-    Apply rotary embeddings to input tensors using the given frequencies.
-    Args:
-        hidden_states: [B, L, H, D]
-        freqs: [B, L, 1, D/2] complex tensor
-    """
-    # [B, L, H, D] -> [B, L, H, D/2, 2] -> [B, L, H, D/2] (complex)
-    x_rotated = torch.view_as_complex(hidden_states.float().reshape(*hidden_states.shape[:-1], -1, 2))
-    
-    # [B, L, H, D/2] * [B, L, 1, D/2] (broadcast) -> [B, L, H, D/2]
-    x_rotated = x_rotated * freqs
-    
-    # [B, L, H, D/2] -> [B, L, H, D/2, 2] -> [B, L, H, D]
-    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)
-
-class RePo3DRotaryEmbedding(nn.Module):
-    def __init__(self, model_channels: int, num_heads: int, head_dim: int, repo_hidden_ratio: float = 0.125, max_freq: float = 16.0):
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = head_dim
-        repo_hidden_size = int(model_channels * repo_hidden_ratio)
-        
-        # Lightweight Residual Network
-        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)
-        
-        # Split head_dim into 3 parts for x, y, z (same as Rotary3DPoseEmbedding)
-        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
-        self.max_freq = max_freq
-        
-        dims = [self.dim_0, self.dim_1, self.dim_2]
-        freqs_list = []
-        for d in dims:
-            freq_dim = d // 2
-            freqs = torch.linspace(1.0, float(max_freq), steps=freq_dim, dtype=torch.float32)
-            freqs_list.append(freqs)
-        
-        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 initialize_weights(self):
-        # Initialize proj_out with small values
-        nn.init.constant_(self.final_map.weight, 0)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        # hidden_states: [B, L, C]
-        # base_rotary_emb: [1, L, 1, D/2] (or similar broadcastable shape)
-        # pos: [B, L, 3] original 3D coordinates
-        hidden_states = hidden_states.float()
-        # Lightweight fusion logic
-        h = self.norm(hidden_states)
-        feat = self.act(self.gate_map(h)) * self.content_map(h)
-        out = self.final_map(feat)
-        
-        # Reshape to [B, L, 3 * H] -> [B, L, H, 3]
-        B, L, _ = out.shape
-        delta_pos = out.reshape(B, L, self.num_heads, 3)
-        
-        # Compute angles for delta (rope like)
-        # delta_pos[..., 0] is [B, L, H]
-        # freqs_0 is [D0/2]
-        # We want [B, L, H, D0/2]
-        # delta_pos gradient scale clamping, avoid large freq gradient exploding
-        # forward equivalent to delta_pos * self.freqs_I * torch.pi
-        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
-        
-        # Concatenate angles [B, L, H, D/2]
-        ang = torch.cat([ang_0, ang_1, ang_2], dim=-1)
-        
-        # Complex rotation for delta
-        delta_freqs_cis = torch.polar(torch.ones_like(ang), ang)
-        
-        # Apply residual rotation: base * delta
-        # base_rotary_emb will broadcast to [B, L, H, D/2]
-        return delta_freqs_cis.type(torch.complex64)
-
-
-class Rotary3DPoseEmbedding(nn.Module):
-    def __init__(self, head_dim: int, max_res: int = 10):
-        super().__init__()
-        self.head_dim = head_dim
-        
-        # Split head_dim into 3 parts for x, y, z
-        # Each part must be divisible by 2 for complex RoPE
-        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
-            logs = torch.linspace(0.0, float(max_res), steps=freq_dim, dtype=torch.float32)
-            freqs = torch.pow(2.0, logs)
-            freqs_list.append(freqs)
-            
-        self.register_buffer("freqs_0", freqs_list[0])
-        self.register_buffer("freqs_1", freqs_list[1])
-        self.register_buffer("freqs_2", freqs_list[2])
-
-    def forward(self, pos: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            pos: [B, L, 3] tensor of spatial positions in [0, 1]
-        Returns:
-            rope_emb: [B, L, 1, D/2] complex tensor
-        """
-        # pos[..., 0]: [B, L]
-        # freqs_0: [D0/2]
-        # ang: [B, L, D0/2]
-        
-        ang_0 = torch.outer(pos[..., 0].flatten(), self.freqs_0).view(*pos.shape[:-1], -1) * torch.pi
-        ang_1 = torch.outer(pos[..., 1].flatten(), self.freqs_1).view(*pos.shape[:-1], -1) * torch.pi
-        ang_2 = torch.outer(pos[..., 2].flatten(), self.freqs_2).view(*pos.shape[:-1], -1) * torch.pi
-        
-        # Concatenate angles [B, L, D/2]
-        ang = torch.cat([ang_0, ang_1, ang_2], dim=-1)
-        
-        # Convert to complex exponentials
-        freqs_cis = torch.polar(torch.ones_like(ang), ang)
-        
-        # Add head dimension for broadcasting: [B, L, 1, D/2]
-        return freqs_cis.unsqueeze(-2).type(torch.complex64)
-
-
-class RopeMultiHeadAttention(nn.Module):
-    def __init__(
-        self,
-        channels: int,
-        num_heads: int,
-        ctx_channels: Optional[int]=None,
-        type: Literal["self", "cross"] = "self",
-        qkv_bias: bool = True,
-        qk_rms_norm: bool = False,
-        use_rope: bool = 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
-        
-        assert channels % num_heads == 0
-
-        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: torch.Tensor, context: Optional[torch.Tensor] = None, rope_emb: Optional[torch.Tensor] = None) -> torch.Tensor:
-        """
-        Args:
-            x: [B, L, C] query tensor
-            context: [B, S, C] context tensor (for cross attention)
-            rope_emb: [B, L, 1, D/2] precomputed complex RoPE embedding
-        """
-        B, L, C = x.shape
-        
-        if self._type == "self":
-            # [B, L, 3, H, D]
-            qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim)
-            q, k, v = qkv.unbind(2) # [B, L, H, D]
-            
-            if self.use_rope:
-                q = apply_rotary_emb(q, rope_emb)
-                k = apply_rotary_emb(k, rope_emb)
-        else:
-            # Cross attention
-            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)
-            # No RoPE for cross attention per instructions
-            
-        if self.qk_rms_norm:
-            q = self.q_norm(q)
-            k = self.k_norm(k)
-            
-        h = scaled_dot_product_attention(q, k, v)
-        
-        # [B, L, H, D] -> [B, L, C]
-        h = h.reshape(B, L, C)
-        return self.out(h)

idol/modules/mlp.py

@@ -1,16 +0,0 @@
-import torch
-import torch.nn as nn
-from typing import *
-
-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)

idol/modules/norm.py

@@ -1,29 +0,0 @@
-import torch
-import torch.nn as nn
-
-
-class LayerNorm32(nn.LayerNorm):
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return super().forward(x.float()).type(x.dtype)
-    
-class RMSNorm32(nn.RMSNorm):
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return super().forward(x.float()).type(x.dtype)
-    
-
-class GroupNorm32(nn.GroupNorm):
-    """
-    A GroupNorm layer that converts to float32 before the forward pass.
-    """
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return super().forward(x.float()).type(x.dtype)
-    
-    
-class ChannelLayerNorm32(LayerNorm32):
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        DIM = x.dim()
-        x = x.permute(0, *range(2, DIM), 1).contiguous()
-        x = super().forward(x)
-        x = x.permute(0, DIM-1, *range(1, DIM-1)).contiguous()
-        return x
-    

idol/modules/transformer/__init__.py

@@ -1,2 +0,0 @@
-from .blocks import *
-from .modulated import *

idol/modules/transformer/blocks.py

@@ -1,177 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-from ..attention import MultiHeadAttention
-from ..norm import LayerNorm32
-
-
-class AbsolutePositionEmbedder(nn.Module):
-    """
-    Embeds spatial positions into vector representations.
-    """
-    def __init__(self, channels: int, in_channels: int = 3):
-        super().__init__()
-        self.channels = channels
-        self.in_channels = in_channels
-        self.freq_dim = channels // in_channels // 2
-        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
-        self.freqs = 1.0 / (10000 ** self.freqs)
-        
-    def _sin_cos_embedding(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Create sinusoidal position embeddings.
-
-        Args:
-            x: a 1-D Tensor of N indices
-
-        Returns:
-            an (N, D) Tensor of positional embeddings.
-        """
-        self.freqs = self.freqs.to(x.device)
-        out = torch.outer(x, self.freqs)
-        out = torch.cat([torch.sin(out), torch.cos(out)], dim=-1)
-        return out
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            x (torch.Tensor): (N, D) tensor of spatial positions
-        """
-        N, D = x.shape
-        assert D == self.in_channels, "Input dimension must match number of input channels"
-        embed = self._sin_cos_embedding(x.reshape(-1))
-        embed = embed.reshape(N, -1)
-        if embed.shape[1] < self.channels:
-            embed = torch.cat([embed, torch.zeros(N, self.channels - embed.shape[1], device=embed.device)], dim=-1)
-        return embed
-
-
-class FeedForwardNet(nn.Module):
-    def __init__(self, channels: int, mlp_ratio: float = 4.0, channels_out: Optional[int] = 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: torch.Tensor) -> torch.Tensor:
-        return self.mlp(x)
-
-class TransformerBlock(nn.Module):
-    """
-    Transformer block (MSA + FFN).
-    """
-    def __init__(
-        self,
-        channels: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        use_checkpoint: bool = False,
-        qk_rms_norm: bool = False,
-        qkv_bias: bool = True,
-    ):
-        super().__init__()
-        self.use_checkpoint = use_checkpoint
-        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
-        self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
-        self.attn = MultiHeadAttention(
-            channels,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            qk_rms_norm=qk_rms_norm,
-        )
-        self.mlp = FeedForwardNet(
-            channels,
-            mlp_ratio=mlp_ratio,
-        )
-
-    def _forward(self, x: torch.Tensor, indices=None) -> torch.Tensor:
-        h = self.norm1(x)
-        h = self.attn(h, indices=indices)
-        x = x + h
-        h = self.norm2(x)
-        h = self.mlp(h)
-        x = x + h
-        return x
-
-    def forward(self, x: torch.Tensor, indices=None) -> torch.Tensor:
-        if self.use_checkpoint:
-            return torch.utils.checkpoint.checkpoint(self._forward, x, indices, use_reentrant=False)
-        else:
-            return self._forward(x, indices=indices)
-
-
-class TransformerCrossBlock(nn.Module):
-    """
-    Transformer cross-attention block (MSA + MCA + FFN).
-    """
-    def __init__(
-        self,
-        channels: int,
-        ctx_channels: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        use_checkpoint: bool = False,
-        qk_rms_norm: bool = False,
-        qk_rms_norm_cross: bool = False,
-        qkv_bias: bool = True,
-        use_2_cross_block: bool = False,
-    ):
-        super().__init__()
-        self.use_checkpoint = use_checkpoint
-        self.use_2_cross_block = use_2_cross_block
-        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",
-            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",
-            qkv_bias=qkv_bias,
-            qk_rms_norm=qk_rms_norm_cross,
-        )
-        if self.use_2_cross_block:
-            self.norm2_1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
-            self.cross_attn2 = MultiHeadAttention(
-                channels,
-                ctx_channels=ctx_channels,
-                num_heads=num_heads,
-                type="cross",
-                qkv_bias=qkv_bias,
-                qk_rms_norm=qk_rms_norm_cross,
-            )
-        self.mlp = FeedForwardNet(
-            channels,
-            mlp_ratio=mlp_ratio,
-        )
-
-    def _forward(self, x: torch.Tensor, context: torch.Tensor, context2=None, indices=None, context_indices=None, context2_indices=None) -> torch.Tensor:
-        h = self.norm1(x)
-        h = self.self_attn(h, indices=indices)
-        x = x + h
-        h = self.norm2(x)
-        h = self.cross_attn(h, context, context_indices=context_indices)
-        x = x + h
-        if self.use_2_cross_block:
-            h = self.norm2_1(x)
-            h = self.cross_attn2(h, context2, context_indices=context2_indices)
-            x = x + h
-        h = self.norm3(x)
-        h = self.mlp(h)
-        x = x + h
-        return x
-
-    def forward(self, x: torch.Tensor, context: torch.Tensor, context2=None, indices=None, context_indices=None, context2_indices=None) -> torch.Tensor:
-        if self.use_checkpoint:
-            return torch.utils.checkpoint.checkpoint(self._forward, x, context, context2, indices, context_indices, context2_indices, use_reentrant=False)
-        else:
-            return self._forward(x, context, context2=context2, indices=indices, context_indices=context_indices, context2_indices=context2_indices)

idol/modules/transformer/modulated.py

@@ -1,381 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-
-from ..attention.modules import RopeMultiHeadAttention
-from ..attention import MultiHeadAttention
-from ..norm import LayerNorm32, RMSNorm32
-from .blocks import FeedForwardNet
-
-
-class ModulatedTransformerBlock(nn.Module):
-    """
-    Transformer block (MSA + FFN) with adaptive layer norm conditioning.
-    """
-    def __init__(
-        self,
-        channels: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        use_checkpoint: bool = False,
-        use_rope: bool = False,
-        qk_rms_norm: bool = False,
-        qkv_bias: bool = True,
-        share_mod: bool = False,
-    ):
-        super().__init__()
-        self.use_checkpoint = use_checkpoint
-        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.attn = MultiHeadAttention(
-            channels,
-            num_heads=num_heads,
-            qkv_bias=qkv_bias,
-            qk_rms_norm=qk_rms_norm,
-        )
-        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: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:
-        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)
-        h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
-        h = self.attn(h)
-        h = h * gate_msa.unsqueeze(1)
-        x = x + h
-        h = self.norm2(x)
-        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
-        h = self.mlp(h)
-        h = h * gate_mlp.unsqueeze(1)
-        x = x + h
-        return x
-
-    def forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:
-        if self.use_checkpoint:
-            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, use_reentrant=False)
-        else:
-            return self._forward(x, mod)
-
-
-class ModulatedTransformerCrossBlock(nn.Module):
-    """
-    Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.
-    """
-    def __init__(
-        self,
-        channels: int,
-        ctx_channels: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        use_checkpoint: bool = False,
-        use_rope: bool = False,
-        qk_rms_norm: bool = False,
-        qk_rms_norm_cross: bool = False,
-        qkv_bias: bool = True,
-        share_mod: bool = False,
-        use_shift_table: bool = False,
-    ):
-        super().__init__()
-        self.use_checkpoint = use_checkpoint
-        self.share_mod = share_mod
-        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",
-            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",
-            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)
-            )
-        if use_shift_table:
-            self.shift_table = nn.Parameter(torch.randn(1, 6 * channels) / channels**0.5)
-        else:
-            self.shift_table = None
-
-    def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor):
-        if not self.share_mod:
-            mod = self.adaLN_modulation(mod)
-        if 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.self_attn(h)
-        h = h * gate_msa.unsqueeze(1)
-        x = x + h
-        h = self.norm2(x)
-        h = self.cross_attn(h, context)
-        x = x + h
-        h = self.norm3(x)
-        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
-        h = self.mlp(h)
-        h = h * gate_mlp.unsqueeze(1)
-        x = x + h
-        return x
-
-    def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor):
-        if self.use_checkpoint:
-            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, use_reentrant=False)
-        else:
-            return self._forward(x, mod, context)
-
-
-class UnifiedTransformerBlock(nn.Module):
-    """
-    Transformer block (MSA + FFN) with optional adaptive layer norm conditioning and RoPE support.
-    Used for Unified MMDiT.
-    """
-    def __init__(
-        self,
-        channels: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        use_checkpoint: bool = False,
-        use_rope: bool = False,
-        qk_rms_norm: bool = False,
-        qkv_bias: bool = True,
-        modulation: bool = True,
-        share_mod: bool = False,
-        use_shift_table: bool = False,
-    ):
-        super().__init__()
-        self.use_checkpoint = use_checkpoint
-        self.modulation = modulation
-        self.share_mod = share_mod
-        
-        # If modulation is used, we usually disable elementwise_affine in LN because modulation handles it.
-        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",
-            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)
-                )
-            if use_shift_table:
-                self.shift_table = nn.Parameter(torch.randn(1, 6 * channels) / channels**0.5)
-            else:
-                self.shift_table = None
-
-    def _forward(self, x: torch.Tensor, mod: Optional[torch.Tensor] = None, rotary_emb: Optional[torch.Tensor] = None):
-        if self.modulation:
-            assert mod is not None, "Modulation tensor must be provided when modulation is enabled"
-            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)
-            h = h * gate_msa.unsqueeze(1)
-            x = x + h
-            
-            h = self.norm2(x)
-            h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
-            h = self.mlp(h)
-            h = h * gate_mlp.unsqueeze(1)
-            x = x + h
-        else:
-            h = self.norm1(x)
-            h = self.attn(h, rope_emb=rotary_emb)
-            x = x + h
-            
-            h = self.norm2(x)
-            h = self.mlp(h)
-            x = x + h
-            
-        return x
-
-    def forward(self, x: torch.Tensor, mod: Optional[torch.Tensor] = None, rotary_emb: Optional[torch.Tensor] = None):
-        if self.use_checkpoint:
-            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, rotary_emb, use_reentrant=False)
-        else:
-            return self._forward(x, mod, rotary_emb)
-
-
-class ModulatedTransformerCrossBlockv2(nn.Module):
-    """
-    Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.
-    """
-    def __init__(
-        self,
-        channels: int,
-        ctx_channels: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        use_checkpoint: bool = False,
-        use_rope: bool = False,
-        qk_rms_norm: bool = False,
-        qk_rms_norm_cross: bool = False,
-        qkv_bias: bool = True,
-        share_mod: bool = False,
-        use_shift_table: bool = False,
-    ):
-        super().__init__()
-        self.use_checkpoint = use_checkpoint
-        self.share_mod = share_mod
-        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 = RopeMultiHeadAttention(
-            channels,
-            num_heads=num_heads,
-            type="self",
-            qkv_bias=qkv_bias,
-            use_rope=use_rope,
-            qk_rms_norm=qk_rms_norm,
-        )
-        self.cross_attn = RopeMultiHeadAttention(
-            channels,
-            ctx_channels=ctx_channels,
-            num_heads=num_heads,
-            type="cross",
-            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)
-            )
-        if use_shift_table:
-            self.shift_table = nn.Parameter(torch.randn(1, 6 * channels) / channels**0.5)
-        else:
-            self.shift_table = None
-
-    def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor, rotary_emb: Optional[torch.Tensor] = None):
-        if not self.share_mod:
-            mod = self.adaLN_modulation(mod)
-        if 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.self_attn(h, rope_emb=rotary_emb)
-        h = h * gate_msa.unsqueeze(1)
-        x = x + h
-        h = self.norm2(x)
-        h = self.cross_attn(h, context)
-        x = x + h
-        h = self.norm3(x)
-        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
-        h = self.mlp(h)
-        h = h * gate_mlp.unsqueeze(1)
-        x = x + h
-        return x
-
-    def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor, rotary_emb: Optional[torch.Tensor] = None):
-        if self.use_checkpoint:
-            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, rotary_emb, use_reentrant=False)
-        else:
-            return self._forward(x, mod, context, rotary_emb)
-    
-class ModulatedTransformerCrossOnlyBlock(nn.Module):
-    """
-    Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.
-    """
-    def __init__(
-        self,
-        channels: int,
-        ctx_channels: int,
-        num_heads: int,
-        mlp_ratio: float = 4.0,
-        use_checkpoint: bool = False,
-        qk_rms_norm_cross: bool = True,
-        qkv_bias: bool = True,
-        share_mod: bool = False,
-    ):
-        super().__init__()
-        self.use_checkpoint = use_checkpoint
-        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",
-            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: torch.Tensor, mod: torch.Tensor, context: torch.Tensor):
-        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)
-        h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
-        h = self.cross_attn(h, context)
-        h = h * gate_msa.unsqueeze(1)
-        x = x + h
-        h = self.norm2(x)
-        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
-        h = self.mlp(h)
-        h = h * gate_mlp.unsqueeze(1)
-        x = x + h
-        return x
-
-    def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor):
-        if self.use_checkpoint:
-            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, use_reentrant=False)
-        else:
-            return self._forward(x, mod, context)

idol/modules/utils.py

@@ -1,50 +0,0 @@
-import torch.nn as nn
-
-FP16_MODULES = (
-    nn.Conv1d,
-    nn.Conv2d,
-    nn.Conv3d,
-    nn.ConvTranspose1d,
-    nn.ConvTranspose2d,
-    nn.ConvTranspose3d,
-    nn.Linear,
-)
-
-def convert_module_to_f16(l):
-    """
-    Convert primitive modules to float16.
-    """
-    if isinstance(l, FP16_MODULES):
-        for p in l.parameters():
-            p.data = p.data.half()
-
-
-def convert_module_to_f32(l):
-    """
-    Convert primitive modules to float32, undoing convert_module_to_f16().
-    """
-    if isinstance(l, FP16_MODULES):
-        for p in l.parameters():
-            p.data = p.data.float()
-
-
-def zero_module(module):
-    """
-    Zero out the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().zero_()
-    return module
-
-
-def scale_module(module, scale):
-    """
-    Scale the parameters of a module and return it.
-    """
-    for p in module.parameters():
-        p.detach().mul_(scale)
-    return module
-
-
-def modulate(x, shift, scale):
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

idol/pipeline.py

@@ -2,7 +2,7 @@
 import json
 import os
 from dataclasses import dataclass
-from typing import Callable, Optional, Union
+from typing import Optional, Union
 
 import numpy as np
 import torch
@@ -15,7 +15,7 @@ from .conditioning import ImageEncoder
 from .gaussian import Gaussian
 from .models import BiRefNet, GaussianFixedlenDecoder, LatentSeqMMFlowModel, OctreeProbabilityFixedlenDecoder
 from .preprocess import CROP_MARGIN, preprocess_image
-from .sampler import FlowEulerCfgSampler
+from .sampler import sample_flow
 
 DEFAULT_REPO = 'speridlabs/idol'
 WEIGHT_FILES = ['config.json', 'denoiser.safetensors', 'decoder.safetensors', 'decoder_gs.safetensors',
@@ -29,34 +29,25 @@ class IdolOutput:
     camera: dict                     # predicted input camera: yaw, pitch (radians), radius, fov (degrees)
 
 
-def _load(model, path):
-    model.load_state_dict(load_file(path), strict=True)
-    return model
-
-
 class IdolPipeline:
     def __init__(self, weights_dir: str, device: Optional[str] = None, dinov3_dtype: str = 'bf16'):
         self.device = torch.device(device or ('cuda' if torch.cuda.is_available() else 'cpu'))
         with open(os.path.join(weights_dir, 'config.json')) as f:
             self.config = json.load(f)
-        config = self.config
 
-        def weights(name):
-            return os.path.join(weights_dir, f'{name}.safetensors')
+        def load(model, name):
+            model.load_state_dict(load_file(os.path.join(weights_dir, f'{name}.safetensors')), strict=True)
+            return model.eval().requires_grad_(False).to(self.device)
 
-        self.denoiser = _load(LatentSeqMMFlowModel(**config['denoiser']['args']), weights('denoiser'))
-        self.decoder = _load(OctreeProbabilityFixedlenDecoder(**config['decoder']['args']), weights('decoder'))
-        self.decoder_gs = _load(GaussianFixedlenDecoder(**config['decoder_gs']['args']), weights('decoder_gs'))
-        self.background_remover = BiRefNet()
-        self.background_remover.load_safetensors(weights('birefnet'))
-        self.image_encoder = ImageEncoder(config['image_encoders']['dinov3'], config['image_encoders']['flux2_vae'],
-                                          resolution=config['image_encoders']['resolution'],
-                                          dinov3_dtype=dinov3_dtype, device=self.device)
-        for model in (self.denoiser, self.decoder, self.decoder_gs):
-            model.eval().requires_grad_(False).to(self.device)
-        remover_dtype = torch.float16 if self.device.type == 'cuda' else torch.float32
-        self.background_remover = self.background_remover.eval().to(self.device, remover_dtype)
-        self.sampler = FlowEulerCfgSampler(sigma_min=config['sigma_min'])
+        self.denoiser = load(LatentSeqMMFlowModel(**self.config['denoiser']['args']), 'denoiser')
+        self.decoder = load(OctreeProbabilityFixedlenDecoder(**self.config['decoder']['args']), 'decoder')
+        self.decoder_gs = load(GaussianFixedlenDecoder(**self.config['decoder_gs']['args']), 'decoder_gs')
+        self.background_remover = load(BiRefNet(), 'birefnet')
+        if self.device.type == 'cuda':
+            self.background_remover.half()
+        encoders = self.config['image_encoders']
+        self.image_encoder = ImageEncoder(encoders['dinov3'], encoders['flux2_vae'], dinov3_dtype=dinov3_dtype,
+                                          device=self.device)
 
     @classmethod
     def from_pretrained(cls, repo_or_dir: str = DEFAULT_REPO, **kwargs) -> 'IdolPipeline':
@@ -66,21 +57,14 @@ class IdolPipeline:
         from huggingface_hub import snapshot_download
         return cls(snapshot_download(repo_or_dir, allow_patterns=WEIGHT_FILES), **kwargs)
 
-    def to(self, device) -> 'IdolPipeline':
-        self.device = torch.device(device)
-        for model in (self.denoiser, self.decoder, self.decoder_gs, self.background_remover):
-            model.to(self.device)
-        self.image_encoder.to(self.device)
-        return self
-
     def preprocess(self, image: Union[str, Image.Image], mode: str = 'synthetic',
-                   crop_margin: Optional[float] = None, erode_radius: int = 1) -> Image.Image:
+                   crop_margin: Optional[float] = None) -> Image.Image:
         """mode 'synthetic' (renders, generated images of one object: crop 1.2) or 'real' (photos: crop 1.7);
         crop_margin overrides the mode's value."""
         if isinstance(image, str):
             image = Image.open(image)
         margin = crop_margin if crop_margin is not None else CROP_MARGIN[mode]
-        return preprocess_image(image, background_remover=self.background_remover, erode_radius=erode_radius,
+        return preprocess_image(image, background_remover=self.background_remover,
                                 image_size=self.config['image_size'], crop_margin=margin)
 
     @torch.no_grad()
@@ -94,34 +78,23 @@ class IdolPipeline:
         rescale_t: float = 3.0,
         num_gaussians: int = 262144,
         crop_margin: Optional[float] = None,
-        erode_radius: int = 1,
-        preprocess: bool = True,
-        callback: Optional[Callable[[int, int], None]] = None,
         verbose: bool = False,
     ) -> IdolOutput:
         """Generate the 3D Gaussians of the object in ``image``.
 
-        mode / crop_margin / erode_radius: see ``preprocess``. preprocess=False feeds the image as given (it
-        must already be a square, centred object on black).
-        num_gaussians: 32,768 to 1,048,576, rounded down to a multiple of 32.
-        callback(step_done, steps) is called after every sampling step."""
-        if preprocess:
-            prepared = self.preprocess(image, mode=mode, crop_margin=crop_margin, erode_radius=erode_radius)
-        else:
-            prepared = (Image.open(image) if isinstance(image, str) else image).convert('RGB')
+        mode / crop_margin: see ``preprocess``.
+        num_gaussians: 32,768 to 1,048,576, rounded down to a multiple of 32."""
+        prepared = self.preprocess(image, mode=mode, crop_margin=crop_margin)
         pixels = torch.from_numpy(np.array(prepared)).permute(2, 0, 1).float().div(255.0)[None].to(self.device)
 
         torch.manual_seed(seed)
-        np.random.seed(seed)
         cond = self.image_encoder(pixels)
-        neg_cond = cond.apply(torch.zeros_like)
         noise = TensorDict({
             'latent': torch.randn(1, self.denoiser.q_token_length, self.denoiser.in_channels, device=self.device),
             'camera': torch.randn(1, 1, self.denoiser.cam_channels, device=self.device),
         }, batch_size=1, device=self.device)
-        sample = self.sampler.sample(self.denoiser, noise, cond=cond, neg_cond=neg_cond, steps=steps,
-                                     rescale_t=rescale_t, cfg_strength=cfg_strength, verbose=verbose,
-                                     callback=callback)
+        sample = sample_flow(self.denoiser, noise, cond=cond, neg_cond=cond.apply(torch.zeros_like), steps=steps,
+                             rescale_t=rescale_t, cfg_strength=cfg_strength, verbose=verbose)
         gaussian = self.decode(sample['latent'], num_gaussians)
         return IdolOutput(gaussian=gaussian, preprocessed_image=prepared, camera=decode_camera(sample['camera'][0, 0]))
 
@@ -130,8 +103,7 @@ class IdolPipeline:
         """Latent [1, 8192, 16] -> Gaussians: the octree decoder places num_gaussians / 32 anchor points, the
         Gaussian decoder turns each anchor into 32 Gaussians."""
         per_anchor = self.decoder_gs.rep_config['num_gaussians']
-        anchors = OctreeProbabilityFixedlenDecoder.sample(
-            self.decoder, latent, num_points=max(1, int(num_gaussians) // per_anchor),
-            level=self.config['max_voxel_level'], temperature=1.0, algo='systematic')
-        features = self.decoder_gs(x=anchors, cond=latent)
-        return self.decoder_gs.to_representation(anchors, features)[0]
+        anchors = OctreeProbabilityFixedlenDecoder.sample(self.decoder, latent,
+                                                          num_points=max(1, int(num_gaussians) // per_anchor),
+                                                          level=self.config['max_voxel_level'])
+        return self.decoder_gs.to_representation(anchors, self.decoder_gs(anchors, latent))[0]

idol/preprocess.py

@@ -11,22 +11,19 @@ def has_real_alpha(image: Image.Image) -> bool:
     return image.mode == 'RGBA' and np.array(image.getchannel(3), dtype=np.int32).min() < 255
 
 
-def preprocess_image(image: Image.Image, background_remover=None, erode_radius: int = 1, image_size: int = 1024,
+def preprocess_image(image: Image.Image, background_remover, image_size: int = 1024,
                      crop_margin: float = 1.2) -> Image.Image:
     """Return the RGB image the encoders see.
 
     Resize so the short side is ``image_size``; if the image has no real alpha, ``background_remover`` (BiRefNet)
-    predicts one; erode the alpha by ``erode_radius`` px so no background-coloured fringe survives; crop a square of
+    predicts one; erode the alpha by 1 px so no background-coloured fringe survives; crop a square of
     side ``crop_margin`` x the object's bounding box; resize to ``image_size``; composite on black."""
     width, height = image.size
     scale = image_size / min(width, height)
     image = image.resize((max(1, int(round(width * scale))), max(1, int(round(height * scale)))), Image.LANCZOS)
     if not has_real_alpha(image):
-        if background_remover is None:
-            raise ValueError('the image has no alpha channel: pass a background remover or an RGBA cutout')
         image = background_remover.remove_background(image.convert('RGB'))
-    if erode_radius > 0:
-        image.putalpha(image.getchannel(3).filter(ImageFilter.MinFilter(2 * erode_radius + 1)))
+    image.putalpha(image.getchannel(3).filter(ImageFilter.MinFilter(3)))
     alpha = np.array(image.getchannel(3))
     ys, xs = np.nonzero(alpha)
     if len(xs) == 0:

idol/sampler.py

@@ -1,38 +1,28 @@
 """Euler sampler for the rectified-flow denoiser, with classifier-free guidance."""
-from typing import Callable, Optional
-
 import numpy as np
 import torch
 from tqdm import tqdm
 
 
-class FlowEulerCfgSampler:
-    def __init__(self, sigma_min: float):
-        self.sigma_min = sigma_min
+@torch.no_grad()
+def sample_flow(model, noise, cond, neg_cond, steps: int = 50, rescale_t: float = 3.0, cfg_strength: float = 3.0,
+                verbose: bool = True):
+    """Integrate from noise (t = 1) to data (t = 0) in ``steps`` Euler steps.
 
-    def _velocity(self, model, x_t, t: float, cond):
-        t = torch.tensor([1000 * t] * x_t.shape[0], device=x_t.device, dtype=torch.float32)
-        return model(x_t, t, cond)
+    rescale_t shifts the time grid towards t = 1 (more steps at high noise).
+    cfg_strength: guided velocity = (1 + w) * v(cond) - w * v(neg_cond)."""
 
-    @torch.no_grad()
-    def sample(self, model, noise, cond, neg_cond, steps: int = 50, rescale_t: float = 3.0,
-               cfg_strength: float = 3.0, verbose: bool = True,
-               callback: Optional[Callable[[int, int], None]] = None):
-        """Integrate from noise (t = 1) to data (t = 0) in ``steps`` Euler steps.
+    def velocity(sample, t, condition):
+        return model(sample, torch.tensor([1000 * t] * sample.shape[0], device=sample.device, dtype=torch.float32),
+                     condition)
 
-        rescale_t shifts the time grid towards t = 1 (more steps at high noise).
-        cfg_strength: guided velocity = (1 + w) * v(cond) - w * v(neg_cond).
-        callback(step_done, steps) is called after every step."""
-        sample = noise
-        t_seq = np.linspace(1, 0, steps + 1)
-        t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
-        for i in tqdm(range(steps), desc='Sampling', disable=not verbose):
-            t, t_prev = t_seq[i], t_seq[i + 1]
-            velocity = self._velocity(model, sample, t, cond)
-            if cfg_strength > 0:
-                negative_velocity = self._velocity(model, sample, t, neg_cond)
-                velocity = (1 + cfg_strength) * velocity - cfg_strength * negative_velocity
-            sample = sample - velocity * (t - t_prev)
-            if callback is not None:
-                callback(i + 1, steps)
-        return sample
+    sample = noise
+    t_seq = np.linspace(1, 0, steps + 1)
+    t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
+    for i in tqdm(range(steps), desc='Sampling', disable=not verbose):
+        t, t_prev = t_seq[i], t_seq[i + 1]
+        v = velocity(sample, t, cond)
+        if cfg_strength > 0:
+            v = (1 + cfg_strength) * v - cfg_strength * velocity(sample, t, neg_cond)
+        sample = sample - v * (t - t_prev)
+    return sample

idol/sampling_utils.py

@@ -1,150 +0,0 @@
-import numpy as np
-import torch
-from typing import *
-
-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: torch.Tensor,
-    counts: torch.Tensor,
-    algo: Literal["iid", "residual", "systematic"] = "systematic",
-    group_rows: bool = False,
-) -> torch.Tensor:
-    """
-    probs  : [..., P] nonnegative; rows need not be normalized
-    counts : [...] nonnegative ints
-    returns: [..., P] long tensor of counts per bin
-    handle properly for conunts = 0 case
-    """
-    batch_shape = counts.shape
-    B = counts.numel()
-    P = probs.size(-1)
-    assert B * P == probs.numel() and probs.dim() - counts.dim() == 1
-    device = probs.device
-    probs = probs.view(B, P)
-    counts = counts.view(B)
-
-    # Normalize with uniform fallback for zero-sum rows
-    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)
-
-    if algo == "iid":
-        # Pure multinomial IID, batched by unique k
-        unique_k, inv = counts.unique(sorted=False, return_inverse=True)
-        for i, k in enumerate(unique_k.tolist()):
-            if k == 0:
-                continue
-            rows = (inv == i).nonzero(as_tuple=False).squeeze(1)  # indices for this k
-            samples = torch.multinomial(
-                probs.index_select(0, rows), num_samples=int(k), replacement=True
-            )
-            buf = torch.zeros(rows.numel(), P, dtype=torch.float32, device=device)
-            buf.scatter_add_(1, samples, torch.ones_like(samples, dtype=buf.dtype))
-            out.index_copy_(0, rows, buf.to(torch.long))
-
-    elif algo in ("residual", "owen_scrambled_residual"):
-        # Step 1: deterministic floor allocation (vectorized)
-        expected = probs * counts.unsqueeze(1).to(torch.float32)  # [B, P]
-        base = torch.floor(expected)  # [B, P]
-        out += base.to(torch.long)
-
-        # Step 2: compute residuals
-        residual_mass = expected - base  # [B, P]
-        m = counts - base.sum(1).to(torch.long)  # [B]
-
-        # Guard: skip rows with no residuals
-        active_mask = m > 0
-        if active_mask.any():
-            rows = active_mask.nonzero(as_tuple=False).squeeze(1)
-            m_active = m[rows]
-            res_probs = residual_mass[rows]
-            res_probs /= res_probs.sum(1, keepdim=True)  # normalize residuals
-
-            if algo == "residual":
-                # --- Batched multinomial ---
-                unique_m, invm = m_active.unique(sorted=False, return_inverse=True)
-                for j, mm in enumerate(unique_m.tolist()):
-                    rows_m = (invm == j).nonzero(as_tuple=False).squeeze(1)  # indices for this m
-                    rsel = rows[rows_m]
-                    if mm == 0:
-                        continue
-                    samples = torch.multinomial(
-                        res_probs.index_select(0, rows_m),
-                        num_samples=int(mm),
-                        replacement=True,
-                    )
-                    buf = torch.zeros(rows_m.numel(), P, dtype=torch.float32, device=device)
-                    buf.scatter_add_(1, samples, torch.ones_like(samples, dtype=buf.dtype))
-                    out.index_copy_(0, rsel, out[rsel] + buf.to(torch.long))
-            else:
-                raise NotImplementedError("Owen-scrambled residual not implemented yet")
-        assert out.sum() == counts.sum(), f"sampling count mismatch, input probs max: {probs.max().item()}, counts: {counts.sum().item()}, output sum: {out.sum().item()}"
-
-    elif algo == "systematic":
-        cdf = probs.cumsum(dim=1).clamp(max=1.0 - 1e-12)
-        if group_rows:
-            # Stable grouping retains both ascending count order (torch.unique's
-            # existing order) and original row order within each group. Therefore
-            # torch.rand sees exactly the same sequence of call shapes and rows.
-            # One metadata transfer replaces one synchronizing nonzero per group.
-            sorted_counts, row_order = torch.sort(counts, stable=True)
-            unique_n, group_sizes = torch.unique_consecutive(sorted_counts, return_counts=True)
-            group_metadata = torch.stack((unique_n, group_sizes), dim=1).tolist()
-            row_ranges = []
-            row_offset = 0
-            for count, group_size in group_metadata:
-                row_ranges.append((count, row_order[row_offset:row_offset + group_size]))
-                row_offset += group_size
-        else:
-            unique_n, inv = counts.unique(sorted=False, return_inverse=True)
-            row_ranges = ((count, (inv == index).nonzero(as_tuple=False).squeeze(1))
-                          for index, count in enumerate(unique_n.tolist()) if count != 0)
-        for n, rows in row_ranges:
-            if n == 0:
-                continue
-            r = rows.numel()
-
-            # random offset per row in [0,1/n)
-            U0 = torch.rand(r, 1, device=device) / float(n)                 # [r,1]
-            # stratified points
-            grid = torch.arange(n, device=device, dtype=torch.float32)[None, :] / float(n)  # [1,n]
-            us = (U0 + grid).clamp(max=1.0 - 1e-12)                         # [r,n]
-
-            cdf_rows = cdf.index_select(0, rows)                             # [r,P]
-            idx = torch.searchsorted(cdf_rows, us).clamp_max(probs.size(1) - 1)                           # [r,n]
-
-            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))
-
-    else:
-        raise ValueError(f"Unknown algo {algo}")
-    
-    out = out.view(*batch_shape, P)
-    return out