MotifVAE

Co-authored-by: gkalstn0 <gkalstn0@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,44 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/korean_ocr_recon.png filter=lfs diff=lfs merge=lfs -text
+assets/natural_recon.png filter=lfs diff=lfs merge=lfs -text
+assets/ucf101_fvd_curves.png filter=lfs diff=lfs merge=lfs -text
+assets/video_recon.gif filter=lfs diff=lfs merge=lfs -text
+assets/omnidoc_recon.png filter=lfs diff=lfs merge=lfs -text
+assets/hero_video_recon.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/ocr_recon.png filter=lfs diff=lfs merge=lfs -text
+assets/video_recon_2.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/video_recon_1.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,206 @@
+---
+license: mit
+library_name: diffusers
+pipeline_tag: image-to-image
+tags:
+- video-vae
+- vae
+- video-generation
+- tokenizer
+- pytorch
+---
+
+# MotifVAE
+
+MotifVAE is a 3D causal video VAE with 4x temporal and 32x spatial compression and a
+128-channel latent space, built as a high-compression tokenizer for latent video diffusion
+models. A 256x256x17-frame clip encodes to a `(128, 5, 8, 8)` latent, which is 1/16 the
+spatio-temporal tokens of an 8x-spatial VAE, so downstream diffusion-transformer training and
+inference cost less. The decoder is much larger than the encoder (about 4:1 in parameters),
+which helps it keep fine detail such as small text in documents.
+
+It also runs as a plain image VAE at `T=1`. In its compression class it has the best document
+text fidelity, video reconstruction, and diffusability (how well a diffusion model converges on
+its latents) of the tokenizers we compared.
+
+## Architecture
+
+| | |
+|---|---|
+| Type | 3D causal VAE, attention-free (pure convolution) |
+| Compression | temporal 4x (T = 4n+1 grammar) / spatial 32x32 |
+| Latent | 128 channels, deterministic encoder |
+| Parameters | 1.244B total: encoder 242M, decoder 1,002M (1:4.15) |
+| Encoder | 5 down stages (3x spatial-only, 2x spatio-temporal), 2 resblocks/stage, parameter-free global skip connections |
+| Decoder | 1.5x channel width, 3 resblocks/stage, sub-pixel (PixelShuffle + ICNR) spatial upsampling, trilinear temporal upsampling |
+| Norm / act | LayerNorm / SiLU |
+
+## Usage
+
+The model code ships inside this repository (`modeling_motifvae.py`) and loads through the
+stock diffusers `AutoModel` — no extra package beyond `torch`, `diffusers`, `einops`:
+
+```python
+import torch
+from diffusers import AutoModel
+
+vae = AutoModel.from_pretrained("Motif-Technologies/motif-vae", trust_remote_code=True)
+vae = vae.to("cuda", dtype=torch.bfloat16).eval()
+
+# video: (B, 3, T, H, W) in [-1, 1], T = 4n+1 (1, 5, 9, 13, 17, ...),
+# H/W multiples of 32.  T=1 works as a plain image VAE.
+x = torch.randn(1, 3, 17, 256, 256, device="cuda", dtype=torch.bfloat16)
+with torch.no_grad():
+    z = vae.encode(x).latent_dist.sample()   # (1, 128, 5, 8, 8)
+    rec = vae.decode(z).sample               # (1, 3, 17, 256, 256)
+```
+
+For diffusion training, whiten the latent to roughly unit variance with the per-channel
+statistics shipped in the config — `shift` is the channel mean and `scale` the channel std
+(128 values each, measured over Kinetics-400 video (256x256x17) and ImageNet-1K images (256x256)).
+Apply per channel, reshaping both to `(1, 128, 1, 1, 1)`:
+
+```python
+z = vae.encode(x).latent_dist.sample()
+shift = torch.tensor(vae.config.shift, device=z.device).view(1, -1, 1, 1, 1)
+scale = torch.tensor(vae.config.scale, device=z.device).view(1, -1, 1, 1, 1)
+z = (z - shift) / scale          # whiten for the diffusion model
+# ... sample/denoise in this space, then invert before decoding:
+z = z * scale + shift
+rec = vae.decode(z).sample
+```
+
+Recompute these statistics on your own data if its distribution differs from natural video.
+
+## Evaluation
+
+Baseline numbers are either the published values cited in each table or measured by us under
+the same protocol. MotifVAE is a **32x spatial** tokenizer, so the comparison is against the
+other 32x tokenizers (LTX-Video, LTX-2, and the f32 image VAEs below); the lower-compression
+16x Wan2.2 appears only as a reference upper bound.
+
+### Image reconstruction — ImageNet-1K val (256x256) & FFHQ (1024x1024)
+
+All models below are **32x (f32) image VAEs**. MotifVAE is measured by us with per-image
+PSNR / SSIM (the published-table convention); baseline rows are the published numbers from the
+**Qwen-Image-VAE-2.0** technical report (Table 2, same ImageNet-256 / FFHQ protocol). Our own
+runs of the baselines reproduce those values to within 0.2 dB, so the comparison is direct.
+
+| Model | Setting | ImageNet PSNR ↑ | SSIM ↑ | FFHQ PSNR ↑ | SSIM ↑ |
+|---|---|---|---|---|---|
+| **MotifVAE** | **f32c128** | **30.27** | **0.852** | **36.17** | **0.921** |
+| Qwen-Image-VAE-2.0 | f32c128 | 29.69 | 0.842 | 35.91 | 0.918 |
+| LTX-Video | f32c128 | 29.57 | 0.833 | 35.56 | 0.905 |
+| HunyuanImage-2.1 | f32c64 | 28.67 | 0.820 | 35.30 | 0.911 |
+| LTX-2 | f32c128 | 26.06 | 0.793 | 33.63 | 0.906 |
+| DC-AE (Sana) | f32c32 | 24.82 | 0.690 | 31.35 | 0.830 |
+
+At the f32c128 channel budget MotifVAE leads both benchmarks among 32x image VAEs, above
+Qwen-Image-VAE-2.0's f32c128 (29.69 / 35.91 dB), LTX-Video, LTX-2, and the lighter f32c64 /
+f32c32 models.
+
+### Qualitative comparison (real photographs)
+
+GT vs frozen-tokenizer reconstruction on ImageNet-1K val photographs; the red box marks the
+zoom strip shown below each row.
+
+![Natural-image reconstruction (ImageNet-1K val): GT vs MotifVAE vs LTX-Video vs LTX-2 vs Wan2.2](assets/natural_recon.png)
+
+These rows are high-frequency patterns that high-compression VAEs find hard: an armadillo's
+banded shell, a dense solar-panel grid, brickwork, and vehicle markings. LTX-Video and LTX-2
+(same 32x class) lose the pattern, and the solar grid collapses for them entirely. MotifVAE
+keeps it and stays close to the 16x Wan2.2 reference, which also struggles on some of these.
+
+### Document text fidelity — OmniDoc-TokenBench (256x256)
+
+f32-compression VAEs on OmniDoc-TokenBench (~3K text-rich documents). Baseline rows are the
+published numbers from the **Qwen-Image-VAE-2.0** technical report (Table 3, same benchmark
+/ resolution / OCR-NED protocol); sorted by NED.
+
+| Model | Setting | SSIM ↑ | PSNR ↑ | LPIPS ↓ | FID ↓ | OCR-NED ↑ |
+|---|---|---|---|---|---|---|
+| **MotifVAE** | **f32c128** | **0.849** | **22.40** | **0.052** | **3.43** | **0.825** |
+| Qwen-Image-VAE-2.0 | f32c128 | 0.844 | 22.13 | 0.064 | 3.36 | 0.707 |
+| LTX-Video | f32c128 | 0.806 | 20.92 | 0.119 | 17.10 | 0.565 |
+| HunyuanImage-2.1 | f32c64 | 0.781 | 19.85 | 0.096 | 5.19 | 0.490 |
+| LTX-2 | f32c128 | 0.735 | 18.41 | 0.119 | 9.94 | 0.357 |
+| DC-AE (Sana) | f32c32 | 0.526 | 15.62 | 0.144 | 7.26 | 0.069 |
+
+MotifVAE has the highest OCR-NED of any f32 VAE here: +0.12 over Qwen's f32c128 (same channel
+budget) and about 2.3x that of LTX-Video / LTX-2. It also leads every pixel metric in its class.
+The figure below covers English print, Korean print, and Korean handwriting; each row shows the
+full page and a zoom-in on the red-boxed region. The top two rows are OmniDoc-TokenBench
+(English); the bottom two are Korean printed text and handwriting from AIHUB (the table metrics
+above are OmniDoc-only):
+
+![Document reconstruction — top two rows OmniDoc-TokenBench (English), bottom two rows Korean print + handwriting from AIHUB; GT vs MotifVAE vs LTX-Video vs LTX-2 vs Wan2.2](assets/ocr_recon.png)
+
+### Video reconstruction — Kinetics-400 (256x256) & OpenVid-1M (720p)
+
+The same four frozen tokenizers on both benchmarks, under one pipeline: the three 32x spatial
+tokenizers (MotifVAE, LTX-Video, LTX-2) and the 16x Wan2.2 as a reference.
+
+**Kinetics-400 val — 17x256x256**
+
+| Model | compression | PSNR ↑ | SSIM ↑ | LPIPS ↓ |
+|---|---|---|---|---|
+| **MotifVAE** | **4x32x32** | **35.37** | **0.954** | **0.046** |
+| LTX-Video | 8x32x32 | 32.58 | 0.919 | 0.101 |
+| LTX-2 | 8x32x32 | 28.05 | 0.888 | 0.128 |
+| Wan2.2 | 4x16x16 | 37.41 | 0.966 | 0.034 |
+
+**OpenVid-1M — native-720p, 17x1280x704**
+
+| Model | compression | PSNR ↑ | SSIM ↑ | LPIPS ↓ |
+|---|---|---|---|---|
+| **MotifVAE** | **4x32x32** | **36.66** | **0.965** | **0.044** |
+| LTX-Video | 8x32x32 | 34.04 | 0.942 | 0.093 |
+| LTX-2 | 8x32x32 | 31.95 | 0.937 | 0.094 |
+| Wan2.2 | 4x16x16 | 38.13 | 0.972 | 0.035 |
+
+Among the 32x spatial tokenizers MotifVAE leads both tables: +2.7 dB over LTX-Video and
++7.3 dB over LTX-2 on K400, and +2.6 / +4.7 dB on 720p. At 720p it stays above 30 dB /
+0.05 LPIPS (36.7 / 0.044), and at 512x512 it reaches PSNR 39.68. The 16x Wan2.2 reference is a
+little ahead; the gap to the other 32x VAEs is largest on high-motion clips (examples below).
+
+> **Protocol.** K400: 256 random val clips; OpenVid-1M: 30 native-1280x720 web clips
+> centre-cropped to 1280x704. Both use 17 frames/clip and report per-video PSNR / SSIM / LPIPS
+> (pooled per clip, then averaged) under one pipeline, so each table's ranking is exact;
+> absolute values aren't comparable across papers.
+
+### Video reconstruction examples
+
+Frozen-tokenizer reconstruction on high-motion Kinetics-400 clips. Each clip is one strip:
+Input / MotifVAE / LTX-Video / LTX-2 / Wan2.2 (labelled). MotifVAE keeps the fast motion close
+to the 16x Wan2.2; LTX-Video and LTX-2 blur it.
+
+<video controls autoplay loop muted src="https://huggingface.co/Motif-Technologies/motif-vae/resolve/main/assets/hero_video_recon.mp4"></video>
+
+<video controls autoplay loop muted src="https://huggingface.co/Motif-Technologies/motif-vae/resolve/main/assets/video_recon_1.mp4"></video>
+
+<video controls autoplay loop muted src="https://huggingface.co/Motif-Technologies/motif-vae/resolve/main/assets/video_recon_2.mp4"></video>
+
+### Diffusability — UCF101 class-conditional generation
+
+Latte-XL/1 (patch-1, rectified flow, 250K steps, no-CFG FVD-2048) over each frozen
+tokenizer, under the Latte-official **stride-3** protocol. The Wan2.2 control (FVD 204) is
+in line with values reported for this setting in the literature, confirming the harness
+reproduces the expected scale.
+
+![UCF101 FVD convergence over frozen tokenizers](assets/ucf101_fvd_curves.png)
+
+| Tokenizer | FVD @50K | FVD @100K | FVD @250K |
+|---|---|---|---|
+| **MotifVAE** (4x32x32) | **636** | **385** | **276** |
+| LTX-2 (8x32x32) | 952 | 556 | 450 |
+| LTX-Video (8x32x32) | 1251 | 732 | 593 |
+| Wan2.2 (4x16x16) | 656 | 304 | 204 |
+
+Among the 32x spatial tokenizers MotifVAE leads at every milestone, finishing 39% below LTX-2
+and 53% below LTX-Video. It reaches LTX-2's final 250K FVD by about 75K steps (~3x faster) and
+stays close to the 16x Wan2.2 reference. Reconstruction quality and diffusability are separate
+properties: LTX-Video reconstructs better than LTX-2 but is the worst of the three for diffusion.
+
+## License
+
+MIT

assets/hero_video_recon.mp4

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

assets/video_recon_1.mp4

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

assets/video_recon_2.mp4

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

config.json

@@ -0,0 +1,311 @@
+{
+  "_class_name": "MotifVAE",
+  "_diffusers_version": "0.38.0",
+  "auto_map": {
+    "AutoModel": "modeling_motifvae.MotifVAE"
+  },
+  "base_channels": [
+    96,
+    192,
+    384,
+    768,
+    768,
+    768
+  ],
+  "decoder_base_channels": [
+    144,
+    288,
+    576,
+    1152,
+    1152,
+    1152
+  ],
+  "decoder_num_resblocks": 3,
+  "decoder_variant": "V4",
+  "dropout": 0.0,
+  "encoder_num_resblocks": 2,
+  "input_type": "video",
+  "latent_dim": 128,
+  "layer_res_type": [
+    "ResnetBlock2D",
+    "ResnetBlock2D",
+    "ResnetBlock2D",
+    "ResnetBlock3D",
+    "ResnetBlock3D"
+  ],
+  "layer_type": [
+    "hw",
+    "hw",
+    "hw",
+    "thw",
+    "thw"
+  ],
+  "mid_layers_type": [
+    "ResnetBlock3D",
+    "Identity",
+    "ResnetBlock3D"
+  ],
+  "norm_type": "layernorm",
+  "scale": [
+    0.193781,
+    0.160648,
+    0.212569,
+    0.295691,
+    0.157944,
+    0.285008,
+    0.255515,
+    0.229688,
+    0.221633,
+    0.232136,
+    0.139637,
+    0.244213,
+    0.19225,
+    0.209093,
+    0.24903,
+    0.33651,
+    0.224575,
+    0.192788,
+    0.177301,
+    0.191538,
+    0.181008,
+    0.224443,
+    0.21792,
+    0.219838,
+    0.15724,
+    0.215918,
+    0.197021,
+    0.228155,
+    0.165341,
+    0.202019,
+    0.196588,
+    0.242306,
+    0.170205,
+    0.169329,
+    0.271,
+    0.237309,
+    0.172621,
+    0.303556,
+    0.190824,
+    0.175744,
+    0.194385,
+    0.189157,
+    0.224383,
+    0.187539,
+    0.159795,
+    0.209762,
+    0.17788,
+    0.160627,
+    0.189383,
+    0.137677,
+    0.191022,
+    0.181854,
+    0.183199,
+    0.181121,
+    0.189024,
+    0.183419,
+    0.18505,
+    0.185571,
+    0.203301,
+    0.260823,
+    0.166004,
+    0.143982,
+    0.556997,
+    0.252643,
+    0.296122,
+    0.373026,
+    0.162889,
+    0.261822,
+    0.18582,
+    0.210161,
+    0.228212,
+    0.217485,
+    0.165954,
+    0.189418,
+    0.210086,
+    0.186101,
+    0.350991,
+    0.174668,
+    0.199964,
+    0.22358,
+    0.172738,
+    0.159074,
+    0.284967,
+    0.148048,
+    0.192643,
+    0.185355,
+    0.264092,
+    0.148031,
+    0.285184,
+    0.177816,
+    0.198399,
+    0.222337,
+    0.230656,
+    0.239789,
+    0.141241,
+    0.360838,
+    0.194351,
+    0.152777,
+    0.197615,
+    0.166155,
+    0.194582,
+    0.166021,
+    0.170577,
+    0.318022,
+    0.271652,
+    0.200244,
+    0.18932,
+    0.179213,
+    0.179843,
+    0.246742,
+    0.237725,
+    0.189994,
+    0.174361,
+    0.206378,
+    0.178746,
+    0.174695,
+    0.22207,
+    0.165952,
+    0.161002,
+    0.242472,
+    0.240569,
+    0.257974,
+    0.297275,
+    0.225371,
+    0.175571,
+    0.177474,
+    0.155566,
+    0.179106
+  ],
+  "shift": [
+    -0.164127,
+    0.092961,
+    -0.081433,
+    -0.086702,
+    -0.130157,
+    -0.008459,
+    -0.068426,
+    -0.192902,
+    0.051399,
+    -0.047744,
+    -0.066346,
+    0.171836,
+    0.059536,
+    -0.168096,
+    -0.031525,
+    -0.071924,
+    0.187083,
+    0.021963,
+    0.186838,
+    -0.039138,
+    -0.220178,
+    0.139961,
+    0.088891,
+    0.054164,
+    -0.067481,
+    0.165844,
+    -0.132874,
+    0.00983,
+    0.22256,
+    0.060354,
+    0.068444,
+    -0.010999,
+    0.194253,
+    -0.27535,
+    -0.10218,
+    -0.077211,
+    0.000231,
+    -0.145139,
+    -0.249444,
+    0.22325,
+    0.016289,
+    -0.084124,
+    -0.087526,
+    0.210053,
+    -0.215475,
+    -0.15729,
+    -0.124322,
+    -0.173046,
+    0.182702,
+    -0.212797,
+    0.141195,
+    0.155048,
+    -0.039869,
+    -0.103284,
+    0.077888,
+    0.256192,
+    -0.125461,
+    0.197733,
+    0.219629,
+    0.007501,
+    -0.153288,
+    0.199373,
+    0.072393,
+    0.044252,
+    -0.002293,
+    -0.032271,
+    -0.276078,
+    0.040138,
+    -0.150994,
+    0.128878,
+    -0.031116,
+    -0.129651,
+    -0.149721,
+    0.128694,
+    0.040513,
+    0.044896,
+    0.04819,
+    -0.122257,
+    0.094014,
+    0.024673,
+    -0.085999,
+    0.269909,
+    -0.016975,
+    -0.143176,
+    -0.143257,
+    0.100177,
+    0.049843,
+    -0.141697,
+    -0.148462,
+    0.191348,
+    0.05774,
+    -0.066514,
+    0.06819,
+    -0.019659,
+    -0.15781,
+    0.167692,
+    -0.044084,
+    0.185789,
+    0.169589,
+    0.085906,
+    -0.156088,
+    -0.045654,
+    0.181734,
+    -0.090986,
+    0.122933,
+    -0.186479,
+    -0.035989,
+    0.098177,
+    0.121564,
+    -0.093159,
+    -0.142898,
+    0.112587,
+    -0.242465,
+    0.127805,
+    0.142723,
+    0.022688,
+    -0.178088,
+    0.091873,
+    -0.125931,
+    -0.04513,
+    -0.050111,
+    -0.107747,
+    0.06705,
+    0.040281,
+    0.135995,
+    0.16423,
+    0.210625,
+    0.208269
+  ],
+  "t_interpolation": "trilinear",
+  "upsample_residual": true
+}

diffusion_pytorch_model.safetensors

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

modeling_motifvae.py

@@ -0,0 +1,1132 @@
+"""MotifVAE — a 3D causal video VAE for HuggingFace ``AutoModel``.
+
+MotifVAE compresses video by 4x in time and 32x in each spatial dimension
+into a 128-channel deterministic latent. The
+encoder is fully 3D causal and uses a parameter-free Global Skip Connection
+(GSC) at every down stage; the decoder mirrors the stage layout but is
+deliberately wider (asymmetric, decoder-heavy). Spatial upsampling is
+sub-pixel convolution with ICNR initialisation to suppress checkerboard
+artifacts.
+
+The whole model lives in this single file so the released checkpoint loads
+straight from the Hub via remote code::
+
+    import diffusers
+    vae = diffusers.AutoModel.from_pretrained("<repo>", trust_remote_code=True)
+"""
+
+from dataclasses import dataclass
+from typing import List, Literal, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from diffusers import ConfigMixin, ModelMixin
+from diffusers.configuration_utils import register_to_config
+from diffusers.utils import BaseOutput
+
+
+# ---------------------------------------------------------------------------
+# Block-type resolution
+# ---------------------------------------------------------------------------
+#
+# The config encodes residual / mid block choices as strings (e.g.
+# ``"ResnetBlock2D"``). ``resolve_str_to_obj`` maps those strings to the
+# classes defined in this module. ``"Identity"`` maps to ``nn.Identity`` so a
+# mid stage can be made attention-free while keeping the surrounding key
+# layout intact.
+
+_STR_TO_OBJ_OVERRIDES = {"Identity": nn.Identity}
+
+
+def resolve_str_to_obj(name: str):
+    """Resolve a block-type string from the config to its class."""
+    if name in _STR_TO_OBJ_OVERRIDES:
+        return _STR_TO_OBJ_OVERRIDES[name]
+    try:
+        return globals()[name]
+    except KeyError as exc:
+        raise AttributeError(
+            f"resolve_str_to_obj: unknown class name {name!r} "
+            f"(not defined in this module)"
+        ) from exc
+
+
+class VideoBaseAE(ModelMixin, ConfigMixin):
+    """Base class wiring diffusers' ``ModelMixin`` + ``ConfigMixin`` together."""
+
+    config_name = "config.json"
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+
+# ---------------------------------------------------------------------------
+# Primitive ops
+# ---------------------------------------------------------------------------
+
+
+def video_to_image(func):
+    """Wrap a 2D forward so it also accepts 5D ``(B, C, T, H, W)`` tensors.
+
+    Frames are folded into the batch dimension, processed independently, then
+    unfolded — letting purely spatial layers operate transparently on video.
+    """
+
+    def wrapper(self, x, *args, **kwargs):
+        if x.dim() == 5:
+            t = x.shape[2]
+            x = rearrange(x, "b c t h w -> (b t) c h w").contiguous()
+            x = func(self, x, *args, **kwargs)
+            x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+        else:
+            x = func(self, x, *args, **kwargs)
+        return x
+
+    return wrapper
+
+
+def nonlinearity(x: torch.Tensor) -> torch.Tensor:
+    """SiLU / swish activation used throughout the network."""
+    return x * torch.sigmoid(x)
+
+
+def cast_tuple(value, length: int = 1):
+    """Broadcast a scalar to a length-``length`` tuple; pass tuples through."""
+    return value if isinstance(value, (tuple, list)) else ((value,) * length)
+
+
+# ---------------------------------------------------------------------------
+# Convolutions
+# ---------------------------------------------------------------------------
+
+
+class Conv2d(nn.Conv2d):
+    """``nn.Conv2d`` that also accepts 5D video tensors (per-frame)."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int]] = 3,
+        stride: Union[int, Tuple[int]] = 1,
+        padding: Union[str, int, Tuple[int]] = 0,
+        dilation: Union[int, Tuple[int]] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = "zeros",
+        device=None,
+        dtype=None,
+    ) -> None:
+        super().__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding,
+            dilation,
+            groups,
+            bias,
+            padding_mode,
+            device,
+            dtype,
+        )
+
+    @video_to_image
+    def forward(self, x):
+        return super().forward(x)
+
+
+class CausalConv3d(nn.Module):
+    """A plain 3D convolution wrapped to keep input dtype aligned with weights.
+
+    The wrapper exists so spatial/temporal padding is configured consistently
+    across the network and so bf16/fp32 mismatches under AMP are resolved by
+    casting the input to the convolution's weight dtype.
+    """
+
+    def __init__(
+        self,
+        chan_in: int,
+        chan_out: int,
+        kernel_size: Union[int, Tuple[int, int, int]],
+        bias: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        self.kernel_size = cast_tuple(kernel_size, 3)
+        self.time_kernel_size = self.kernel_size[0]
+        self.chan_in = chan_in
+        self.chan_out = chan_out
+        stride = cast_tuple(kwargs.pop("stride", 1), 3)
+        padding = list(cast_tuple(kwargs.pop("padding", 0), 3))  # (T, H, W)
+        self.stride = stride
+        self.padding = padding
+        self.conv = nn.Conv3d(
+            chan_in,
+            chan_out,
+            self.kernel_size,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+        )
+
+    def forward(self, x):
+        if x.dtype is not self.conv.weight.dtype:
+            x = x.to(dtype=self.conv.weight.dtype)
+        return self.conv(x)
+
+
+# ---------------------------------------------------------------------------
+# Normalisation
+# ---------------------------------------------------------------------------
+
+
+class LayerNorm(nn.Module):
+    """Channel-last layer normalisation for both 4D and 5D tensors."""
+
+    def __init__(self, num_channels, eps=1e-6, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.norm = torch.nn.LayerNorm(num_channels, eps=eps, elementwise_affine=True)
+
+    def forward(self, x):
+        if x.dim() == 5:
+            x = rearrange(x, "b c t h w -> b t h w c")
+            x = self.norm(x)
+            x = rearrange(x, "b t h w c -> b c t h w")
+        else:
+            x = rearrange(x, "b c h w -> b h w c")
+            x = self.norm(x)
+            x = rearrange(x, "b h w c -> b c h w")
+        return x
+
+
+def Normalize(in_channels, num_groups=32, norm_type="groupnorm"):
+    """Build the normalisation layer selected by ``norm_type``."""
+    if norm_type == "groupnorm":
+        return torch.nn.GroupNorm(
+            num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
+        )
+    elif norm_type == "layernorm":
+        return LayerNorm(num_channels=in_channels, eps=1e-6)
+
+
+# ---------------------------------------------------------------------------
+# Up / down sampling
+# ---------------------------------------------------------------------------
+#
+# Spatial upsampling is sub-pixel convolution everywhere: a stride-1 conv that
+# expands channels by ``r**2`` followed by a pixel shuffle and ICNR
+# initialisation. This is the only spatial upsample mode the model ships.
+
+
+def spatial_pixel_shuffle_5d(x: torch.Tensor, upscale_factor: int = 2) -> torch.Tensor:
+    """Apply a spatial-only pixel shuffle to a 5D ``(B, C*r**2, T, H, W)`` tensor.
+
+    Returns ``(B, C, T, H*r, W*r)`` — equivalent to running
+    ``F.pixel_shuffle`` independently on every frame.
+    """
+    return rearrange(
+        x,
+        "b (c r1 r2) t h w -> b c t (h r1) (w r2)",
+        r1=upscale_factor,
+        r2=upscale_factor,
+    )
+
+
+def icnr_init_(weight: torch.Tensor, upscale_factor: int = 2) -> None:
+    """Initialise a sub-pixel conv weight with ICNR (Aitken et al., 2017).
+
+    The ``r**2`` output-channel groups that the pixel shuffle redistributes
+    into one ``r x r`` block are initialised identically, so the layer behaves
+    as nearest-neighbour upsampling at initialisation (no checkerboard).
+
+    Works for 2D and 3D kernels alike: only the leading channel dimension that
+    the shuffle reorders is touched, leaving any temporal axis untouched.
+    """
+    out_full = weight.shape[0]
+    r2 = upscale_factor ** 2
+    if out_full % r2 != 0:
+        raise ValueError(
+            f"icnr_init_: out_channels={out_full} must be divisible by "
+            f"upscale_factor**2={r2}"
+        )
+    sub = out_full // r2
+    base = torch.empty(sub, *weight.shape[1:], device=weight.device, dtype=weight.dtype)
+    nn.init.kaiming_normal_(base)
+    with torch.no_grad():
+        weight.copy_(base.repeat_interleave(r2, dim=0))
+
+
+class Upsample(nn.Module):
+    """2x spatial upsample (image branch) via sub-pixel convolution.
+
+    ``Conv2d(out=C*4) -> PixelShuffle(2)`` with ICNR init. When ``residual``
+    is set, a parameter-free nearest-neighbour identity skip is added on top,
+    scaled by a non-learnable ``residual_alpha`` buffer. The buffer is
+    non-persistent so a checkpoint saved without it still loads cleanly.
+    """
+
+    def __init__(self, in_channels, out_channels, residual: bool = False, **kwargs):
+        super().__init__()
+        self.upscale_factor = 2
+        self.conv = torch.nn.Conv2d(
+            in_channels,
+            out_channels * (self.upscale_factor ** 2),
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+        icnr_init_(self.conv.weight, self.upscale_factor)
+        if self.conv.bias is not None:
+            nn.init.zeros_(self.conv.bias)
+        self.shuffle = nn.PixelShuffle(self.upscale_factor)
+
+        self.residual = residual
+        self._in_channels = int(in_channels)
+        self._out_channels = int(out_channels)
+        if residual:
+            self.register_buffer("residual_alpha", torch.tensor(1.0), persistent=False)
+
+    def _identity_residual(self, x: torch.Tensor) -> torch.Tensor:
+        """Nearest 2x spatial upsample plus channel match to ``out_channels``.
+
+        The canonical configs use ``in_c == out_c`` for every up block, so only
+        the passthrough branch is exercised; the channel-average / repeat
+        branches are kept for configs whose stages change channel count.
+        """
+        in_c = x.shape[1]
+        out_c = self._out_channels
+        if in_c == out_c:
+            r = x
+        elif in_c > out_c:
+            if in_c % out_c != 0:
+                raise ValueError(
+                    f"residual channel-average failed: in_c={in_c} not divisible "
+                    f"by out_c={out_c}; configure matching channels or disable residual."
+                )
+            r = x.reshape(x.shape[0], out_c, in_c // out_c, *x.shape[2:]).mean(dim=2)
+        else:
+            if out_c % in_c != 0:
+                raise ValueError(
+                    f"residual channel-repeat failed: out_c={out_c} not divisible "
+                    f"by in_c={in_c}; configure matching channels or disable residual."
+                )
+            r = x.repeat_interleave(out_c // in_c, dim=1)
+        return F.interpolate(r, scale_factor=2.0, mode="nearest")
+
+    @video_to_image
+    def forward(self, x):
+        h = self.shuffle(self.conv(x))
+        if self.residual:
+            # Cast alpha to the main path's dtype so an AMP-bf16 forward is not
+            # silently promoted to fp32 by the float buffer.
+            alpha = self.residual_alpha.to(dtype=h.dtype, device=h.device)
+            h = h + alpha * self._identity_residual(x).to(dtype=h.dtype)
+        return h
+
+
+class Downsample(nn.Module):
+    """2x spatial downsample (encoder): asymmetric zero-pad then stride-2 conv."""
+
+    def __init__(self, in_channels, out_channels, **kwargs):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=2, padding=0
+        )
+
+    @video_to_image
+    def forward(self, x):
+        x = F.pad(x, (0, 1, 0, 1), mode="constant", value=0)
+        return self.conv(x)
+
+
+class Spatial2xTime2x3DDownsample(nn.Module):
+    """Joint 2x spatial + 2x temporal downsample via a stride-2 causal conv."""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.conv = CausalConv3d(
+            in_channels, out_channels, kernel_size=3, padding=0, stride=2
+        )
+
+    def forward(self, x):
+        x = F.pad(x, (0, 1, 0, 1, 1, 1), mode="constant", value=0)
+        return self.conv(x)
+
+
+class Spatial2xTime2x3DUpsample(nn.Module):
+    """Joint spatial + temporal 2x upsample.
+
+    The temporal axis is stretched with ``F.interpolate`` (cheap and smooth);
+    the spatial 2x is sub-pixel convolution — a ``CausalConv3d`` expanding
+    channels by 4 followed by a per-frame pixel shuffle — applied after the
+    temporal stretch. An optional parameter-free identity residual (see
+    :class:`Upsample`) is computed against the post-temporal tensor.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        t_interpolation="trilinear",
+        residual: bool = False,
+    ):
+        super().__init__()
+        self.t_interpolation = t_interpolation
+        self.upscale_factor = 2
+        # Produce 4x out_channels so the per-frame pixel shuffle recovers
+        # out_channels at 2x spatial resolution.
+        self.conv = CausalConv3d(
+            in_channels,
+            out_channels * (self.upscale_factor ** 2),
+            kernel_size=3,
+            padding=1,
+        )
+        icnr_init_(self.conv.conv.weight, self.upscale_factor)
+        if self.conv.conv.bias is not None:
+            nn.init.zeros_(self.conv.conv.bias)
+
+        self.residual = residual
+        self._in_channels = int(in_channels)
+        self._out_channels = int(out_channels)
+        if residual:
+            self.register_buffer("residual_alpha", torch.tensor(1.0), persistent=False)
+
+    def _identity_residual_5d(self, x: torch.Tensor) -> torch.Tensor:
+        """Nearest 2x spatial upsample on a 5D tensor plus channel match.
+
+        Called on the post-temporal tensor so its ``T`` axis already matches
+        the main path. See :meth:`Upsample._identity_residual` for the
+        channel-match rule.
+        """
+        in_c = x.shape[1]
+        out_c = self._out_channels
+        if in_c == out_c:
+            r = x
+        elif in_c > out_c:
+            if in_c % out_c != 0:
+                raise ValueError(
+                    f"residual channel-average failed: in_c={in_c} not divisible "
+                    f"by out_c={out_c}; configure matching channels or disable residual."
+                )
+            r = x.reshape(x.shape[0], out_c, in_c // out_c, *x.shape[2:]).mean(dim=2)
+        else:
+            if out_c % in_c != 0:
+                raise ValueError(
+                    f"residual channel-repeat failed: out_c={out_c} not divisible "
+                    f"by in_c={in_c}; configure matching channels or disable residual."
+                )
+            r = x.repeat_interleave(out_c // in_c, dim=1)
+        return F.interpolate(r, scale_factor=(1, 2, 2), mode="nearest")
+
+    def forward(self, x):
+        # Temporal upsample. A single-frame (image) input has nothing to
+        # stretch; an even frame count upsamples directly; an odd frame count
+        # is split so the leading frame is preserved before stretching.
+        if x.size(2) > 1:
+            T = x.size(2)
+            if T % 2 == 0:
+                x = F.interpolate(x, scale_factor=(2, 1, 1), mode=self.t_interpolation)
+            else:
+                x, x_tail = x[:, :, :1], x[:, :, 1:]
+                x_tail = F.interpolate(
+                    x_tail, scale_factor=(2, 1, 1), mode=self.t_interpolation
+                )
+                x = torch.cat([x, x_tail], dim=2)
+
+        # Spatial sub-pixel upsample. The temporal stretch is shared between
+        # the main and residual paths, so the residual only mirrors the
+        # spatial branch.
+        h = spatial_pixel_shuffle_5d(self.conv(x), self.upscale_factor)
+        if self.residual:
+            alpha = self.residual_alpha.to(dtype=h.dtype, device=h.device)
+            h = h + alpha * self._identity_residual_5d(x).to(dtype=h.dtype)
+        return h
+
+
+# ---------------------------------------------------------------------------
+# Residual blocks
+# ---------------------------------------------------------------------------
+
+
+class ResnetBlock2D(nn.Module):
+    """Pre-activation residual block over 2D (or per-frame 5D) features."""
+
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        norm_type,
+        dropout,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels, norm_type=norm_type)
+        self.conv1 = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        self.norm2 = Normalize(out_channels, norm_type=norm_type)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=1, padding=1
+        )
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(
+                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                )
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(
+                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                )
+
+    @video_to_image
+    def forward(self, x):
+        h = self.norm1(x)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+        return x + h
+
+
+class ResnetBlock3D(nn.Module):
+    """Pre-activation residual block over 3D features (causal convs)."""
+
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        dropout,
+        norm_type,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels if out_channels is None else out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels, norm_type=norm_type)
+        self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = Normalize(out_channels, norm_type=norm_type)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1)
+            else:
+                self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0)
+
+    def forward(self, x):
+        h = self.norm1(x)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+        return x + h
+
+
+# ---------------------------------------------------------------------------
+# Output containers
+# ---------------------------------------------------------------------------
+#
+# The encoder is deterministic, so ``latent_dist`` always carries a
+# ``DeterministicLatent`` (defined further down). The annotation is a forward
+# reference; it has no runtime effect.
+
+
+@dataclass
+class AutoencoderKLOutput(BaseOutput):
+    latent_dist: "DeterministicLatent"
+    extra_output: Optional[tuple] = None
+
+
+@dataclass
+class DecoderOutput(BaseOutput):
+    sample: torch.Tensor
+    extra_output: Optional[tuple] = None
+
+
+@dataclass
+class ForwardOutput(BaseOutput):
+    sample: torch.Tensor
+    latent_dist: "DeterministicLatent"
+    extra_output: Optional[tuple] = None
+    # The exact latent consumed by the decoder.
+    sampled_latent: Optional[torch.Tensor] = None
+
+
+# ---------------------------------------------------------------------------
+# Latent wrapper
+# ---------------------------------------------------------------------------
+
+
+class DeterministicLatent:
+    """Latent wrapper exposing the posterior API used by ``encode``.
+
+    The encoder emits the latent directly (no mean/log-variance split). This
+    wrapper exposes ``sample`` / ``mode`` (both returning the latent verbatim)
+    and a ``mean`` alias.
+    """
+
+    def __init__(self, x: torch.Tensor):
+        self.x = x
+        self.mean = x
+        self.parameters = x
+
+    def sample(self) -> torch.Tensor:
+        return self.x
+
+    def mode(self) -> torch.Tensor:
+        return self.x
+
+
+# ---------------------------------------------------------------------------
+# Encoder / decoder building blocks
+# ---------------------------------------------------------------------------
+
+
+def build_mid_layer(
+    layer_type: str, channels: int, dropout: float, norm_type: str
+) -> nn.Module:
+    """Build one entry of a mid stack from its config string.
+
+    ``"Identity"`` (and ``None``) build ``nn.Identity`` so a Res-Attn-Res mid
+    can be turned into an attention-free Res-pass-Res mid while preserving the
+    surrounding residual-block keys. Any other string resolves to a residual
+    block with matching in/out channels.
+    """
+    if layer_type is None or layer_type == "Identity":
+        return nn.Identity()
+    return resolve_str_to_obj(layer_type)(
+        in_channels=channels,
+        out_channels=channels,
+        dropout=dropout,
+        norm_type=norm_type,
+    )
+
+
+def pad_time_to_even(x: torch.Tensor) -> torch.Tensor:
+    """Prepend a repeated leading frame so the temporal axis becomes even.
+
+    This matches the temporal arithmetic of ``Spatial2xTime2x3DDownsample``:
+    folding the time axis by 2 needs an even length, and pre-pending (rather
+    than appending) the repeated frame keeps the fold causal.
+    """
+    if x.shape[2] % 2 == 0:
+        return x
+    first = x[:, :, :1].expand(-1, -1, 1, -1, -1)
+    return torch.cat([first, x], dim=2)
+
+
+class MotifDownBlock(nn.Module):
+    """Encoder down stage with a parameter-free Global Skip Connection.
+
+    The main path runs ``ResBlock(s) -> downsample -> ResBlock`` to reach
+    ``out_channels``. In parallel, the block's own input is folded by
+    space-to-channel and channel-averaged to ``out_channels``, then added to
+    the main path. ``hw`` stages fold space only; ``thw`` stages also fold time
+    by 2 (after padding odd ``T``) to match the main path's joint downsample.
+
+    For the canonical config the space-to-channel fold over-produces channels
+    by an integer factor, so the alignment is always an exact channel-average;
+    a zero-pad fallback covers configs where it does not.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_res_blocks: int = 2,
+        down_type: Literal["thw", "hw"] = "thw",
+        res_block: nn.Module = ResnetBlock3D,
+        dropout: float = 0.0,
+        norm_type: str = "layernorm",
+    ) -> None:
+        super().__init__()
+        assert num_res_blocks >= 2, "num_res_blocks too small (need >= 2)"
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.down_type = down_type
+
+        # Pre-downsample residual stack (channel-preserving).
+        self.res_block = nn.Sequential(
+            *[
+                res_block(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    dropout=dropout,
+                    norm_type=norm_type,
+                )
+                for _ in range(num_res_blocks - 1)
+            ]
+        )
+
+        if down_type == "thw":
+            self.down = Spatial2xTime2x3DDownsample(
+                in_channels=in_channels, out_channels=in_channels
+            )
+        elif down_type == "hw":
+            self.down = Downsample(in_channels=in_channels, out_channels=in_channels)
+        else:
+            raise ValueError(f"unknown down_type={down_type}")
+
+        # Post-downsample residual carrying the main path to out_channels.
+        self.out_res_block = res_block(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            dropout=dropout,
+            norm_type=norm_type,
+        )
+
+    def _space_to_channel(self, raw: torch.Tensor) -> torch.Tensor:
+        """Fold spatial (and, for ``thw``, temporal) resolution into channels."""
+        if self.down_type == "hw":
+            return rearrange(
+                raw, "b c t (h p1) (w p2) -> b (c p1 p2) t h w", p1=2, p2=2
+            )
+        raw = pad_time_to_even(raw)
+        return rearrange(
+            raw,
+            "b c (t p0) (h p1) (w p2) -> b (c p0 p1 p2) t h w",
+            p0=2,
+            p1=2,
+            p2=2,
+        )
+
+    def forward(self, h: torch.Tensor):
+        # The skip is tapped from the block input, before any transform.
+        h_in = h
+
+        h = self.res_block(h)
+        h = self.down(h)
+        h_main = self.out_res_block(h)
+
+        skip = self._space_to_channel(h_in)
+        out_c = h_main.shape[1]
+        skip_c = skip.shape[1]
+        if skip_c == out_c:
+            skip_aligned = skip
+        elif skip_c > out_c:
+            assert skip_c % out_c == 0, (
+                f"GSC channel-average failed: skip_c={skip_c} not divisible "
+                f"by out_c={out_c}.  Re-check base_channels and layer_type."
+            )
+            g = skip_c // out_c
+            skip_aligned = skip.reshape(
+                skip.shape[0], out_c, g, *skip.shape[2:]
+            ).mean(dim=2)
+        else:
+            # Zero-pad fallback for configs whose fold under-produces channels.
+            pad = torch.zeros(
+                skip.shape[0],
+                out_c - skip_c,
+                *skip.shape[2:],
+                dtype=skip.dtype,
+                device=skip.device,
+            )
+            skip_aligned = torch.cat([skip, pad], dim=1)
+
+        assert skip_aligned.shape[2:] == h_main.shape[2:], (
+            f"GSC skip shape {skip_aligned.shape} does not match main path "
+            f"{h_main.shape}."
+        )
+        return h_main + skip_aligned
+
+
+class MotifUpBlock(nn.Module):
+    """Decoder up stage: residual stack, spatial/temporal upsample, residual.
+
+    ``upsample_residual=True`` forwards the parameter-free identity skip into
+    the upsample module (see :class:`Upsample` / :class:`Spatial2xTime2x3DUpsample`).
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_res_blocks: int = 2,
+        up_type: Literal["thw", "hw"] = "thw",
+        res_block: nn.Module = ResnetBlock3D,
+        dropout: float = 0.0,
+        norm_type: str = "layernorm",
+        t_interpolation: str = "trilinear",
+        upsample_residual: bool = False,
+    ) -> None:
+        super().__init__()
+        assert num_res_blocks >= 2, "num_res_blocks too small (need >= 2)"
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.up_type = up_type
+
+        # Pre-upsample residual stack (channel-preserving).
+        self.res_block = nn.Sequential(
+            *[
+                res_block(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    dropout=dropout,
+                    norm_type=norm_type,
+                )
+                for _ in range(num_res_blocks - 1)
+            ]
+        )
+
+        if up_type == "thw":
+            self.up = Spatial2xTime2x3DUpsample(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                t_interpolation=t_interpolation,
+                residual=upsample_residual,
+            )
+        elif up_type == "hw":
+            self.up = Upsample(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                residual=upsample_residual,
+            )
+        else:
+            raise ValueError(f"unknown up_type={up_type}")
+
+        self.out_res_block = res_block(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            dropout=dropout,
+            norm_type=norm_type,
+        )
+
+    def forward(self, x: torch.Tensor):
+        x = self.res_block(x)
+        x = self.up(x)
+        return self.out_res_block(x)
+
+
+# ---------------------------------------------------------------------------
+# Encoder / decoder
+# ---------------------------------------------------------------------------
+
+
+class MotifEncoder(VideoBaseAE):
+    """Deterministic 3D causal encoder.
+
+    For a ``(B, 3, T, H, W)`` clip the stem is a stride-1 conv; five down
+    stages reduce resolution (3 spatial-only ``hw`` stages then 2 joint
+    ``thw`` stages), each adding its Global Skip Connection; a mid stack and a
+    final conv emit ``latent_dim`` channels directly at ``(ceil(T/4), H/32,
+    W/32)``.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        latent_dim: int = 128,
+        num_resblocks: int = 2,
+        dropout: float = 0.0,
+        input_type: Literal["video", "image"] = "video",
+        norm_type: str = "layernorm",
+        base_channels: List[int] = [96, 192, 384, 768, 768, 768],
+        mid_layers_type: List[str] = [
+            "ResnetBlock3D",
+            "Identity",
+            "ResnetBlock3D",
+        ],
+        down_layer_type: List[str] = ["hw", "hw", "hw", "thw", "thw"],
+        down_layer_res_type: List[str] = [
+            "ResnetBlock2D",
+            "ResnetBlock2D",
+            "ResnetBlock2D",
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+        ],
+    ) -> None:
+        super().__init__()
+        assert len(base_channels) == len(down_layer_type) + 1, (
+            f"base_channels must have one more entry than down_layer_type; "
+            f"got {len(base_channels)} vs {len(down_layer_type)}+1"
+        )
+        assert len(down_layer_type) == len(down_layer_res_type), (
+            f"down_layer_type and down_layer_res_type length mismatch: "
+            f"{len(down_layer_type)} vs {len(down_layer_res_type)}"
+        )
+        self.input_type = input_type
+        self.conv_in = Conv2d(3, base_channels[0], kernel_size=3, stride=1, padding=1)
+
+        self.down_blocks = nn.ModuleList()
+        for idx, (down_type, down_res_type) in enumerate(
+            zip(down_layer_type, down_layer_res_type)
+        ):
+            self.down_blocks.append(
+                MotifDownBlock(
+                    in_channels=base_channels[idx],
+                    out_channels=base_channels[idx + 1],
+                    num_res_blocks=num_resblocks,
+                    down_type=down_type,
+                    res_block=resolve_str_to_obj(down_res_type),
+                    dropout=dropout,
+                    norm_type=norm_type,
+                )
+            )
+
+        self.mid = nn.Sequential(
+            *[
+                build_mid_layer(t, base_channels[-1], dropout, norm_type)
+                for t in mid_layers_type
+            ]
+        )
+
+        self.norm_out = Normalize(base_channels[-1], norm_type=norm_type)
+        # conv_out emits latent_dim channels directly (no mean/log-var split).
+        if self.input_type == "video":
+            self.conv_out = CausalConv3d(
+                base_channels[-1], latent_dim, kernel_size=3, stride=1, padding=1
+            )
+        else:
+            self.conv_out = Conv2d(
+                base_channels[-1], latent_dim, kernel_size=3, stride=1, padding=1
+            )
+
+    def forward(self, x: torch.Tensor):
+        h = self.conv_in(x)
+        for down_block in self.down_blocks:
+            h = down_block(h)
+        h = self.mid(h)
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        return self.conv_out(h)
+
+
+class MotifDecoder(VideoBaseAE):
+    """Asymmetric decoder mirroring the encoder's stage layout.
+
+    The decoder is intentionally wider than the encoder
+    (``base_channels=[144, 288, 576, 1152, 1152, 1152]`` in the released
+    config). ``upsample_residual=True`` enables the parameter-free identity
+    skip in every spatial upsample, scaled by a non-persistent
+    ``residual_alpha`` buffer.
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        latent_dim: int = 128,
+        num_resblocks: int = 2,
+        dropout: float = 0.0,
+        input_type: Literal["video", "image"] = "video",
+        norm_type: str = "layernorm",
+        t_interpolation: str = "trilinear",
+        base_channels: List[int] = [96, 192, 384, 768, 768, 768],
+        up_layer_type: List[str] = ["hw", "hw", "hw", "thw", "thw"],
+        mid_layers_type: List[str] = [
+            "ResnetBlock3D",
+            "Identity",
+            "ResnetBlock3D",
+        ],
+        up_layer_res_type: List[str] = [
+            "ResnetBlock2D",
+            "ResnetBlock2D",
+            "ResnetBlock2D",
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+        ],
+        upsample_residual: bool = False,
+    ) -> None:
+        super().__init__()
+        assert len(base_channels) == len(up_layer_type) + 1, (
+            f"base_channels must have one more entry than up_layer_type; "
+            f"got {len(base_channels)} vs {len(up_layer_type)}+1"
+        )
+        assert len(up_layer_type) == len(up_layer_res_type), (
+            f"up_layer_type and up_layer_res_type length mismatch: "
+            f"{len(up_layer_type)} vs {len(up_layer_res_type)}"
+        )
+        self.input_type = input_type
+        if self.input_type == "video":
+            self.conv_in = CausalConv3d(
+                latent_dim, base_channels[-1], kernel_size=3, stride=1, padding=1
+            )
+        else:
+            self.conv_in = Conv2d(
+                latent_dim, base_channels[-1], kernel_size=3, stride=1, padding=1
+            )
+
+        self.mid = nn.Sequential(
+            *[
+                build_mid_layer(t, base_channels[-1], dropout, norm_type)
+                for t in mid_layers_type
+            ]
+        )
+
+        self.up_blocks = nn.ModuleList()
+        # Iterate from the deepest stage back up to the shallowest.
+        for idx, (up_type, up_res_type) in enumerate(
+            zip(reversed(up_layer_type), reversed(up_layer_res_type))
+        ):
+            idx = len(up_layer_type) - idx
+            self.up_blocks.append(
+                MotifUpBlock(
+                    in_channels=base_channels[idx],
+                    out_channels=base_channels[idx - 1],
+                    num_res_blocks=num_resblocks,
+                    up_type=up_type,
+                    res_block=resolve_str_to_obj(up_res_type),
+                    t_interpolation=t_interpolation,
+                    dropout=dropout,
+                    norm_type=norm_type,
+                    upsample_residual=upsample_residual,
+                )
+            )
+
+        self.norm_out = Normalize(base_channels[0], norm_type=norm_type)
+        self.conv_out = Conv2d(base_channels[0], 3, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: torch.Tensor):
+        h = self.conv_in(z)
+        h = self.mid(h)
+        for up_block in self.up_blocks:
+            h = up_block(h)
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        return self.conv_out(h)
+
+
+# ---------------------------------------------------------------------------
+# Top-level model
+# ---------------------------------------------------------------------------
+
+
+class MotifVAE(VideoBaseAE):
+    """3D causal video VAE with a 128-channel deterministic latent.
+
+    Compression is 4x temporal and 32x spatial. The encoder (width 96) uses a
+    parameter-free Global Skip Connection at every stage; the decoder is wider
+    (width 144).
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        latent_dim: int = 128,
+        base_channels: List[int] = [96, 192, 384, 768, 768, 768],
+        decoder_base_channels: Optional[List[int]] = [144, 288, 576, 1152, 1152, 1152],
+        decoder_variant: Literal["V4"] = "V4",
+        layer_type: List[str] = ["hw", "hw", "hw", "thw", "thw"],
+        layer_res_type: List[str] = [
+            "ResnetBlock2D",
+            "ResnetBlock2D",
+            "ResnetBlock2D",
+            "ResnetBlock3D",
+            "ResnetBlock3D",
+        ],
+        encoder_num_resblocks: int = 2,
+        decoder_num_resblocks: int = 2,
+        dropout: float = 0.0,
+        norm_type: str = "layernorm",
+        t_interpolation: str = "trilinear",
+        input_type: Literal["video", "image"] = "video",
+        mid_layers_type: List[str] = [
+            "ResnetBlock3D",
+            "Identity",
+            "ResnetBlock3D",
+        ],
+        scale: Optional[List[float]] = None,
+        shift: Optional[List[float]] = None,
+        upsample_residual: bool = False,
+    ) -> None:
+        super().__init__()
+        self.use_tiling = False
+        self.t_chunk_enc = 16
+        self.t_chunk_dec = 4
+        self.use_quant_layer = False
+
+        if scale is None:
+            scale = [0.18215] * latent_dim
+        if shift is None:
+            shift = [0.0] * latent_dim
+        if decoder_base_channels is None:
+            decoder_base_channels = [144, 288, 576, 1152, 1152, 1152]
+
+        self.encoder = MotifEncoder(
+            latent_dim=latent_dim,
+            base_channels=base_channels,
+            num_resblocks=encoder_num_resblocks,
+            dropout=dropout,
+            norm_type=norm_type,
+            down_layer_type=layer_type,
+            down_layer_res_type=layer_res_type,
+            input_type=input_type,
+            mid_layers_type=mid_layers_type,
+        )
+
+        if decoder_variant != "V4":
+            raise ValueError(
+                f"unknown decoder_variant={decoder_variant!r}; expected 'V4'"
+            )
+        self.decoder = MotifDecoder(
+            latent_dim=latent_dim,
+            base_channels=decoder_base_channels,
+            num_resblocks=decoder_num_resblocks,
+            dropout=dropout,
+            norm_type=norm_type,
+            t_interpolation=t_interpolation,
+            up_layer_type=layer_type,
+            up_layer_res_type=layer_res_type,
+            input_type=input_type,
+            mid_layers_type=mid_layers_type,
+            upsample_residual=upsample_residual,
+        )
+
+    def get_encoder(self):
+        return [self.encoder]
+
+    def get_decoder(self):
+        return [self.decoder]
+
+    def encode(self, x, **kwargs):
+        # Deterministic encoder: DeterministicLatent exposes the sample()/mode()
+        # posterior API over the latent.
+        h = self.encoder(x)
+        posterior = DeterministicLatent(h)
+        return AutoencoderKLOutput(latent_dist=posterior, extra_output=None)
+
+    def decode(self, z, **kwargs):
+        dec = self.decoder(z)
+        return DecoderOutput(sample=dec, extra_output=None)
+
+    def forward(self, input, sample_posterior=True):
+        posterior = self.encode(input).latent_dist
+        z = posterior.sample() if sample_posterior else posterior.mode()
+        dec = self.decode(z).sample
+        return ForwardOutput(
+            sample=dec,
+            latent_dist=posterior,
+            sampled_latent=z,
+            extra_output=None,
+        )
+
+    def get_last_layer(self):
+        if hasattr(self.decoder.conv_out, "conv"):
+            return self.decoder.conv_out.conv.weight
+        return self.decoder.conv_out.weight