TT-VidT TT3D encoder (structured resample), transformers remote code

README.md

@@ -0,0 +1,77 @@
+---
+license: apache-2.0
+library_name: transformers
+pipeline_tag: video-classification
+tags:
+  - video
+  - self-supervised
+  - motion
+---
+
+# TT-VidT (TT3D, structured resample)
+
+Video encoder from *TT-VidT: Decoupling the Temporal Axis for Efficient
+Motion-Centric Video Pretraining*. A DINOv3 ViT-B/16 appearance path is paired with
+a compact Temporal Transfer (TT3D) pathway that produces per-frame motion tokens.
+TT3D down/up-samples the spatial tokens with a fixed structured weight plus a
+trainable channel mix: 195.5M parameters instead of 407.8M for the full-linear
+resample of the paper checkpoint.
+
+Code, training recipes and evaluation: see the TT-VidT repository.
+
+## Usage
+
+With `transformers` only:
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("<this repo>", trust_remote_code=True).cuda().eval()
+video = torch.rand(1, 8, 3, 256, 256, device="cuda") * 2 - 1   # [B, T, C, H, W], pixels in [-1, 1]
+with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
+    motion = model(video).motion_output                          # [B, T, 1, 768]
+```
+
+With the TT-VidT repository:
+
+```python
+from ttvidt.hub import load_model
+
+model = load_model("<this repo>", device="cuda")
+motion = model.encoder(video).motion_output
+```
+
+Input: 8 frames at 256x256, normalised with mean = std = 0.5 (frames were sampled at
+6 fps during training).
+
+## Files
+
+| File | Content |
+|---|---|
+| `config.json` | architecture + `TTVidTrainer` arguments, and the `transformers` fields |
+| `model.safetensors` | fp32 encoder (195.5M parameters, incl. the 85.1M DINOv3 spatial path); no decoder |
+| `*.py` | encoder code for `trust_remote_code` (needs `torch` and `transformers` only) |
+
+Loading needs no access to the DINOv3 base weights.
+
+## Training
+
+Initialised from the paper's TT3D + Diff Compression encoder (full-linear resample)
+with the resample replaced, then trained for 30k steps to match that encoder's motion
+and conditioning embeddings (relative MSE) on the pretraining data
+(OpenVid-1M 384 px + Moments-in-Time v2), batch 16, AdamW 5e-4 with muP
+scaling, 200 warmup steps, cosine decay.
+
+## Results
+
+Frozen attentive probe, mean top-1 over 3 seeds (TT-VidT evaluation protocol, 8 frames):
+
+| Encoder | HMDB51 | ARID | IARD | Jester | SSv2 |
+|---|---:|---:|---:|---:|---:|
+| paper TT3D (full-linear resample, 407.8M) | 24.8 | 35.4 | 86.6 | 72.9 | 25.9 |
+| this checkpoint (structured resample, 195.5M) | 25.2 | 36.1 | 86.0 | 72.9 | 25.4 |
+
+## License
+
+Apache-2.0. The spatial path is initialised from DINOv3, which has its own license.

config.json

@@ -0,0 +1,171 @@
+{
+  "format": "ttvidt-release-v1",
+  "model": {
+    "base_model_name": null,
+    "motion_encoder_config": {
+      "temporal_patch_size": 1,
+      "temporal_patch_overlap": 0,
+      "motion_layers_period": 1,
+      "num_motion_tokens": 8,
+      "motion_ffn_type": "gelu",
+      "tt_mode": "3d",
+      "tt_downsample": 4
+    },
+    "motion_decoder_config": {
+      "image_dim": 4,
+      "patch_size": 2,
+      "decoder_style": "dit",
+      "num_layers": 12,
+      "hidden_size": 768,
+      "intermediate_size": 3072,
+      "num_heads": 12,
+      "decode_mode": "diffusion",
+      "qk_norm": true,
+      "attn_bias": false,
+      "use_final_norm": true,
+      "encoder_hidden_size": 768,
+      "temporal_patch_size": 1
+    },
+    "encoder_ae_mode": "video",
+    "decoder_ae_mode": "image",
+    "latent_mean": [
+      -0.69,
+      -0.48,
+      -0.6,
+      0.28
+    ],
+    "latent_std": [
+      12.38,
+      11.22,
+      7.93,
+      21.22
+    ],
+    "patch_size": null,
+    "image_dim": null,
+    "token_counts": [
+      1
+    ],
+    "gradient_checkpointing": false,
+    "train_mode": "regression",
+    "mae_mask_ratio": 0.75,
+    "ar_shift_range": null,
+    "ar_shift_sampling": "gamma",
+    "source_frame_count": 8,
+    "use_ref": false,
+    "unfreeze_backbone": true,
+    "decoder_mask_ratio": 0.0,
+    "mask_patch_size": 4,
+    "diffusion_timestep_sampling": "uniform",
+    "backbone_arch": "ttvidt",
+    "backbone_config": {
+      "transformers_version": "5.14.1",
+      "architectures": [
+        "DINOv3ViTModel"
+      ],
+      "output_hidden_states": false,
+      "return_dict": true,
+      "dtype": "float32",
+      "chunk_size_feed_forward": 0,
+      "is_encoder_decoder": false,
+      "id2label": {
+        "0": "LABEL_0",
+        "1": "LABEL_1"
+      },
+      "label2id": {
+        "LABEL_0": 0,
+        "LABEL_1": 1
+      },
+      "problem_type": null,
+      "patch_size": 16,
+      "hidden_size": 768,
+      "intermediate_size": 3072,
+      "num_hidden_layers": 12,
+      "num_attention_heads": 12,
+      "hidden_act": "gelu",
+      "attention_dropout": 0.0,
+      "initializer_range": 0.02,
+      "layer_norm_eps": 1e-05,
+      "rope_theta": 100.0,
+      "image_size": 224,
+      "num_channels": 3,
+      "query_bias": true,
+      "key_bias": false,
+      "value_bias": true,
+      "proj_bias": true,
+      "mlp_bias": true,
+      "layerscale_value": 1.0,
+      "drop_path_rate": 0.0,
+      "use_gated_mlp": false,
+      "num_register_tokens": 4,
+      "pos_embed_shift": null,
+      "pos_embed_jitter": null,
+      "pos_embed_rescale": 2.0,
+      "apply_layernorm": true,
+      "reshape_hidden_states": true,
+      "stage_names": [
+        "stem",
+        "stage1",
+        "stage2",
+        "stage3",
+        "stage4",
+        "stage5",
+        "stage6",
+        "stage7",
+        "stage8",
+        "stage9",
+        "stage10",
+        "stage11",
+        "stage12"
+      ],
+      "_name_or_path": "facebook/dinov3-vitb16-pretrain-lvd1689m",
+      "model_type": "dinov3_vit",
+      "temporal_patch_size": 1,
+      "temporal_patch_overlap": 0,
+      "motion_layers_period": 1,
+      "num_motion_tokens": 8,
+      "motion_ffn_type": "gelu",
+      "tt_mode": "3d",
+      "tt_downsample": 4,
+      "output_attentions": false,
+      "out_features": [
+        "stage12"
+      ],
+      "out_indices": [
+        12
+      ]
+    },
+    "decoder_lr_multiplier": 1.0,
+    "fused_adam": false,
+    "ema_foreach": false,
+    "nan_guard_interval": 1,
+    "name": "ttvidt-tt3d-diffcomp",
+    "log_interval": 2500,
+    "learning_rate": 0.0005,
+    "base_dim": 256,
+    "decoder_base_dim": null,
+    "weight_decay": 0.01,
+    "betas": [
+      0.9,
+      0.98
+    ],
+    "scheduler_config": {
+      "lr": {
+        "mode": "cosine",
+        "end": 436705,
+        "min_value": 0.01,
+        "warmup": 10000
+      }
+    }
+  },
+  "encoder_only": true,
+  "name": "ttvidt-tt3d-structured",
+  "source": "distilled from ttvidt-tt3d-diffcomp",
+  "model_type": "ttvidt",
+  "architectures": [
+    "TTVidTModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_ttvidt.TTVidTConfig",
+    "AutoModel": "modeling_ttvidt.TTVidTModel"
+  }
+}

dinov3_vit.py

@@ -0,0 +1,395 @@
+from dataclasses import dataclass
+from typing import Optional, Unpack
+
+import torch
+import torch.nn as nn
+from transformers import DINOv3ViTModel, DINOv3ViTConfig
+from transformers.utils import TransformersKwargs
+from transformers.modeling_outputs import ModelOutput
+from transformers.models.dinov3_vit.modeling_dinov3_vit import DINOv3ViTEmbeddings
+from .mup import mup_init, mup_init_output
+
+from .layers import AttentivePooling, RMSNorm
+from .patch import TemporalPatch
+from .tt import TemporalTransfer
+from .tt3d import TemporalTransfer3D
+
+
+@dataclass
+class VideoModelOutputWithMotionTokens(ModelOutput):
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    frames_output: Optional[torch.FloatTensor] = None
+    pooler_output: Optional[torch.FloatTensor] = None
+    motion_output: Optional[torch.FloatTensor] = None
+    cls_output: Optional[torch.FloatTensor] = None
+    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[tuple[torch.FloatTensor, ...]] = None
+
+
+class DINOv3VidTConfig(DINOv3ViTConfig):
+    model_type = "dinov3_vit"
+
+    def __init__(
+        self,
+        patch_size: int = 16,
+        temporal_patch_size: int = 1,
+        temporal_patch_overlap: int = 0,
+        hidden_size: int = 384,
+        intermediate_size: int = 1536,
+        num_hidden_layers: int = 12,
+        motion_layers_period: int = 0,
+        num_attention_heads: int = 6,
+        hidden_act: str = "gelu",
+        attention_dropout: float = 0.0,
+        initializer_range: float = 0.02,
+        layer_norm_eps: float = 1e-5,
+        rope_theta: float = 100.0,
+        image_size: int = 224,
+        num_channels: int = 3,
+        query_bias: bool = True,
+        key_bias: bool = False,
+        value_bias: bool = True,
+        proj_bias: bool = True,
+        mlp_bias: bool = True,
+        layerscale_value: float = 1.0,
+        drop_path_rate: float = 0.0,
+        use_gated_mlp: bool = False,
+        num_register_tokens: int = 0,
+        num_motion_tokens: int = 0,
+        motion_ffn_type: str = "swiglu",  # "swiglu" or "gelu" for TemporalTransfer
+        tt_mode: str = "1d",  # "1d" (motion-only) or "3d" (motion + downsampled spatial)
+        tt_downsample: int = 4,  # spatial downsample factor for tt_mode="3d"
+        # train augs
+        pos_embed_shift: Optional[float] = None,
+        pos_embed_jitter: Optional[float] = None,
+        pos_embed_rescale: Optional[float] = 2.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.temporal_patch_overlap = temporal_patch_overlap
+        self.num_channels = num_channels
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.motion_layers_period = motion_layers_period
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.attention_dropout = attention_dropout
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.layerscale_value = layerscale_value
+        self.drop_path_rate = drop_path_rate
+        self.use_gated_mlp = use_gated_mlp
+        self.rope_theta = rope_theta
+        self.query_bias = query_bias
+        self.key_bias = key_bias
+        self.value_bias = value_bias
+        self.proj_bias = proj_bias
+        self.mlp_bias = mlp_bias
+        self.num_register_tokens = num_register_tokens
+        self.num_motion_tokens = num_motion_tokens
+        self.motion_ffn_type = motion_ffn_type
+        self.tt_mode = tt_mode
+        self.tt_downsample = tt_downsample
+
+        # train augs
+        self.pos_embed_shift = pos_embed_shift
+        self.pos_embed_jitter = pos_embed_jitter
+        self.pos_embed_rescale = pos_embed_rescale
+
+
+class DINOv3VidTModel(DINOv3ViTModel):
+    config: DINOv3VidTConfig
+
+    def __init__(self, config: DINOv3VidTConfig):
+        super().__init__(config)
+        self.num_motion_tokens = config.num_motion_tokens
+        if self.num_motion_tokens > 0:
+            self.motion_tokens = nn.Parameter(
+                torch.randn(config.num_motion_tokens, config.hidden_size)
+                / (config.hidden_size * config.num_motion_tokens) ** 0.5
+            )
+        else:
+            self.motion_tokens = None
+        if config.temporal_patch_size > 1:
+            self.temporal_embeddings = TemporalPatch(
+                config.hidden_size,
+                config.hidden_size,
+                config.temporal_patch_size,
+                config.temporal_patch_overlap,
+            )
+        else:
+            self.temporal_embeddings = nn.Identity()
+        # Attentive pooling on motion tokens -> CLS-like conditioning token
+        if self.num_motion_tokens > 0:
+            self.motion_pooler = AttentivePooling(
+                config.hidden_size, num_heads=config.num_attention_heads
+            )
+            # Pool all TT tokens → 1 motion token for decoder output
+            self.motion_output_pooler = AttentivePooling(
+                config.hidden_size, num_heads=config.num_attention_heads
+            )
+            # Normalize cls_output (AttentivePooling has no output norm; without this
+            # cls_output can reach absmax~50 and destabilize AdaRMSNorm conditioning)
+            self.cls_norm = RMSNorm(config.hidden_size)
+        else:
+            self.motion_pooler = None
+            self.motion_output_pooler = None
+            self.cls_norm = None
+
+        self.motion_layers = nn.ModuleList([])
+        self.tt_mode = config.tt_mode
+
+        layer_added = 0
+        for idx, layer in enumerate(self.model.layer):
+            if (
+                config.motion_layers_period > 0
+                and idx % config.motion_layers_period == 0
+            ):
+                layer_added += 1
+                if config.tt_mode == "3d":
+                    self.motion_layers.append(
+                        TemporalTransfer3D(
+                            config.hidden_size,
+                            config.intermediate_size,
+                            config.num_attention_heads,
+                            downsample_factor=config.tt_downsample,
+                            ffn_type=config.motion_ffn_type,
+                        )
+                    )
+                else:
+                    self.motion_layers.append(
+                        TemporalTransfer(
+                            config.hidden_size,
+                            config.intermediate_size,
+                            config.num_attention_heads,
+                            ffn_type=config.motion_ffn_type,
+                        )
+                    )
+            else:
+                self.motion_layers.append(nn.Identity())
+
+        mup_init(self.motion_layers.parameters())
+        # Restore special inits clobbered by mup_init
+        for layer in self.motion_layers:
+            if isinstance(layer, nn.Identity):
+                continue
+            # Residual output projections: near-zero init
+            if hasattr(layer, 'attn') and hasattr(layer.attn, 'out_proj'):
+                mup_init_output(layer.attn.out_proj.weight)
+            if hasattr(layer, 'out_proj'):  # TT3D has out_proj on the module
+                mup_init_output(layer.out_proj.weight)
+            if hasattr(layer.mlp, 'fc2'):
+                mup_init_output(layer.mlp.fc2.weight)
+            elif hasattr(layer.mlp, 'down'):
+                mup_init_output(layer.mlp.down.weight)
+            # TT3D-specific: zero-init spatial_out_proj, identity resample channel mix, qk_scale
+            if hasattr(layer, 'spatial_out_proj'):
+                nn.init.zeros_(layer.spatial_out_proj.weight)
+                layer.spatial_down.reset_parameters()
+                layer.spatial_up.reset_parameters()
+            for module in layer.modules():
+                if hasattr(module, 'qk_scale'):
+                    module.qk_scale.data.fill_(10.0)
+        print(f"Motion layers added: {layer_added} / {config.num_hidden_layers} (mode={config.tt_mode})")
+
+    @torch.no_grad()
+    def _init_weights(self, module) -> None:
+        super()._init_weights(module)
+        # transformers' from_pretrained builds on the meta device and leaves
+        # non-persistent buffers empty: recompute those of the TT modules
+        if hasattr(module, "reset_buffers"):
+            module.reset_buffers()
+
+    def freeze_pretrained(self, requires_grad: bool = False):
+        """Freeze or unfreeze the pretrained backbone components.
+
+        Pretrained components: embeddings, rope_embeddings, transformer layers, norm.
+        Leaves trainable-from-scratch parts (motion_tokens, motion_layers,
+        temporal_embeddings) untouched.
+        """
+        self.embeddings.requires_grad_(requires_grad)
+        self.rope_embeddings.requires_grad_(requires_grad)
+        self.model.layer.requires_grad_(requires_grad)
+        self.norm.requires_grad_(requires_grad)
+
+    def setup_patch(self, num_channels, patch_size):
+        self.config.num_channels = num_channels
+        self.config.patch_size = patch_size
+        self.embeddings = DINOv3ViTEmbeddings(self.config)
+        mup_init(self.embeddings.parameters())
+
+    @classmethod
+    def from_dino_v3(
+        cls,
+        dino_v3: DINOv3ViTModel,
+        temporal_patch_size: int = 0,
+        temporal_patch_overlap: int = 0,
+        motion_layers_period: int = 0,
+        num_motion_tokens: int = 0,
+        motion_ffn_type: str = "swiglu",
+        tt_mode: str = "1d",
+        tt_downsample: int = 4,
+    ):
+        dino_v3_config = dino_v3.config
+        ttvidt_config = DINOv3VidTConfig(
+            **dino_v3_config.to_dict(),
+            temporal_patch_size=temporal_patch_size,
+            temporal_patch_overlap=temporal_patch_overlap,
+            motion_layers_period=motion_layers_period,
+            num_motion_tokens=num_motion_tokens,
+            motion_ffn_type=motion_ffn_type,
+            tt_mode=tt_mode,
+            tt_downsample=tt_downsample,
+        )
+        new_model = cls(ttvidt_config)
+        new_model.load_state_dict(dino_v3.state_dict(), strict=False)
+        return new_model
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        bool_masked_pos: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        spatial_mask: Optional[torch.Tensor] = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> VideoModelOutputWithMotionTokens:
+        r"""
+        pixel_values: [B, T, C, H, W]
+        spatial_mask: optional bool tensor [L], True = visible/keep (for MAE tube masking)
+        """
+        b, t = pixel_values.shape[:2]
+
+        _, _, _, H_in, W_in = pixel_values.shape
+        ps = self.config.patch_size
+        spatial_h = H_in // ps
+        spatial_w = W_in // ps
+
+        pixel_values = pixel_values.to(
+            self.embeddings.patch_embeddings.weight.dtype
+        ).flatten(0, 1)
+        hidden_states = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
+        position_embeddings = self.rope_embeddings(pixel_values)
+        pixel_length = position_embeddings[0].size(0)
+
+        hidden_states = self.temporal_embeddings(hidden_states.unflatten(0, (b, -1)))
+
+        # MAE tube masking: keep only visible spatial tokens
+        if spatial_mask is not None:
+            # hidden_states may have prefix tokens (CLS + registers) before spatial tokens
+            # spatial_mask: bool [L] where L = num spatial patches
+            n_prefix = hidden_states.shape[2] - spatial_mask.shape[0]
+            if n_prefix > 0:
+                prefix = hidden_states[:, :, :n_prefix]
+                spatial = hidden_states[:, :, n_prefix:]
+                hidden_states = torch.cat([prefix, spatial[:, :, spatial_mask]], dim=2)
+            else:
+                hidden_states = hidden_states[:, :, spatial_mask]
+            # Subsample RoPE position embeddings: (cos[L, head_dim], sin[L, head_dim])
+            cos, sin = position_embeddings
+            position_embeddings = (cos[spatial_mask], sin[spatial_mask])
+            pixel_length = spatial_mask.sum().item()
+        b, ot, seq_len, dim = hidden_states.shape
+        if self.motion_tokens is not None:
+            hidden_states = torch.concat(
+                [self.motion_tokens[None, None].repeat(b, ot, 1, 1), hidden_states],
+                dim=2,
+            )
+
+        hidden_states = hidden_states.flatten(0, 1)
+        for i, (layer_module, motion_module) in enumerate(
+            zip(self.model.layer, self.motion_layers)
+        ):
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            hidden_states = layer_module(
+                hidden_states,
+                attention_mask=layer_head_mask,
+                position_embeddings=position_embeddings,
+            )
+            if not isinstance(motion_module, nn.Identity):
+                hidden_states = hidden_states.unflatten(0, (b, ot))
+                motion_tokens = hidden_states[:, :, : self.num_motion_tokens]
+                vit_state = hidden_states[:, :, self.num_motion_tokens :]
+
+                if isinstance(motion_module, TemporalTransfer3D):
+                    # TT3D: pass spatial tokens (last pixel_length of vit_state)
+                    # vit_state layout: [CLS, REG_0..REG_3, S_0..S_L]
+                    spatial_tokens = vit_state[:, :, -pixel_length:]
+
+                    if spatial_mask is not None:
+                        # MAE mode: pad visible tokens back to full spatial grid
+                        # for pixel unshuffle/shuffle in TT3D, then extract visible
+                        L_full = spatial_h * spatial_w
+                        full_spatial = spatial_tokens.new_zeros(
+                            b, ot, L_full, dim
+                        )
+                        full_spatial[:, :, spatial_mask] = spatial_tokens
+                        motion_tokens, full_spatial_out = motion_module(
+                            motion_tokens, full_spatial, spatial_h, spatial_w
+                        )
+                        spatial_out = full_spatial_out[:, :, spatial_mask]
+                    else:
+                        motion_tokens, spatial_out = motion_module(
+                            motion_tokens, spatial_tokens, spatial_h, spatial_w
+                        )
+
+                    # Write spatial back (only last pixel_length of vit_state)
+                    vit_state = torch.cat(
+                        [vit_state[:, :, :-pixel_length], spatial_out], dim=2
+                    )
+                else:
+                    # TT1D: motion tokens only
+                    motion_tokens = motion_module(motion_tokens)
+
+                hidden_states = torch.concat([motion_tokens, vit_state], dim=2)
+                hidden_states = hidden_states.flatten(0, 1)
+
+        sequence_output = self.norm(hidden_states).unflatten(0, (b, ot))
+        frames_output = sequence_output[:, :, -pixel_length:, :]
+        pooled_output = sequence_output[:, :, self.num_motion_tokens, :]
+        if self.motion_tokens is None:
+            motion_output = None
+            cls_output = None
+        else:
+            all_tt_tokens = sequence_output[
+                :, :, : self.num_motion_tokens, :
+            ]  # [B, T, num_motion_tokens, dim]
+            cls_output = self.cls_norm(self.motion_pooler(all_tt_tokens))  # [B, T, dim]
+            # Pool all TT tokens → 1 motion token for decoder
+            motion_output = self.motion_output_pooler(all_tt_tokens).unsqueeze(2)  # [B, T, 1, dim]
+
+        return VideoModelOutputWithMotionTokens(
+            last_hidden_state=sequence_output,
+            frames_output=frames_output,
+            pooler_output=pooled_output,
+            motion_output=motion_output,
+            cls_output=cls_output,
+        )
+
+
+if __name__ == "__main__":
+    dino = DINOv3ViTModel.from_pretrained(
+        "facebook/dinov3-vitb16-pretrain-lvd1689m"
+    ).cuda()
+    dino_vidt = DINOv3VidTModel.from_dino_v3(
+        dino,
+        temporal_patch_size=4,
+        temporal_patch_overlap=1,
+        motion_layers_period=1,
+        num_motion_tokens=16,
+    ).cuda()
+
+    print(dino)
+    print(dino_vidt)
+
+    test_x = torch.randn(1, 61, 3, 256, 256).cuda()
+    with torch.autocast("cuda"):
+        result = dino_vidt(test_x)
+
+    print(result.frames_output.shape)
+    print(result.pooler_output.shape)
+    print(result.motion_output.shape)

env.py

@@ -0,0 +1,3 @@
+import os
+TORCH_COMPILE = os.environ.get("TORCH_COMPILE", "1") == "1"
+COMPILE_KWARG = {"mode": "default", "dynamic": False}

layers.py

@@ -0,0 +1,341 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import compile_wrapper
+
+
+class RMSNorm(nn.RMSNorm):
+    @compile_wrapper
+    def forward(self, x):
+        return F.rms_norm(
+            x.to(self.weight), self.normalized_shape, self.weight, self.eps
+        )
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super(SwiGLU, self).__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.up = nn.Linear(hidden_size, intermediate_size * 2)
+        self.down = nn.Linear(intermediate_size, hidden_size)
+
+    @compile_wrapper
+    def forward(self, x):
+        x, gate = self.up(x).chunk(2, dim=-1)
+        x = F.silu(x) * gate
+        x = self.down(x)
+        return x
+
+
+class GELUMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.fc1 = nn.Linear(hidden_size, intermediate_size)
+        self.act = nn.GELU()
+        self.fc2 = nn.Linear(intermediate_size, hidden_size)
+
+    @compile_wrapper
+    def forward(self, x):
+        return self.fc2(self.act(self.fc1(x)))
+
+
+def _qk_norm(q, k, scale, eps=1e-6):
+    """Cosine-similarity attention with learnable temperature (DisMo-style).
+
+    Logits = scale * cos(q, k).
+    Computes in float32 for numerical stability, uses rsqrt(sum_sq + eps).
+    Scale is clamped to min=eps to prevent NaN from sqrt of negative values.
+    """
+    dtype = q.dtype
+    q, k, scale = q.float(), k.float(), scale.float()
+    sum_sq_q = torch.sum(q ** 2, dim=-1, keepdim=True)
+    sum_sq_k = torch.sum(k ** 2, dim=-1, keepdim=True)
+    sqrt_scale = torch.sqrt(scale.clamp(min=eps))
+    scale_q = sqrt_scale * torch.rsqrt(sum_sq_q + eps)
+    scale_k = sqrt_scale * torch.rsqrt(sum_sq_k + eps)
+    return (q * scale_q).to(dtype), (k * scale_k).to(dtype)
+
+
+class Attention(nn.Module):
+    def __init__(self, hidden_size, num_heads, qk_norm=False, bias=True):
+        super(Attention, self).__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.qk_norm = qk_norm
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        if qk_norm:
+            self.qk_scale = nn.Parameter(torch.full([num_heads, 1, 1], 10.0))
+
+    @compile_wrapper
+    def forward(self, x, rope=None, positions=None):
+        """
+        Args:
+            x: [B, T, SL, D] or [B, SL, D]
+            rope: optional RoPE3D module (from tt3d.py)
+            positions: optional [SL, 3] positions for 3D RoPE
+        """
+        # [B, T, SL, D]
+        *b, seq_len, dim = x.shape
+        x = x.flatten(1, 2)
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+        q = q.view(-1, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(-1, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(-1, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Optional 3D RoPE
+        if rope is not None and positions is not None:
+            q, k = rope(q, k, positions)
+
+        if self.qk_norm:
+            q, k = _qk_norm(q, k, self.qk_scale)
+            attn = F.scaled_dot_product_attention(q, k, v, scale=1.0)
+        else:
+            attn = F.scaled_dot_product_attention(q, k, v)
+
+        x = self.out_proj(attn.transpose(1, 2).flatten(2, 3))
+        return x.reshape(*b, seq_len, dim)
+
+
+class TransformerBlock(nn.Module):
+    """Standard pre-norm transformer block: RMSNorm + Attention + FFN (SwiGLU or GELU)."""
+
+    def __init__(self, hidden_size: int, intermediate_size: int, num_heads: int,
+                 ffn_type: str = "swiglu", qk_norm: bool = False):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size)
+        self.attn = Attention(hidden_size, num_heads, qk_norm=qk_norm)
+        self.norm2 = RMSNorm(hidden_size)
+        if ffn_type == "gelu":
+            self.mlp = GELUMLP(hidden_size, intermediate_size)
+        else:
+            self.mlp = SwiGLU(hidden_size, intermediate_size)
+
+    def forward(self, x: torch.Tensor, rope=None, positions=None) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x), rope=rope, positions=positions)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class AdaRMSNorm(nn.Module):
+    """Adaptive RMSNorm: scale-only modulation from conditioning signal (DisMo-style)."""
+
+    def __init__(self, hidden_size: int, cond_size: int):
+        super().__init__()
+        self.norm = nn.RMSNorm(hidden_size)
+        self.linear = nn.Linear(cond_size, hidden_size, bias=False)
+        nn.init.zeros_(self.linear.weight)
+
+    @compile_wrapper
+    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [..., D] input tensor
+            c: [..., C] conditioning tensor (broadcastable to x batch dims)
+        Returns:
+            [..., D] modulated tensor
+        """
+        return (self.linear(F.silu(c)) + 1) * self.norm(x.to(self.norm.weight))
+
+
+class CrossAttention(nn.Module):
+    """Cross-attention: Q from main stream, K/V from context."""
+
+    def __init__(self, hidden_size: int, num_heads: int, context_size: int | None = None, qk_norm: bool = False, bias: bool = True):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.qk_norm = qk_norm
+        context_size = context_size or hidden_size
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.k_proj = nn.Linear(context_size, hidden_size, bias=bias)
+        self.v_proj = nn.Linear(context_size, hidden_size, bias=bias)
+        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        if qk_norm:
+            self.qk_scale = nn.Parameter(torch.full([num_heads, 1, 1], 10.0))
+
+    @compile_wrapper
+    def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, S, D] main stream (query)
+            context: [B, T, Sc, Dc] context (key/value)
+        Returns:
+            [B, T, S, D]
+        """
+        bt = x.shape[0] * x.shape[1]
+        seq_q = x.shape[2]
+        seq_kv = context.shape[2]
+
+        x_flat = x.reshape(bt, seq_q, self.hidden_size)
+        ctx_flat = context.reshape(bt, seq_kv, -1)
+
+        q = self.q_proj(x_flat).view(bt, seq_q, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(ctx_flat).view(bt, seq_kv, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(ctx_flat).view(bt, seq_kv, self.num_heads, self.head_dim).transpose(1, 2)
+
+        if self.qk_norm:
+            q, k = _qk_norm(q, k, self.qk_scale)
+            attn = F.scaled_dot_product_attention(q, k, v, scale=1.0)
+        else:
+            attn = F.scaled_dot_product_attention(q, k, v)
+        out = self.out_proj(attn.transpose(1, 2).flatten(2, 3))
+        return out.view(x.shape)
+
+
+class AttentivePooling(nn.Module):
+    """Learnable query token attends to M tokens to produce a CLS-like token."""
+
+    def __init__(self, hidden_size: int, num_heads: int = 1):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.query = nn.Parameter(torch.randn(1, 1, 1, hidden_size) * 0.02)
+        self.k_proj = nn.Linear(hidden_size, hidden_size)
+        self.v_proj = nn.Linear(hidden_size, hidden_size)
+        self.out_proj = nn.Linear(hidden_size, hidden_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, M, D] motion tokens
+        Returns:
+            [B, T, D] pooled CLS-like token
+        """
+        B, T, M, D = x.shape
+        bt = B * T
+
+        q = self.query.expand(B, T, -1, -1).reshape(bt, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        x_flat = x.reshape(bt, M, D)
+        k = self.k_proj(x_flat).view(bt, M, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x_flat).view(bt, M, self.num_heads, self.head_dim).transpose(1, 2)
+
+        attn = F.scaled_dot_product_attention(q, k, v)  # [bt, H, 1, head_dim]
+        out = self.out_proj(attn.transpose(1, 2).flatten(2, 3))  # [bt, 1, D]
+        return out.squeeze(2).view(B, T, D)
+
+
+class DiTTransformerBlock(nn.Module):
+    """DiT-style transformer block with AdaRMSNorm conditioning and optional cross-attention."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        num_heads: int,
+        cond_size: int,
+        use_cross_attn: bool = False,
+        context_size: int | None = None,
+        use_rope: bool = False,
+        qk_norm: bool = False,
+        bias: bool = True,
+    ):
+        super().__init__()
+        from .pos_embed import RoPE2DAttention
+
+        # Self-attention sub-layer (with optional 2D RoPE)
+        self.ada_norm1 = AdaRMSNorm(hidden_size, cond_size)
+        self.use_rope = use_rope
+        if use_rope:
+            self.attn = RoPE2DAttention(hidden_size, num_heads, qk_norm=qk_norm, bias=bias)
+        else:
+            self.attn = Attention(hidden_size, num_heads, qk_norm=qk_norm, bias=bias)
+
+        # Optional cross-attention sub-layer
+        self.use_cross_attn = use_cross_attn
+        if use_cross_attn:
+            self.ada_norm_cross = AdaRMSNorm(hidden_size, cond_size)
+            self.cross_attn = CrossAttention(hidden_size, num_heads, context_size, qk_norm=qk_norm, bias=bias)
+
+        # FFN sub-layer
+        self.ada_norm2 = AdaRMSNorm(hidden_size, cond_size)
+        self.mlp = SwiGLU(hidden_size, intermediate_size)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        c: torch.Tensor,
+        context: torch.Tensor | None = None,
+        h: int | None = None,
+        w: int | None = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, S, D] main stream
+            c: [B, T, 1, Dc] conditioning (broadcast over S)
+            context: [B, T, Sc, Dc_ctx] optional cross-attn context
+            h, w: spatial dims (required when use_rope=True)
+        Returns:
+            [B, T, S, D]
+        """
+        if self.use_rope:
+            x = x + self.attn(self.ada_norm1(x, c), h, w)
+        else:
+            x = x + self.attn(self.ada_norm1(x, c))
+        if self.use_cross_attn and context is not None:
+            x = x + self.cross_attn(self.ada_norm_cross(x, c), context)
+        x = x + self.mlp(self.ada_norm2(x, c))
+        return x
+
+
+class LatentPatchEmbed(nn.Module):
+    """
+    Patch embedding for latent-space video input.
+
+    Expects input of shape [B, T, C_in, H, W] where the spatial dimensions are
+    already at patch resolution (e.g. from a VAE encoder).  If the input channel
+    dimension differs from hidden_size, a linear projection is applied per-patch.
+
+    When spatial_patch_size > 1 the spatial grid is further down-sampled via a
+    Conv2d with kernel_size = stride = spatial_patch_size.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_size: int,
+        spatial_patch_size: int = 1,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_size = hidden_size
+        self.spatial_patch_size = spatial_patch_size
+
+        if spatial_patch_size > 1:
+            self.proj = nn.Conv2d(
+                in_channels,
+                hidden_size,
+                kernel_size=spatial_patch_size,
+                stride=spatial_patch_size,
+            )
+        else:
+            self.proj = nn.Linear(in_channels, hidden_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, C, H, W]
+        Returns:
+            [B, T, H'*W', D]  where H' = H // spatial_patch_size, etc.
+        """
+        b, t, c, h, w = x.shape
+        if self.spatial_patch_size > 1:
+            x = x.reshape(b * t, c, h, w)
+            x = self.proj(x)  # [B*T, D, H', W']
+            _, d, hp, wp = x.shape
+            x = x.flatten(2).transpose(1, 2)  # [B*T, H'*W', D]
+            x = x.reshape(b, t, hp * wp, d)
+        else:
+            # x: [B, T, C, H, W] -> [B, T, H*W, C] -> Linear -> [B, T, H*W, D]
+            x = x.permute(0, 1, 3, 4, 2).reshape(b, t, h * w, c)
+            x = self.proj(x)
+        return x

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03f7f365511775818645e02e4b4d56758a981d2c5c9122c4758fcb22fccbfce5
+size 781976200

modeling_ttvidt.py

@@ -0,0 +1,62 @@
+"""Hugging Face ``transformers`` wrapper around the TT-VidT encoder.
+
+``scripts/tools/export_hf.py`` copies this file (as ``modeling_ttvidt.py``) and the
+encoder sources next to an exported model, so the same directory / Hub repo loads with
+
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained("<repo>", trust_remote_code=True)
+    motion = model(video).motion_output      # video: [B, T, 3, 256, 256] in [-1, 1]
+
+and with ``ttvidt.hub.load_model("<repo>")``: ``config.json`` holds the release format
+(``format``, ``model`` = TTVidTrainer arguments) plus the ``transformers`` fields, and
+``model.safetensors`` uses the TTVidTrainer key names (``encoder.*``).
+"""
+
+from transformers import PretrainedConfig, PreTrainedModel
+
+from .dinov3_vit import DINOv3VidTConfig, DINOv3VidTModel
+
+
+class TTVidTConfig(PretrainedConfig):
+    """``model``: the TTVidTrainer keyword arguments of the release (``backbone_config``
+    holds the full DINOv3VidTConfig, so no DINOv3 download is needed)."""
+
+    model_type = "ttvidt"
+
+    def __init__(self, format: str | None = None, model: dict | None = None,
+                 encoder_only: bool = False, **kwargs):
+        self.format = format
+        self.model = model or {}
+        self.encoder_only = encoder_only
+        super().__init__(**kwargs)
+
+    def encoder_config(self) -> DINOv3VidTConfig:
+        if self.model.get("backbone_arch", "ttvidt") != "ttvidt":
+            raise ValueError(f"only TT-VidT encoders are supported, got {self.model.get('backbone_arch')!r}")
+        cfg = dict(self.model["backbone_config"])
+        for k, v in (self.model.get("motion_encoder_config") or {}).items():
+            cfg.setdefault(k, v)  # same merge as TTVidTrainer
+        return DINOv3VidTConfig(**cfg)
+
+
+class TTVidTModel(PreTrainedModel):
+    """The TT-VidT video encoder (``DINOv3VidTModel``); ``forward`` returns its output
+    (``motion_output``: [B, T, K, D] motion tokens per frame)."""
+
+    config_class = TTVidTConfig
+    base_model_prefix = "ttvidt"
+    main_input_name = "pixel_values"
+    # the decoder and the frame-VAE latent statistics of full releases are not used
+    _keys_to_ignore_on_load_unexpected = [r"^decoder\.", r"^latent_(mean|std)$"]
+
+    def __init__(self, config: TTVidTConfig):
+        super().__init__(config)
+        self.encoder = DINOv3VidTModel(config.encoder_config())
+        self.post_init()
+
+    def _init_weights(self, module):
+        pass  # DINOv3VidTModel initialises itself; pretrained weights come from the file
+
+    def forward(self, pixel_values, **kwargs):
+        """pixel_values: [B, T, 3, H, W] in [-1, 1]."""
+        return self.encoder(pixel_values, **kwargs)

mup.py

@@ -0,0 +1,18 @@
+"""muP initialisation (same as ``optimfactory.mup_init`` / ``mup_init_output``)."""
+
+import math
+
+import torch
+
+
+def mup_init(params, is_output: bool = False) -> None:
+    for param in params:
+        if param.ndim == 1:
+            continue
+        fan_in = math.prod(param.shape[1:])
+        std = (1 / fan_in) ** (1 if is_output else 0.5)
+        torch.nn.init.normal_(param, mean=0.0, std=std)
+
+
+def mup_init_output(param: torch.Tensor) -> None:
+    mup_init([param], is_output=True)

patch.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import compile_wrapper
+
+
+class TemporalPatch(nn.Module):
+    def __init__(
+        self,
+        input_channels,
+        output_channels,
+        patch_size=2,
+        overlap=1,
+        init_gamma=1.1,
+    ):
+        super(TemporalPatch, self).__init__()
+        self.patch_size = patch_size
+        self.overlap = overlap
+        self.input_channels = input_channels
+        self.output_channels = output_channels
+        if input_channels != output_channels:
+            self.fproj = nn.Conv1d(input_channels, output_channels, 1, 1, 0)
+        else:
+            self.fproj = nn.Identity()
+        self.tproj = nn.Conv1d(
+            output_channels, output_channels, patch_size + overlap, stride=patch_size
+        )
+        self.init_gamma = init_gamma
+        self.init_weight()
+
+    def init_weight(self):
+        if isinstance(self.fproj, nn.Conv1d):
+            nn.init.constant_(self.fproj.bias, 0)
+            nn.init.normal_(self.fproj.weight, std=1 / self.input_channels**0.5)
+
+        nn.init.constant_(self.tproj.bias, 0)
+        tproj_patch = self.tproj.weight.size(-1)
+        I = torch.eye(self.tproj.weight.size(0)).to(self.tproj.weight)[..., None]
+        weight = I.repeat(1, 1, tproj_patch) * self.init_gamma
+        scale = [self.init_gamma**i for i in range(tproj_patch)]
+        scale = torch.tensor(scale).to(self.tproj.weight) / sum(scale)
+        self.tproj.weight.data.copy_(weight * scale[None, None, :])
+
+    @compile_wrapper
+    def forward(self, x):
+        """
+        x: [B, T, L, C]
+        """
+        b, t, l, c = x.shape
+        x = x.permute(0, 2, 3, 1).contiguous().view(b * l, c, t)
+
+        x = self.fproj(x)
+        first_frame = x[..., 0:1]
+        x = self.tproj(x)
+        x = torch.concat([first_frame, x], dim=-1)
+
+        _, oc, ot = x.shape
+        x = x.view(b, l, oc, ot).permute(0, 3, 1, 2).contiguous()
+        return x
+
+
+if __name__ == "__main__":
+    x = torch.randn(2, 1 + 15 * 8, 3, 224, 224)
+    m = TemporalPatch(3, 4, 8, 1)
+    y = m(x)  # [B, 1 + (T-1)//p, 4, H, W]
+    print(y.shape)

pos_embed.py

@@ -0,0 +1,358 @@
+"""
+Positional embedding modules for TT-VidT.
+
+Includes:
+- Sinusoidal 1D/2D position embeddings
+- Adaptive RoPE (Rotary Position Embedding) with aspect-ratio-aware positions
+- Temporal distance embeddings
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import compile_wrapper
+
+# ============================================================================
+# Sinusoidal Position Embeddings
+# ============================================================================
+
+
+def get_1d_sincos_pos_embed(
+    embed_dim: int, pos: torch.Tensor, max_period: int = 10000
+) -> torch.Tensor:
+    """
+    Generate 1D sinusoidal positional embeddings.
+
+    Args:
+        embed_dim: Embedding dimension
+        pos: Position tensor [N]
+        max_period: Maximum period for frequency bands
+
+    Returns:
+        Sinusoidal embeddings [N, embed_dim]
+    """
+    half_dim = embed_dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(half_dim, dtype=torch.float32, device=pos.device)
+        / half_dim
+    )
+    args = pos.unsqueeze(-1) * freqs  # [N, D/2]
+    return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)  # [N, D]
+
+
+def get_2d_sincos_pos_embed(
+    embed_dim: int, h: int, w: int, max_period: int = 10000
+) -> torch.Tensor:
+    """
+    Generate 2D sinusoidal positional embeddings.
+
+    Args:
+        embed_dim: Embedding dimension
+        h, w: Grid height and width
+        max_period: Maximum period for frequency bands
+
+    Returns:
+        Sinusoidal embeddings [H*W, embed_dim]
+    """
+    grid_h = torch.arange(h, dtype=torch.float32)
+    grid_w = torch.arange(w, dtype=torch.float32)
+    grid = torch.stack(
+        torch.meshgrid(grid_h, grid_w, indexing="ij"), dim=-1
+    )  # [H, W, 2]
+    grid = grid.reshape(-1, 2)  # [H*W, 2]
+
+    # Split embedding dim for h and w
+    half_dim = embed_dim // 2
+    emb_h = get_1d_sincos_pos_embed(half_dim, grid[:, 0], max_period)
+    emb_w = get_1d_sincos_pos_embed(half_dim, grid[:, 1], max_period)
+    return torch.cat([emb_h, emb_w], dim=-1)  # [H*W, D]
+
+
+def get_adaptive_2d_pos(
+    h: int, w: int, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generate adaptive 2D positions with aspect-ratio-aware normalization.
+
+    Position mapping is normalized such that:
+    - max(H) * max(W) = 1 (positions are in normalized range)
+    - max(H) / max(W) = h / w (preserves aspect ratio)
+
+    This means: max_H = sqrt(h/w), max_W = sqrt(w/h)
+
+    Args:
+        h, w: Grid dimensions
+        device: Torch device
+
+    Returns:
+        (grid_h, grid_w): Flattened position grids, each [H*W]
+    """
+    aspect = h / w
+    max_H = math.sqrt(aspect)
+    max_W = math.sqrt(1.0 / aspect)
+
+    pos_h = torch.linspace(0, max_H, h, device=device, dtype=torch.float32)
+    pos_w = torch.linspace(0, max_W, w, device=device, dtype=torch.float32)
+
+    grid_h, grid_w = torch.meshgrid(pos_h, pos_w, indexing="ij")
+    return grid_h.reshape(-1), grid_w.reshape(-1)
+
+
+# ============================================================================
+# Temporal Distance Embedding
+# ============================================================================
+
+
+class TemporalDistanceEmbedding(nn.Module):
+    """
+    Temporal distance embedding as a concatenated token using sincos encoding.
+
+    Instead of adding temporal pos encoding, we create a dedicated token
+    that encodes the distance from source frame to target motion frame.
+    This token is concatenated: [frame_emb, motion_emb, time_info]
+
+    Uses sinusoidal positional encoding (like standard transformer timestep)
+    which generalizes to arbitrary frame distances without a fixed max.
+    """
+
+    def __init__(self, hidden_size: int, max_period: int = 10000):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.max_period = max_period
+
+        # Precompute frequency bands for sincos encoding
+        half_dim = hidden_size // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(half_dim, dtype=torch.float32)
+            / half_dim
+        )
+        self.register_buffer("freqs", freqs)
+
+    def forward(
+        self,
+        batch_size: int,
+        num_frames: int,
+        device: torch.device,
+        start_distance: int = 1,
+    ) -> torch.Tensor:
+        """
+        Generate temporal distance tokens using sincos encoding.
+
+        Args:
+            batch_size: batch size
+            num_frames: number of motion frames (T)
+            device: torch device
+            start_distance: distance of first motion frame from reference
+
+        Returns:
+            [B, T, 1, D] temporal distance tokens
+        """
+        distances = torch.arange(
+            start_distance,
+            start_distance + num_frames,
+            device=device,
+            dtype=torch.float32,
+        )
+
+        # Sincos encoding: [T, D/2] -> [T, D]
+        args = distances.unsqueeze(-1) * self.freqs.to(device)  # [T, D/2]
+        time_tokens = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)  # [T, D]
+
+        # Expand: [1, T, 1, D] -> [B, T, 1, D]
+        return time_tokens.unsqueeze(0).unsqueeze(2).expand(batch_size, -1, -1, -1)
+
+
+# ============================================================================
+# Additive RoPE
+# ============================================================================
+
+
+class AdditiveRoPE2D(nn.Module):
+    """
+    Additive 2D Rotary Position Embedding with adaptive aspect-ratio-aware positions.
+
+    x = x + rope(learnable_token, pos)
+
+    Position mapping is normalized such that:
+    - max(H) * max(W) = 1 (positions are in [0, 1] range)
+    - max(H) / max(W) = vid_h / vid_w (preserves aspect ratio)
+
+    This removes the need for max_h/max_w parameters and generalizes to any resolution.
+    """
+
+    def __init__(self, hidden_size: int, max_period: int = 10000):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.max_period = max_period
+
+        # Learnable token that gets modulated by position
+        self.learnable_token = nn.Parameter(torch.randn(1, 1, hidden_size) * 0.02)
+
+        # Precompute frequency bands for sincos encoding (half for h, half for w)
+        quarter_dim = hidden_size // 4
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(quarter_dim, dtype=torch.float32)
+            / quarter_dim
+        )
+        self.register_buffer("freqs", freqs)
+
+    def _get_adaptive_pos_embed(
+        self, h: int, w: int, device: torch.device
+    ) -> torch.Tensor:
+        """Generate position embeddings with adaptive aspect-ratio-aware mapping."""
+        grid_h, grid_w = get_adaptive_2d_pos(h, w, device)
+
+        # Sincos encoding for h and w positions
+        args_h = grid_h.unsqueeze(-1) * self.freqs.to(device)  # [L, D/4]
+        emb_h = torch.cat([torch.sin(args_h), torch.cos(args_h)], dim=-1)  # [L, D/2]
+
+        args_w = grid_w.unsqueeze(-1) * self.freqs.to(device)  # [L, D/4]
+        emb_w = torch.cat([torch.sin(args_w), torch.cos(args_w)], dim=-1)  # [L, D/2]
+
+        return torch.cat([emb_h, emb_w], dim=-1)  # [L, D]
+
+    def forward(self, x: torch.Tensor, h: int, w: int) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, L, D] or [B, L, D] input tensor
+            h, w: spatial dimensions
+
+        Returns:
+            x with additive RoPE applied
+        """
+        pos = self._get_adaptive_pos_embed(h, w, x.device).unsqueeze(0)  # [1, L, D]
+        rope_emb = self.learnable_token * pos  # [1, L, D]
+
+        if x.ndim == 4:
+            rope_emb = rope_emb.unsqueeze(1)  # [1, 1, L, D]
+
+        return x + rope_emb
+
+
+# ============================================================================
+# Rotary RoPE Attention
+# ============================================================================
+
+
+class RoPE2DAttention(nn.Module):
+    """
+    Attention with 2D Rotary Position Embedding applied to Q and K.
+
+    Uses standard RoPE formulation extended to 2D with adaptive aspect-ratio-aware positions.
+    Position mapping is normalized such that max(H) * max(W) = 1 and preserves aspect ratio.
+    """
+
+    def __init__(self, hidden_size: int, num_heads: int, max_period: int = 10000, qk_norm: bool = False, bias: bool = True):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.max_period = max_period
+        self.qk_norm = qk_norm
+        if qk_norm:
+            self.qk_scale = nn.Parameter(torch.full([num_heads, 1, 1], 10.0))
+
+        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.v_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=bias)
+
+        # Precompute frequency bands (quarter for h, quarter for w)
+        quarter_head_dim = self.head_dim // 4
+        freqs = 1.0 / (
+            max_period
+            ** (
+                torch.arange(0, quarter_head_dim, dtype=torch.float32)
+                / quarter_head_dim
+            )
+        )
+        self.register_buffer("freqs", freqs)
+
+    def _compute_adaptive_rope_freqs(self, h: int, w: int, device: torch.device):
+        """Compute RoPE frequencies for 2D grid with adaptive aspect-ratio mapping."""
+        grid_h, grid_w = get_adaptive_2d_pos(h, w, device)
+
+        freqs_h = torch.outer(grid_h, self.freqs.to(device))  # [L, head_dim/4]
+        freqs_w = torch.outer(grid_w, self.freqs.to(device))  # [L, head_dim/4]
+        freqs = torch.cat([freqs_h, freqs_w], dim=-1)  # [L, head_dim/2]
+
+        return torch.cos(freqs), torch.sin(freqs)
+
+    def _apply_rope(
+        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+    ) -> torch.Tensor:
+        """Apply rotary position embedding."""
+        # x: [B, num_heads, L, head_dim]
+        # cos, sin: [L, head_dim/2]
+        x_half = x.shape[-1] // 2
+        x1, x2 = x[..., :x_half], x[..., x_half:]
+
+        cos = cos.unsqueeze(0).unsqueeze(0)  # [1, 1, L, head_dim/2]
+        sin = sin.unsqueeze(0).unsqueeze(0)
+
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+    @compile_wrapper
+    def forward(self, x: torch.Tensor, h: int, w: int) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, S, D] input (T=temporal, S=sequence length)
+               S may be h*w (spatial only) or h*w + extra tokens (motion, timestep, etc.)
+            h, w: spatial dimensions (h*w <= S)
+        """
+        B, T, S, D = x.shape
+        L = h * w  # spatial token count
+        total = T * S
+
+        x_flat = x.flatten(1, 2)  # [B, T*S, D]
+
+        q = self.q_proj(x_flat)
+        k = self.k_proj(x_flat)
+        v = self.v_proj(x_flat)
+
+        q = q.view(B, total, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, total, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, total, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Get adaptive RoPE frequencies for spatial positions
+        cos_spatial, sin_spatial = self._compute_adaptive_rope_freqs(h, w, x.device)
+        # cos_spatial, sin_spatial: [L, head_dim/2]
+
+        # Build per-token RoPE: spatial tokens get position encoding,
+        # extra tokens (motion, timestep) get zero rotation (cos=1, sin=0)
+        half = cos_spatial.shape[-1]
+        if S > L:
+            extra = S - L
+            ones = torch.ones(extra, half, device=x.device, dtype=cos_spatial.dtype)
+            zeros = torch.zeros(extra, half, device=x.device, dtype=cos_spatial.dtype)
+            cos_frame = torch.cat([cos_spatial, ones], dim=0)   # [S, half]
+            sin_frame = torch.cat([sin_spatial, zeros], dim=0)  # [S, half]
+        else:
+            cos_frame = cos_spatial
+            sin_frame = sin_spatial
+
+        # Tile for all temporal frames
+        cos = cos_frame.unsqueeze(0).expand(T, -1, -1).reshape(total, -1)
+        sin = sin_frame.unsqueeze(0).expand(T, -1, -1).reshape(total, -1)
+
+        # Apply RoPE to Q and K
+        q = self._apply_rope(q, cos, sin)
+        k = self._apply_rope(k, cos, sin)
+
+        # QK-norm (after RoPE)
+        if self.qk_norm:
+            from .layers import _qk_norm
+            q, k = _qk_norm(q, k, self.qk_scale)
+            attn = F.scaled_dot_product_attention(q, k, v, scale=1.0)
+        else:
+            attn = F.scaled_dot_product_attention(q, k, v)
+
+        out = attn.transpose(1, 2).flatten(2, 3)  # [B, T*S, D]
+        out = self.out_proj(out)
+
+        return out.view(B, T, S, D)

tt.py

@@ -0,0 +1,167 @@
+"""
+TemporalTransfer (TT1D): Block-causal temporal attention on motion tokens with 1D temporal RoPE.
+
+Each frame's M motion tokens attend to current and past frames' motion tokens.
+1D RoPE encodes temporal position (frame index) so tokens know their temporal order.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .utils import compile_wrapper
+from .layers import SwiGLU, GELUMLP, RMSNorm
+
+
+# =============================================================================
+# Block-causal mask (cached)
+# =============================================================================
+
+_mask_cache: dict[tuple, torch.Tensor] = {}
+
+
+def get_block_causal_mask(t, m, device):
+    """Bool mask [T*M, T*M] where frame i attends to frames 0..i."""
+    key = (t, m, str(device))
+    if key not in _mask_cache:
+        causal = torch.tril(torch.ones(t, t, device=device, dtype=torch.bool))
+        block = causal[:, :, None, None].expand(-1, -1, m, m)
+        mask = block.permute(0, 2, 1, 3).reshape(t * m, t * m)
+        _mask_cache[key] = mask
+    return _mask_cache[key]
+
+
+# =============================================================================
+# 1D Temporal RoPE
+# =============================================================================
+
+class TemporalRoPE1D(nn.Module):
+    """1D Rotary Position Embedding for temporal dimension.
+
+    Applies RoPE to half the head_dim (temporal), leaves the other half unchanged.
+    """
+
+    def __init__(self, head_dim: int, max_period: float = 10000.0):
+        super().__init__()
+        self.head_dim = head_dim
+        # Use half for temporal RoPE, half unchanged
+        self.rope_dim = head_dim // 2
+        self.unused_dim = head_dim - self.rope_dim
+        self.max_period = max_period
+        self.register_buffer("freqs", self._freqs(), persistent=False)
+
+    def _freqs(self) -> torch.Tensor:
+        half = self.rope_dim // 2
+        return torch.exp(
+            -math.log(self.max_period) * torch.arange(half, dtype=torch.float32) / half
+        )
+
+    def reset_buffers(self) -> None:
+        """Recompute the non-persistent buffers (e.g. after transformers' meta-device loading)."""
+        self.freqs.copy_(self._freqs())
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, t: int, m: int):
+        """
+        Apply 1D temporal RoPE to q and k.
+
+        Args:
+            q, k: [B, H, T*M, head_dim]
+            t: number of frames
+            m: tokens per frame
+        """
+        # Build temporal positions: each token in frame i gets position i
+        # [0,0,...,0, 1,1,...,1, 2,2,...,2, ...] repeated m times per frame
+        positions = torch.arange(t, device=q.device).repeat_interleave(m)  # [T*M]
+
+        angles = positions.unsqueeze(-1).float() * self.freqs.to(q.device)  # [T*M, half]
+        cos = torch.cos(angles).to(q.dtype)
+        sin = torch.sin(angles).to(q.dtype)
+
+        def apply(x):
+            x_rope = x[..., :self.rope_dim]
+            x_pass = x[..., self.rope_dim:]
+            half = self.rope_dim // 2
+            x1, x2 = x_rope[..., :half], x_rope[..., half:]
+            x_rope = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+            return torch.cat([x_rope, x_pass], dim=-1)
+
+        return apply(q), apply(k)
+
+
+# =============================================================================
+# Block-causal temporal attention with 1D RoPE
+# =============================================================================
+
+class BlockCausalTemporalAttention(nn.Module):
+    def __init__(self, hidden_size, num_heads, qk_norm=False):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.qk_norm = qk_norm
+        self.q_proj = nn.Linear(hidden_size, hidden_size)
+        self.k_proj = nn.Linear(hidden_size, hidden_size)
+        self.v_proj = nn.Linear(hidden_size, hidden_size)
+        self.out_proj = nn.Linear(hidden_size, hidden_size)
+        if qk_norm:
+            self.qk_scale = nn.Parameter(torch.full([num_heads, 1, 1], 10.0))
+
+        # 1D temporal RoPE
+        self.rope = TemporalRoPE1D(self.head_dim)
+
+    @compile_wrapper
+    def forward(self, x):
+        from .layers import _qk_norm
+        b, t, m, _ = x.shape
+        x = x.flatten(1, 2)
+        q = self.q_proj(x)
+        k = self.k_proj(x)
+        v = self.v_proj(x)
+        q = q.view(b, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(b, -1, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(b, -1, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # Apply 1D temporal RoPE
+        q, k = self.rope(q, k, t, m)
+
+        mask = get_block_causal_mask(t, m, q.device)
+        if self.qk_norm:
+            q, k = _qk_norm(q, k, self.qk_scale)
+            attn = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=1.0)
+        else:
+            attn = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+
+        x = self.out_proj(attn.transpose(1, 2).flatten(2, 3))
+        return x.unflatten(1, (t, m))
+
+
+# =============================================================================
+# TemporalTransfer (TT1D)
+# =============================================================================
+
+class TemporalTransfer(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, num_heads, ffn_type="swiglu", qk_norm=False):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size)
+        self.attn = BlockCausalTemporalAttention(hidden_size, num_heads, qk_norm=qk_norm)
+        self.norm2 = RMSNorm(hidden_size)
+        if ffn_type == "gelu":
+            self.mlp = GELUMLP(hidden_size, intermediate_size)
+        else:
+            self.mlp = SwiGLU(hidden_size, intermediate_size)
+
+    @compile_wrapper
+    def forward(self, x):
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+if __name__ == "__main__":
+    tt_layer = TemporalTransfer(768, 3072, 12)
+    x = torch.randn(2, 8, 8, 768)
+    y = tt_layer(x)
+    print(f"TT1D: {x.shape} -> {y.shape}")
+    assert y.shape == x.shape
+    print("TT1D smoke test passed!")

tt3d.py

@@ -0,0 +1,505 @@
+"""
+TemporalTransfer3D: Temporal attention with downsampled spatial context and 3D RoPE.
+
+Each frame contributes M motion tokens + S downsampled spatial tokens.
+Block-causal attention across time with 3D RoPE (x, y, t).
+Both motion and spatial tokens receive attention residual.
+Spatial tokens are downsampled via pixel unshuffle + a linear map,
+and upsampled back via a linear map + pixel shuffle for residual add-back.
+
+3D RoPE head_dim split: [x, y, t, unused] with 1/4 each.
+Motion tokens use position (0, 0, t) — no spatial, only temporal.
+Spatial tokens use position (x, y, t) — full 3D.
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import compile_wrapper
+from .layers import SwiGLU, GELUMLP, RMSNorm
+
+
+# =============================================================================
+# Block-causal mask (cached)
+# =============================================================================
+
+_mask_cache: dict[tuple, torch.Tensor] = {}
+
+
+def get_block_causal_mask(t: int, tokens_per_frame: int, device: torch.device) -> torch.Tensor:
+    """Bool mask [T*N, T*N] where frame i attends to frames 0..i."""
+    key = (t, tokens_per_frame, str(device))
+    if key not in _mask_cache:
+        causal = torch.tril(torch.ones(t, t, device=device, dtype=torch.bool))
+        n = tokens_per_frame
+        block = causal[:, :, None, None].expand(-1, -1, n, n)
+        mask = block.permute(0, 2, 1, 3).reshape(t * n, t * n)
+        _mask_cache[key] = mask
+    return _mask_cache[key]
+
+
+# =============================================================================
+# 3D RoPE
+# =============================================================================
+
+def _compute_freqs(dim: int, max_period: float = 10000.0) -> torch.Tensor:
+    """Frequency bands for RoPE. Returns [dim//2]."""
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(half, dtype=torch.float32) / half
+    )
+    return freqs
+
+
+def _apply_rope_1d(x: torch.Tensor, freqs: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+    """
+    Apply RoPE to a slice of x along the last dim.
+
+    Args:
+        x: [..., dim] — the slice of q or k to rotate
+        freqs: [dim//2] — precomputed frequency bands
+        positions: [...] — positions for each token
+
+    Returns:
+        [..., dim] rotated tensor
+    """
+    half = x.shape[-1] // 2
+    angles = positions.unsqueeze(-1).float() * freqs.to(x.device)
+    cos = torch.cos(angles).to(x.dtype)
+    sin = torch.sin(angles).to(x.dtype)
+    x1, x2 = x[..., :half], x[..., half:]
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class RoPE3D(nn.Module):
+    """
+    3D Rotary Position Embedding for (x, y, t) positions.
+
+    Head dim split into 4 equal parts: [x_rope, y_rope, t_rope, unused].
+    The unused portion passes through unchanged (identity).
+    """
+
+    def __init__(self, head_dim: int, max_period: float = 10000.0):
+        super().__init__()
+        self.head_dim = head_dim
+        self.dim_x = head_dim // 4
+        self.dim_y = head_dim // 4
+        self.dim_t = head_dim // 4
+        self.dim_unused = head_dim - self.dim_x - self.dim_y - self.dim_t
+        self.max_period = max_period
+
+        self.register_buffer("freqs_x", _compute_freqs(self.dim_x, max_period), persistent=False)
+        self.register_buffer("freqs_y", _compute_freqs(self.dim_y, max_period), persistent=False)
+        self.register_buffer("freqs_t", _compute_freqs(self.dim_t, max_period), persistent=False)
+
+    def reset_buffers(self) -> None:
+        """Recompute the non-persistent buffers (e.g. after transformers' meta-device loading)."""
+        self.freqs_x.copy_(_compute_freqs(self.dim_x, self.max_period))
+        self.freqs_y.copy_(_compute_freqs(self.dim_y, self.max_period))
+        self.freqs_t.copy_(_compute_freqs(self.dim_t, self.max_period))
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, positions: torch.Tensor):
+        """
+        Args:
+            q, k: [B, H, S, head_dim]
+            positions: [S, 3] — (x, y, t) per token
+        Returns:
+            q_rot, k_rot: same shape
+        """
+        pos_x = positions[:, 0]
+        pos_y = positions[:, 1]
+        pos_t = positions[:, 2]
+
+        def apply(tensor):
+            t_x = tensor[..., :self.dim_x]
+            t_y = tensor[..., self.dim_x:self.dim_x + self.dim_y]
+            t_t = tensor[..., self.dim_x + self.dim_y:self.dim_x + self.dim_y + self.dim_t]
+            t_u = tensor[..., self.dim_x + self.dim_y + self.dim_t:]
+
+            t_x = _apply_rope_1d(t_x, self.freqs_x, pos_x)
+            t_y = _apply_rope_1d(t_y, self.freqs_y, pos_y)
+            t_t = _apply_rope_1d(t_t, self.freqs_t, pos_t)
+
+            return torch.cat([t_x, t_y, t_t, t_u], dim=-1)
+
+        return apply(q), apply(k)
+
+
+# =============================================================================
+# Pixel unshuffle/shuffle spatial resampling
+# =============================================================================
+
+def _hadamard(n: int) -> torch.Tensor:
+    h = torch.ones(1, 1, dtype=torch.float64)
+    while h.shape[0] < n:
+        h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0)
+    if h.shape[0] != n:
+        raise ValueError(f"Walsh-Hadamard needs a power-of-two size, got {n}")
+    return h
+
+
+def _dct(n: int) -> torch.Tensor:
+    """Orthonormal DCT-II ``[frequency, index]``."""
+    k = torch.arange(n, dtype=torch.float64)[:, None]
+    c = torch.arange(n, dtype=torch.float64)[None]
+    m = torch.cos(math.pi * (c + 0.5) * k / n) * math.sqrt(2 / n)
+    m[0] /= math.sqrt(2)
+    return m
+
+
+def _chirp_signs(n: int, i: int) -> torch.Tensor:
+    """A deterministic +-1 pattern per frequency ``i`` (quadratic chirp, never 0)."""
+    c = torch.arange(n, dtype=torch.float64)
+    s = torch.sign(torch.cos(math.pi * (c * c * (2 * i + 1) + 3 * i * c) / n + 0.25 * i))
+    s[s == 0] = 1.0
+    return s
+
+
+def structured_weight(dim: int, factor: int) -> torch.Tensor:
+    """The fixed structured down weight ``[D, D*f^2]`` (column index ``c*f^2 + p``,
+    the pixel-unshuffle channel order).
+
+    Over the f^2 positions of a patch, take orthonormal Walsh-Hadamard frequencies
+    ``z_i = sum_p h_i[p] x[:, p]``; then ``y = (1/f) sum_i Q_i z_i`` with
+    ``Q_i = C^T diag(chi_i) C`` (C the orthonormal DCT-II over channels, chi_i
+    deterministic chirp signs): one orthogonal D x D map per position frequency.
+    Depends only on ``D`` and ``f``: nothing is stored.
+    """
+    f2 = factor * factor
+    h = _hadamard(f2) / factor                                   # [i, p], orthonormal rows
+    c = _dct(dim)
+    q = torch.stack([c.T @ (_chirp_signs(dim, i)[:, None] * c) for i in range(f2)])
+    w = torch.einsum("ioc,ip->ocp", q, h) / factor               # [o, c, p]
+    return w.reshape(dim, dim * f2).float()
+
+
+def _to_unshuffled(x: torch.Tensor, h: int, w: int, f: int) -> torch.Tensor:
+    """[B, T, H*W, D] -> [B, T, H'W', D*f*f] (index c*f*f + p)."""
+    B, T = x.shape[:2]
+    x = x.unflatten(2, (h, w)).permute(0, 1, 4, 2, 3).flatten(0, 1)  # [BT, D, H, W]
+    x = F.pixel_unshuffle(x, f)                                      # [BT, D*f*f, H', W']
+    return x.flatten(2).transpose(1, 2).unflatten(0, (B, T))         # [B, T, H'W', D*f*f]
+
+
+def _from_unshuffled(x: torch.Tensor, h: int, w: int, f: int) -> torch.Tensor:
+    """[B, T, H'W', D*f*f] -> [B, T, H*W, D] (inverse of ``_to_unshuffled``)."""
+    B, T = x.shape[:2]
+    x = x.transpose(-1, -2).unflatten(-1, (h // f, w // f)).flatten(0, 1)  # [BT, D*f*f, H', W']
+    x = F.pixel_shuffle(x, f)                                              # [BT, D, H, W]
+    return x.unflatten(0, (B, T)).flatten(3, 4).transpose(-1, -2)          # [B, T, H*W, D]
+
+
+class SpatialDownsample(nn.Module):
+    """
+    [B, T, H*W, D] -> [B, T, (H/f)*(W/f), D]: pixel unshuffle, the fixed dense
+    ``structured_weight`` (no parameters, rebuilt from D and f, not saved), then a
+    trainable D x D channel mix ``I + mix`` (``mix`` zero at init). Params: D².
+    """
+
+    def __init__(self, hidden_size: int, factor: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.factor = factor
+        self.register_buffer("fixed", self._fixed(), persistent=False)
+        self.mix = nn.Parameter(torch.zeros(hidden_size, hidden_size))
+
+    def _fixed(self) -> torch.Tensor:
+        return structured_weight(self.hidden_size, self.factor)
+
+    def reset_parameters(self) -> None:
+        """Channel mix back to the identity (after a generic init such as mup_init)."""
+        nn.init.zeros_(self.mix)
+
+    def reset_buffers(self) -> None:
+        """Recompute the fixed weight (e.g. after transformers' meta-device loading)."""
+        self.fixed.copy_(self._fixed())
+
+    def forward(self, x: torch.Tensor, h: int, w: int) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, H*W, D]
+            h, w: spatial grid dims
+        Returns:
+            [B, T, (H//f)*(W//f), D]
+        """
+        y = F.linear(_to_unshuffled(x, h, w, self.factor), self.fixed)  # [B, T, H'W', D]
+        return y + F.linear(y, self.mix)
+
+
+class SpatialUpsample(nn.Module):
+    """
+    [B, T, (H/f)*(W/f), D] -> [B, T, H*W, D], the counterpart of SpatialDownsample:
+    channel mix ``I + mix``, then ``f`` times the transpose of the fixed down weight
+    (writes go back in the basis that was read), pixel shuffle. Params: D².
+    """
+
+    def __init__(self, hidden_size: int, factor: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.factor = factor
+        self.register_buffer("fixed", self._fixed(), persistent=False)
+        self.mix = nn.Parameter(torch.zeros(hidden_size, hidden_size))
+
+    def _fixed(self) -> torch.Tensor:
+        return self.factor * structured_weight(self.hidden_size, self.factor).T.contiguous()
+
+    def reset_parameters(self) -> None:
+        """Channel mix back to the identity (after a generic init such as mup_init)."""
+        nn.init.zeros_(self.mix)
+
+    def reset_buffers(self) -> None:
+        """Recompute the fixed weight (e.g. after transformers' meta-device loading)."""
+        self.fixed.copy_(self._fixed())
+
+    def forward(self, x: torch.Tensor, h: int, w: int) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, H'·W', D]    where H' = h // f, W' = w // f
+            h, w: ORIGINAL spatial grid dims (output size)
+        Returns:
+            [B, T, H*W, D]
+        """
+        x = x + F.linear(x, self.mix)
+        return _from_unshuffled(F.linear(x, self.fixed), h, w, self.factor)
+
+
+# =============================================================================
+# TemporalTransfer3D
+# =============================================================================
+
+class TemporalTransfer3D(nn.Module):
+    """
+    Temporal attention with downsampled spatial context and 3D RoPE.
+
+    Both motion and spatial tokens participate in block-causal attention.
+    Spatial tokens are downsampled (pixel unshuffle + linear), attend temporally,
+    then upsampled (linear + pixel shuffle) for residual add-back to the spatial stream.
+
+    Args:
+        hidden_size: model dimension
+        intermediate_size: FFN hidden dim
+        num_heads: attention heads
+        downsample_factor: spatial downsample factor (e.g. 4 → 16x16 → 4x4)
+        ffn_type: "swiglu" or "gelu"
+        qk_norm: use cosine-similarity attention
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        num_heads: int,
+        downsample_factor: int = 4,
+        ffn_type: str = "swiglu",
+        qk_norm: bool = False,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.downsample_factor = downsample_factor
+        self.qk_norm = qk_norm
+
+        # Spatial resampling: fixed structured weight + D x D channel mix
+        self.spatial_down = SpatialDownsample(hidden_size, downsample_factor)
+        self.spatial_up = SpatialUpsample(hidden_size, downsample_factor)
+        # Zero-init output projection for spatial add-back (identity at init)
+        self.spatial_out_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        nn.init.zeros_(self.spatial_out_proj.weight)
+
+        # Pre-norm
+        self.norm1 = RMSNorm(hidden_size)
+        self.norm2 = RMSNorm(hidden_size)
+
+        # Attention projections
+        self.q_proj = nn.Linear(hidden_size, hidden_size)
+        self.k_proj = nn.Linear(hidden_size, hidden_size)
+        self.v_proj = nn.Linear(hidden_size, hidden_size)
+        self.out_proj = nn.Linear(hidden_size, hidden_size)
+
+        if qk_norm:
+            self.qk_scale = nn.Parameter(torch.full([num_heads, 1, 1], 10.0))
+
+        # 3D RoPE
+        self.rope = RoPE3D(self.head_dim)
+
+        # FFN
+        if ffn_type == "gelu":
+            self.mlp = GELUMLP(hidden_size, intermediate_size)
+        else:
+            self.mlp = SwiGLU(hidden_size, intermediate_size)
+
+        # Position cache
+        self._pos_cache: dict[tuple, torch.Tensor] = {}
+
+    def _build_positions(
+        self, M: int, ds_h: int, ds_w: int, T: int, device: torch.device
+    ) -> torch.Tensor:
+        """
+        Build 3D positions for all tokens across all frames.
+
+        Per-frame layout: [MT_0..MT_{M-1}, DS_(0,0)..DS_(ds_w-1,ds_h-1)]
+        Motion: (0, 0, t)    Spatial: (x, y, t)
+
+        Returns: [T * (M + ds_h*ds_w), 3]
+        """
+        key = (M, ds_h, ds_w, T, str(device))
+        if key in self._pos_cache:
+            return self._pos_cache[key]
+
+        tokens_per_frame = M + ds_h * ds_w
+        positions = torch.zeros(T, tokens_per_frame, 3, device=device)
+
+        for t in range(T):
+            # Motion tokens: (0, 0, t)
+            positions[t, :M, 2] = t
+            # Spatial tokens: (x, y, t)
+            idx = M
+            for row in range(ds_h):
+                for col in range(ds_w):
+                    positions[t, idx, 0] = col
+                    positions[t, idx, 1] = row
+                    positions[t, idx, 2] = t
+                    idx += 1
+
+        positions = positions.reshape(T * tokens_per_frame, 3)
+        self._pos_cache[key] = positions
+        return positions
+
+    @compile_wrapper
+    def _attention(self, x: torch.Tensor, positions: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Block-causal attention with 3D RoPE.
+
+        Args:
+            x: [B, S, D] — normed, flattened (T * tokens_per_frame)
+            positions: [S, 3] — (x, y, t) per token
+            mask: [S, S] — block-causal bool mask
+        """
+        from .layers import _qk_norm
+
+        B, S, D = x.shape
+        q = self.q_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # 3D RoPE
+        q, k = self.rope(q, k, positions)
+
+        # QK-norm (cosine-sim attention)
+        if self.qk_norm:
+            q, k = _qk_norm(q, k, self.qk_scale)
+            attn = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=1.0)
+        else:
+            attn = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+
+        return self.out_proj(attn.transpose(1, 2).reshape(B, S, D))
+
+    def forward(
+        self,
+        motion_tokens: torch.Tensor,
+        spatial_tokens: torch.Tensor,
+        spatial_h: int,
+        spatial_w: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            motion_tokens:  [B, T, M, D] — motion tokens from DINOv3 stream
+            spatial_tokens: [B, T, H*W, D] — spatial tokens from DINOv3 stream
+            spatial_h, spatial_w: spatial grid dimensions (e.g. 16, 16)
+
+        Returns:
+            motion_tokens:  [B, T, M, D] — enriched motion tokens
+            spatial_tokens: [B, T, H*W, D] — spatial tokens with temporal residual
+        """
+        B, T, M, D = motion_tokens.shape
+        f = self.downsample_factor
+        ds_h, ds_w = spatial_h // f, spatial_w // f
+        tokens_per_frame = M + ds_h * ds_w
+
+        # 1. Downsample spatial: pixel unshuffle + linear
+        ds_spatial = self.spatial_down(spatial_tokens, spatial_h, spatial_w)
+        # ds_spatial: [B, T, ds_h*ds_w, D]
+
+        # 2. Concat motion + downsampled spatial per frame
+        x = torch.cat([motion_tokens, ds_spatial], dim=2)  # [B, T, M+S_down, D]
+
+        # 3. Pre-norm + flatten temporal dim
+        x_normed = self.norm1(x).flatten(1, 2)  # [B, T*(M+S_down), D]
+
+        # 4. Positions and mask
+        positions = self._build_positions(M, ds_h, ds_w, T, x.device)
+        mask = get_block_causal_mask(T, tokens_per_frame, x.device)
+
+        # 5. Block-causal attention with 3D RoPE
+        attn_out = self._attention(x_normed, positions, mask)
+        attn_out = attn_out.unflatten(1, (T, tokens_per_frame))
+
+        # 6. Attention residual on ALL tokens (motion + spatial)
+        x = x + attn_out
+
+        # 7. FFN on ALL tokens (motion + spatial)
+        x = x + self.mlp(self.norm2(x))
+
+        # 8. Split back into motion and spatial
+        motion_tokens = x[:, :, :M]
+        spatial_part = x[:, :, M:]
+
+        # 9. Spatial add-back: zero-init proj → upsample → residual
+        #    spatial_out_proj is zero-init, so initially this is a no-op
+        spatial_residual = self.spatial_up(
+            self.spatial_out_proj(spatial_part), spatial_h, spatial_w
+        )
+        spatial_tokens = spatial_tokens + spatial_residual
+
+        return motion_tokens, spatial_tokens
+
+
+# =============================================================================
+# Smoke test
+# =============================================================================
+
+if __name__ == "__main__":
+    B, T, M, D = 2, 8, 8, 768
+    H, W = 16, 16
+
+    layer = TemporalTransfer3D(
+        hidden_size=D,
+        intermediate_size=D * 4,
+        num_heads=12,
+        downsample_factor=4,
+    )
+
+    motion = torch.randn(B, T, M, D)
+    spatial = torch.randn(B, T, H * W, D)
+
+    motion_out, spatial_out = layer(motion, spatial, H, W)
+    print(f"Input:  motion={motion.shape}, spatial={spatial.shape}")
+    print(f"Output: motion={motion_out.shape}, spatial={spatial_out.shape}")
+    assert motion_out.shape == (B, T, M, D)
+    assert spatial_out.shape == (B, T, H * W, D)
+
+    # Verify spatial add-back is zero-init (residual starts as no-op)
+    diff = (spatial_out - spatial).abs().max().item()
+    print(f"Spatial diff at init (should be ~0 from zero-init out proj): {diff:.6f}")
+
+    # Check gradient flows through spatial
+    motion.requires_grad_(True)
+    spatial.requires_grad_(True)
+    m_out, s_out = layer(motion, spatial, H, W)
+    loss = m_out.sum() + s_out.sum()
+    loss.backward()
+    print(f"Gradient on spatial: {spatial.grad is not None}, norm={spatial.grad.norm():.4f}")
+    print(f"Gradient on motion: {motion.grad is not None}, norm={motion.grad.norm():.4f}")
+
+    # Different downsample factors
+    for ds in [1, 2, 4, 8]:
+        l = TemporalTransfer3D(D, D * 4, 12, downsample_factor=ds)
+        m_o, s_o = l(motion.detach(), spatial.detach(), H, W)
+        ds_tokens = (H // ds) * (W // ds)
+        print(f"  ds={ds}: {ds_tokens} spatial tokens/frame, total={M + ds_tokens}/frame")
+
+    print("\nSmoke test passed!")

utils.py

@@ -0,0 +1,28 @@
+import importlib
+import torch
+from . import env
+from .env import TORCH_COMPILE as _  # noqa: F401 (bundles env.py)
+
+
+def import_class(cls_string):
+    if not isinstance(cls_string, str):
+        return cls_string
+    module, cls = cls_string.rsplit(".", 1)
+    module = importlib.import_module(module)
+    return getattr(module, cls)
+
+
+def compile_wrapper(func):
+    """Decorator: compile the function once with torch.compile when TORCH_COMPILE=True."""
+    _compiled = None
+
+    def wrapper(*args, **kwargs):
+        nonlocal _compiled
+        if env.TORCH_COMPILE:
+            if _compiled is None:
+                _compiled = torch.compile(func, **env.COMPILE_KWARG)
+            return _compiled(*args, **kwargs)
+        else:
+            return func(*args, **kwargs)
+
+    return wrapper