``, take its gradient w.r.t. the clip, and
+scatter-add into the full 81-frame pixel-gradient buffer.
+
+Stage B2 exposes two decoder helpers. The released feature-matching model calls
+``replay_decoder_vjp``, which detaches temporal caches between latent frames.
+The optional ``replay_decoder_vjp_exact`` helper preserves temporal derivatives
+using recurrent boundary replay; it is not selected by that training path.
+
+Frozen parameter gradients remain disabled; replay builds only the input and
+explicit recurrent-state graphs required for the requested VJP.
+"""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from typing import Callable, List, Optional, Tuple
+
+import torch
+from torch.utils.checkpoint import checkpoint
+
+from video_fd.replay.scatter_frame_grads import gather_frames, scatter_add_frames, new_frame_grad_buffer
+from video_fd.replay.wan_recurrent_decoder_vjp import (
+ replay_wan_decoder_vjp_exact,
+ wan_recurrent_components,
+)
+
+
+def _amp_ctx(use_amp: bool, amp_dtype: torch.dtype, device_type: str):
+ if use_amp and device_type == "cuda":
+ return torch.autocast(device_type="cuda", dtype=amp_dtype)
+ return nullcontext()
+
+
+def replay_encoder_vjp(
+ encoder,
+ v_clip_detached: torch.Tensor, # [B, T, 3, H, W] in [0,1], detached
+ grad_feat_proj: torch.Tensor, # [B, S, D_proj], dL/d(projected feature)
+ projector: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+ use_amp: bool = False,
+ amp_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+ """One-clip encoder VJP: returns ``g_clip`` = (dF_proj/d v_clip)^T grad_feat_proj."""
+ # Frozen encoder parameters still allow derivatives with respect to pixels.
+ # Recreate only this clip's graph; grad_feat_proj is the cached cotangent.
+ v_leaf = v_clip_detached.detach().clone().requires_grad_(True)
+ with _amp_ctx(use_amp, amp_dtype, v_leaf.device.type):
+ feat = encoder.forward_sparse(v_leaf) # [B, S, D_raw]
+ feat_proj = projector(feat) if projector is not None else feat
+ feat_proj = feat_proj.float()
+ g = grad_feat_proj.detach().to(feat_proj.dtype)
+ if g.shape != feat_proj.shape:
+ raise RuntimeError(
+ f"encoder VJP shape mismatch: feat_proj {tuple(feat_proj.shape)} vs "
+ f"grad {tuple(g.shape)}"
+ )
+ # The inner product injects dL/d feature without recomputing the MMD loss:
+ # its pixel derivative is J_encoder^T (dL/d feature), the required VJP.
+ replay_loss = (feat_proj * g).sum()
+ (g_clip,) = torch.autograd.grad(replay_loss, v_leaf, retain_graph=False, create_graph=False)
+ return g_clip
+
+
+def replay_video_encoder_vjp(
+ video81_detached: torch.Tensor, # [B, 81, 3, H, W] in [0,1], detached
+ specs_and_grads: List[Tuple[object, object, torch.Tensor]],
+ use_amp: bool = False,
+ amp_dtype: torch.dtype = torch.bfloat16,
+) -> torch.Tensor:
+ """Accumulate per-clip encoder VJPs into a full pixel-frame gradient buffer.
+
+ ``specs_and_grads`` is a list of ``(encoder, spec_or_projector_tuple,
+ grad_feat_proj)`` where the middle element is ``(spec, projector)``. Each
+ clip's gradient is scattered into the 81-frame buffer at ``spec.frame_indices``
+ (accumulating overlaps).
+ """
+ # Multiple clips/encoders can touch the same frame. Sum their derivatives
+ # in video coordinates before replaying the VAE decoder.
+ g_video = new_frame_grad_buffer(video81_detached)
+ for encoder, (spec, projector), grad_feat_proj in specs_and_grads:
+ idx = spec.frame_indices
+ v_clip = gather_frames(video81_detached, idx, time_dim=1) # [B, 16, 3, H, W]
+ g_clip = replay_encoder_vjp(
+ encoder, v_clip, grad_feat_proj, projector=projector,
+ use_amp=use_amp, amp_dtype=amp_dtype,
+ )
+ scatter_add_frames(g_video, idx, g_clip, time_dim=1)
+ return g_video
+
+
+@torch.enable_grad()
+def replay_decoder_vjp(
+ vae,
+ z_detached: torch.Tensor, # [B, T_lat, C, H, W] (repo latent layout)
+ g_video01: torch.Tensor, # [B, T_pix, 3, H, W] dL/d(pixels in [0,1])
+ independent_first_frame: bool = True,
+) -> torch.Tensor:
+ """Per-latent-frame decode VJP: ``grad_z = (d pixels01 / d z)^T g_video01``.
+
+ Re-decodes one latent frame at a time with grad, carrying the VAE temporal
+ cache by value but detaching it between frames so peak memory is ~one frame's
+ decode graph. Frames with zero gradient still advance the cache (needed for
+ exact values of later frames) but contribute no latent gradient.
+
+ Temporal cache detachment truncates cross-frame derivatives. This historical
+ path preserves decode values, but is not the full decoder Jacobian.
+ """
+ vae_model = vae.model
+ zdim = vae_model.z_dim
+ device = z_detached.device
+ dtype = z_detached.dtype
+ mean = vae.mean.to(device=device, dtype=dtype).view(1, zdim, 1, 1, 1)
+ inv_std = (1.0 / vae.std.to(device=device, dtype=dtype)).view(1, zdim, 1, 1, 1)
+
+ z_cf = z_detached.detach().permute(0, 2, 1, 3, 4).contiguous() # [B, zdim, T_lat, H, W]
+ T = z_cf.shape[2]
+ grad_z_cf = torch.zeros_like(z_cf, dtype=torch.float32)
+
+ # Recompute one decode segment, inject its pixel cotangent, and release the
+ # segment graph. The returned latent cotangent drives generator replay next.
+ vae_model.clear_cache()
+ px_cursor = 0
+ for i in range(T):
+ vae_model._conv_idx = [0]
+ z_i = z_cf[:, :, i:i + 1].clone().requires_grad_(True)
+ z_i_unscaled = z_i / inv_std + mean
+ x_i = vae_model.conv2(z_i_unscaled)
+ out_i = vae_model.decoder(x_i, feat_cache=vae_model._feat_map, feat_idx=vae_model._conv_idx)
+ out_i = out_i.float().clamp(-1.0, 1.0)
+ pix_i = out_i.shape[2]
+ pixels01_i = ((out_i + 1.0) * 0.5).permute(0, 2, 1, 3, 4) # [B, pix_i, 3, H, W]
+ end = min(px_cursor + pix_i, g_video01.shape[1])
+ # Compute the surrogate in float32 so a large pixel-frame gradient does
+ # not overflow bf16 (the encoder VJP can be large).
+ g_i = g_video01[:, px_cursor:end].float()
+ if g_i.shape[1] == pixels01_i.shape[1] and torch.isfinite(g_i).all() and g_i.abs().sum() > 0:
+ surrogate_i = (pixels01_i.float() * g_i.detach()).sum()
+ (grad_zi,) = torch.autograd.grad(surrogate_i, z_i, retain_graph=False)
+ grad_z_cf[:, :, i:i + 1] += grad_zi.float()
+ # The gradient has already been accumulated. Keeping its graph
+ # root alive until the next decoder call would retain the old
+ # frame leaf and output storage without affecting any later math.
+ del surrogate_i, grad_zi
+ px_cursor += pix_i
+ vae_model._feat_map = [c.detach() if torch.is_tensor(c) else c for c in vae_model._feat_map]
+ # The detached temporal cache above is the only state needed by the
+ # next frame. Release the consumed frame graph and views before that
+ # frame's decoder reaches its allocation peak.
+ del pixels01_i, out_i, x_i, z_i_unscaled, z_i, g_i
+ vae_model.clear_cache()
+ return grad_z_cf.permute(0, 2, 1, 3, 4) # [B, T_lat, C, H, W]
+
+
+@torch.enable_grad()
+def _replay_generic_decoder_vjp_exact(
+ vae,
+ z_detached: torch.Tensor,
+ g_video01: torch.Tensor,
+ *,
+ expected_pixels01: torch.Tensor | None = None,
+) -> torch.Tensor:
+ """Exact full-video decoder VJP with activation replay.
+
+ Wan's causal VAE carries temporal feature caches between all 21 latent
+ frames. Detaching those caches changes the MMD gradient. This path
+ reentrantly checkpoints the complete frozen decode instead: the first
+ forward runs without an autograd graph, then backward recomputes the exact
+ graph while saved activations are moved to pageable CPU memory. Every
+ cross-frame dependency remains in the autograd graph.
+ """
+
+ if z_detached.ndim != 5 or g_video01.ndim != 5:
+ raise ValueError("decoder VJP expects [B,T,C,H,W] latent and pixel tensors")
+ if z_detached.device.type != g_video01.device.type:
+ raise ValueError("decoder latent and pixel cotangent must share a device")
+ if not torch.isfinite(g_video01).all():
+ raise ValueError("decoder pixel cotangent contains non-finite values")
+ z_leaf = z_detached.detach().clone().requires_grad_(True)
+
+ def decode_pixels01(latent: torch.Tensor) -> torch.Tensor:
+ # With reentrant checkpointing the first call is under no_grad; only
+ # the backward recompute saves tensors. Keeping this hook *inside*
+ # the function is required so those saved activations move to CPU.
+ with torch.autograd.graph.save_on_cpu(
+ pin_memory=False, device_type=latent.device.type
+ ):
+ with _amp_ctx(True, torch.bfloat16, latent.device.type):
+ decoded = vae.decode_to_pixel(latent, use_cache=False)
+ pixels01 = ((decoded + 1.0) * 0.5).clamp(0.0, 1.0)
+ return pixels01
+
+ pixels01 = checkpoint(
+ decode_pixels01,
+ z_leaf,
+ use_reentrant=True,
+ preserve_rng_state=False,
+ )
+ if tuple(pixels01.shape) != tuple(g_video01.shape):
+ raise ValueError(
+ f"decoder cotangent shape {tuple(g_video01.shape)} differs from "
+ f"decoded pixels {tuple(pixels01.shape)}"
+ )
+ surrogate = (pixels01.float() * g_video01.detach().float()).sum()
+ if expected_pixels01 is not None and not torch.equal(
+ pixels01.detach(), expected_pixels01
+ ):
+ raise RuntimeError("generic decoder replay differs from expected pixels")
+ surrogate.backward()
+ gradient = z_leaf.grad
+ if gradient is None:
+ raise RuntimeError("exact decoder VJP did not produce a latent gradient")
+ if not torch.isfinite(gradient).all():
+ raise RuntimeError("exact decoder VJP produced non-finite latent gradients")
+ return gradient.detach()
+
+
+@torch.enable_grad()
+def replay_decoder_vjp_exact(
+ vae,
+ z_detached: torch.Tensor,
+ g_video01: torch.Tensor,
+ *,
+ expected_pixels01: torch.Tensor | None = None,
+) -> torch.Tensor:
+ """Exact decoder VJP with bounded-memory Wan recurrent BPTT.
+
+ The official Wan decoder uses an explicit 32-cache recurrent state, K=4
+ pageable-CPU boundaries, and future-state cotangent propagation. This
+ retains the complete temporal Jacobian without materializing or offloading
+ the whole 21-frame graph. Stateless toy/non-Wan decoders retain the
+ generic reentrant checkpoint fallback for test and API compatibility.
+ """
+
+ if z_detached.ndim != 5 or g_video01.ndim != 5:
+ raise ValueError("decoder VJP expects [B,T,C,H,W] latent and pixel tensors")
+ if z_detached.device != g_video01.device:
+ raise ValueError("decoder latent and pixel cotangent must share a device")
+ if not torch.isfinite(g_video01).all():
+ raise ValueError("decoder pixel cotangent contains non-finite values")
+ # This optional Wan path carries future-state cotangents across temporal
+ # boundaries; the generic fallback preserves the full decode graph as well.
+ if wan_recurrent_components(vae) is not None:
+ return replay_wan_decoder_vjp_exact(
+ vae,
+ z_detached,
+ g_video01,
+ expected_pixels01=expected_pixels01,
+ )
+ return _replay_generic_decoder_vjp_exact(
+ vae,
+ z_detached,
+ g_video01,
+ expected_pixels01=expected_pixels01,
+ )
diff --git a/14B/code/video_fd/replay/scatter_frame_grads.py b/14B/code/video_fd/replay/scatter_frame_grads.py
new file mode 100644
index 0000000000000000000000000000000000000000..52b3eb253d42c849d970dabc8ef5315ae82f1fa8
--- /dev/null
+++ b/14B/code/video_fd/replay/scatter_frame_grads.py
@@ -0,0 +1,22 @@
+"""Frame gather / scatter-add for replay over arbitrary selected frame indices.
+
+The stride-4 sampler selects interleaved pixel frames; the replay must (a) gather
+those frames into a 16-frame clip and (b) scatter the clip's pixel gradient back
+into the full 81-frame buffer, **accumulating** when a frame is selected by
+multiple clips / crops / encoders (spec section 11.3).
+"""
+
+from __future__ import annotations
+
+import torch
+
+# Re-export the canonical implementations from the sampler module so there is a
+# single source of truth for the gather/scatter rule.
+from video_fd.samplers.stride4_clip_sampler import gather_frames, scatter_add_frames
+
+__all__ = ["gather_frames", "scatter_add_frames", "new_frame_grad_buffer"]
+
+
+def new_frame_grad_buffer(video: torch.Tensor) -> torch.Tensor:
+ """Zero gradient buffer matching a ``[B, T, 3, H, W]`` video (float32)."""
+ return torch.zeros_like(video, dtype=torch.float32)
diff --git a/14B/code/video_fd/replay/wan_recurrent_decoder_vjp.py b/14B/code/video_fd/replay/wan_recurrent_decoder_vjp.py
new file mode 100644
index 0000000000000000000000000000000000000000..71b0aef6a5c3627fe878d87e04abd4334675c0af
--- /dev/null
+++ b/14B/code/video_fd/replay/wan_recurrent_decoder_vjp.py
@@ -0,0 +1,579 @@
+"""Bounded-memory exact recurrent VJP for the official Wan VAE decoder.
+
+Wan's decoder mutates a Python feature-cache list while consuming one latent
+frame at a time. This module makes that list an explicit recurrent state,
+stores only K=4 boundaries on CPU, and propagates the future-state cotangent
+back through every transition. The resulting VJP is the full temporal
+Jacobian, not the historical detached-cache approximation.
+"""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from typing import Any
+
+import torch
+
+
+CHUNK_LATENT_FRAMES = 4
+REP_SENTINEL = "Rep"
+
+
+def _amp(device_type: str):
+ if device_type == "cuda":
+ return torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+ return nullcontext()
+
+
+def wan_recurrent_components(vae: Any) -> tuple[Any, Any] | None:
+ """Duck-type the official Wan wrapper without creating an import cycle."""
+
+ model = getattr(vae, "model", None)
+ decoder = getattr(model, "decoder", None)
+ if (
+ model is None
+ or decoder is None
+ or not hasattr(model, "conv2")
+ or not hasattr(model, "z_dim")
+ or not hasattr(decoder, "num_res_blocks")
+ or not hasattr(decoder, "dim_mult")
+ or not hasattr(decoder, "temperal_upsample")
+ or not hasattr(vae, "mean")
+ or not hasattr(vae, "std")
+ ):
+ return None
+ return model, decoder
+
+
+def _active_cache_slots(decoder: Any) -> int:
+ stages = len(tuple(decoder.dim_mult))
+ residual_convs = 4 + stages * (int(decoder.num_res_blocks) + 1) * 2
+ temporal_resamplers = sum(
+ bool(value) for value in tuple(decoder.temperal_upsample)[: stages - 1]
+ )
+ # decoder.conv1 + residual convs + temporal resamplers + decoder.head
+ return 1 + residual_convs + temporal_resamplers + 1
+
+
+def _is_formal_480p(model: Any, decoder: Any, latent: torch.Tensor) -> bool:
+ return (
+ int(model.z_dim) == 16
+ and int(decoder.dim) == 96
+ and int(decoder.num_res_blocks) == 2
+ and tuple(decoder.dim_mult) == (1, 2, 4, 4)
+ and tuple(bool(v) for v in decoder.temperal_upsample)
+ == (True, True, False)
+ and tuple(latent.shape) == (1, 21, 16, 60, 104)
+ )
+
+
+def _model_cache_signature(model: Any) -> tuple[Any, ...]:
+ """Identity-only proof that the functional path never touched globals."""
+
+ values = getattr(model, "_feat_map", None)
+ index = getattr(model, "_conv_idx", None)
+ value_signature = None
+ if isinstance(values, list):
+ value_signature = tuple(
+ ("tensor", id(value)) if torch.is_tensor(value) else ("value", value)
+ for value in values
+ )
+ index_signature = tuple(index) if isinstance(index, list) else index
+ return id(values), value_signature, id(index), index_signature
+
+
+def _formal_shapes(
+ *, batch: int, height: int, width: int, frames_consumed: int
+) -> dict[int, tuple[int, ...] | str]:
+ temporal = 1 if frames_consumed == 1 else 2
+ result: dict[int, tuple[int, ...] | str] = {
+ 0: (batch, 16, temporal, height, width),
+ 12: (batch, 192, temporal, height * 2, width * 2),
+ }
+ result.update(
+ {
+ index: (batch, 384, temporal, height, width)
+ for index in range(1, 11)
+ }
+ )
+ result[11] = (
+ REP_SENTINEL
+ if frames_consumed == 1
+ else (batch, 384, 2, height, width)
+ )
+ result.update(
+ {
+ index: (batch, 384, temporal, height * 2, width * 2)
+ for index in range(13, 18)
+ }
+ )
+ result[18] = (
+ REP_SENTINEL
+ if frames_consumed == 1
+ else (batch, 384, 2, height * 2, width * 2)
+ )
+ result.update(
+ {
+ index: (batch, 192, temporal, height * 4, width * 4)
+ for index in range(19, 25)
+ }
+ )
+ result.update(
+ {
+ index: (batch, 96, temporal, height * 8, width * 8)
+ for index in range(25, 32)
+ }
+ )
+ return result
+
+
+def _formal_cache_dtype(index: int) -> torch.dtype:
+ if index not in range(32):
+ raise ValueError(f"formal Wan cache index out of range: {index}")
+ return torch.bfloat16 if index in (0, 11, 18) else torch.float32
+
+
+def _validate_state(
+ state: tuple[Any, ...],
+ *,
+ slots: int,
+ frames_consumed: int,
+ batch: int,
+ height: int,
+ width: int,
+ formal: bool,
+ device_type: str | None = None,
+) -> None:
+ if len(state) != slots:
+ raise RuntimeError(f"Wan state has {len(state)} slots, expected {slots}")
+ for index, value in enumerate(state):
+ if torch.is_tensor(value):
+ if value.ndim != 5 or value.shape[0] != batch:
+ raise RuntimeError(
+ f"Wan cache slot {index} has invalid shape {tuple(value.shape)}"
+ )
+ if not value.is_floating_point():
+ raise TypeError(f"Wan cache slot {index} is not floating point")
+ if device_type is not None and value.device.type != device_type:
+ raise RuntimeError(
+ f"Wan cache slot {index} is on {value.device.type}, "
+ f"expected {device_type}"
+ )
+ elif value is None:
+ if frames_consumed != 0:
+ raise RuntimeError(f"Wan cache slot {index} unexpectedly remained empty")
+ elif value == REP_SENTINEL:
+ if frames_consumed != 1 or (formal and index not in (11, 18)):
+ raise RuntimeError(
+ f"unexpected Rep sentinel in slot {index} after "
+ f"{frames_consumed} frames"
+ )
+ else:
+ raise TypeError(
+ f"unsupported Wan cache value {type(value).__name__} in slot {index}"
+ )
+
+ if not formal:
+ return
+ if slots != 32:
+ raise RuntimeError(f"formal Wan decoder needs 32 cache slots, got {slots}")
+ if frames_consumed == 0:
+ if any(value is not None for value in state):
+ raise RuntimeError("initial formal Wan state is not empty")
+ return
+ expected = _formal_shapes(
+ batch=batch,
+ height=height,
+ width=width,
+ frames_consumed=frames_consumed,
+ )
+ if set(expected) != set(range(32)):
+ raise AssertionError("formal cache layout does not cover all 32 slots")
+ for index, specification in expected.items():
+ value = state[index]
+ if specification == REP_SENTINEL:
+ if value != REP_SENTINEL:
+ raise RuntimeError(f"formal Wan slot {index} must be Rep")
+ continue
+ if not torch.is_tensor(value) or tuple(value.shape) != specification:
+ actual = tuple(value.shape) if torch.is_tensor(value) else value
+ raise RuntimeError(
+ f"formal Wan slot {index} has {actual}, expected {specification}"
+ )
+ expected_dtype = _formal_cache_dtype(index)
+ if value.dtype != expected_dtype:
+ raise TypeError(
+ f"formal Wan slot {index} must be {expected_dtype}, "
+ f"got {value.dtype}"
+ )
+
+
+def _state_to_cpu(state: tuple[Any, ...], *, formal: bool) -> tuple[Any, ...]:
+ result: list[Any] = []
+ for index, value in enumerate(state):
+ if torch.is_tensor(value):
+ if formal and value.dtype != _formal_cache_dtype(index):
+ raise TypeError(
+ f"formal boundary slot {index} must be "
+ f"{_formal_cache_dtype(index)}, got {value.dtype}"
+ )
+ result.append(
+ value.detach().to(device="cpu", non_blocking=False).contiguous()
+ )
+ else:
+ result.append(value)
+ return tuple(result)
+
+
+def _state_to_device(
+ state: tuple[Any, ...], *, device: torch.device
+) -> tuple[Any, ...]:
+ return tuple(
+ value.to(device=device, non_blocking=False).detach()
+ if torch.is_tensor(value)
+ else value
+ for value in state
+ )
+
+
+def _state_as_leaves(state: tuple[Any, ...]) -> tuple[Any, ...]:
+ return tuple(
+ value.detach().requires_grad_(True) if torch.is_tensor(value) else value
+ for value in state
+ )
+
+
+def _transition(
+ *,
+ model: Any,
+ preconv_frame: torch.Tensor,
+ state: tuple[Any, ...],
+ slots: int,
+) -> tuple[torch.Tensor, tuple[Any, ...]]:
+ if preconv_frame.ndim != 5 or preconv_frame.shape[2] != 1:
+ raise ValueError("Wan transition requires one pre-convolved frame")
+ local_cache = list(state)
+ local_index = [0]
+ with _amp(preconv_frame.device.type):
+ decoded_cf = model.decoder(
+ preconv_frame,
+ feat_cache=local_cache,
+ feat_idx=local_index,
+ )
+ # Preserve WanVAEWrapper exactly: FP32 [-1,1] clamp, repository
+ # time-first layout, then the [0,1] affine and clamp.
+ decoded = decoded_cf.float().clamp(-1.0, 1.0).permute(0, 2, 1, 3, 4)
+ pixels01 = ((decoded + 1.0) * 0.5).clamp(0.0, 1.0)
+ if local_index[0] != slots:
+ raise RuntimeError(
+ f"Wan transition consumed {local_index[0]} cache slots, expected {slots}"
+ )
+ if len(local_cache) != slots:
+ raise RuntimeError("Wan transition changed local cache-list length")
+ return pixels01, tuple(local_cache)
+
+
+def _full_preconv(vae: Any, latent: torch.Tensor) -> torch.Tensor:
+ model = vae.model
+ channels = int(model.z_dim)
+ if latent.shape[2] != channels:
+ raise ValueError(
+ f"Wan latent has {latent.shape[2]} channels, expected {channels}"
+ )
+ mean = vae.mean.to(device=latent.device, dtype=latent.dtype).view(
+ 1, channels, 1, 1, 1
+ )
+ inverse_std = (
+ 1.0 / vae.std.to(device=latent.device, dtype=latent.dtype)
+ ).view(1, channels, 1, 1, 1)
+ unscaled = latent.permute(0, 2, 1, 3, 4) / inverse_std + mean
+ return model.conv2(unscaled)
+
+
+def _formal_pixel_span(frame_index: int) -> tuple[int, int]:
+ if frame_index == 0:
+ return 0, 1
+ return 1 + 4 * (frame_index - 1), 1 + 4 * frame_index
+
+
+def replay_wan_decoder_vjp_exact(
+ vae: Any,
+ z_detached: torch.Tensor,
+ g_video01: torch.Tensor,
+ *,
+ expected_pixels01: torch.Tensor | None = None,
+) -> torch.Tensor:
+ """Run exact K=4 recurrent BPTT through an official Wan VAE decoder."""
+
+ components = wan_recurrent_components(vae)
+ if components is None:
+ raise TypeError("recurrent Wan VJP received a non-Wan decoder")
+ model, decoder = components
+ if z_detached.shape[0] != 1:
+ raise ValueError("Wan recurrent decoder VJP requires batch size 1")
+ if model.training or decoder.training:
+ raise RuntimeError("Wan recurrent decoder VJP requires eval mode")
+ if any(parameter.requires_grad for parameter in model.parameters()):
+ raise RuntimeError("Wan recurrent decoder VJP requires frozen parameters")
+ slots = _active_cache_slots(decoder)
+ formal = _is_formal_480p(model, decoder, z_detached)
+ if int(model.z_dim) == 16 and int(decoder.dim) == 96 and slots != 32:
+ raise RuntimeError(f"production Wan exposed {slots} slots, expected 32")
+ if formal and (
+ z_detached.dtype != torch.bfloat16 or g_video01.dtype != torch.float32
+ ):
+ raise TypeError("formal 480P VJP needs BF16 latent and FP32 cotangent")
+ if expected_pixels01 is not None:
+ if expected_pixels01.device != z_detached.device:
+ raise ValueError("expected pixels must remain on the latent device")
+ if tuple(expected_pixels01.shape) != tuple(g_video01.shape):
+ raise ValueError("expected pixels differ from cotangent shape")
+ if not torch.isfinite(expected_pixels01).all():
+ raise ValueError("expected pixels contain non-finite values")
+
+ global_signature = _model_cache_signature(model)
+ batch, latent_frames, _, height, width = z_detached.shape
+ initial_state = tuple(None for _ in range(slots))
+ _validate_state(
+ initial_state,
+ slots=slots,
+ frames_consumed=0,
+ batch=batch,
+ height=height,
+ width=width,
+ formal=formal,
+ )
+
+ # Official decode scales and calls conv2 once at full temporal shape.
+ # Preserve that operation, collect dL/dx_i recurrently, then run one exact
+ # full-shape conv2+scale VJP at the end.
+ with torch.no_grad(), _amp(z_detached.device.type):
+ preconv = _full_preconv(vae, z_detached).detach()
+ if tuple(preconv.shape[:3]) != (
+ batch,
+ int(model.z_dim),
+ latent_frames,
+ ):
+ raise RuntimeError(f"unexpected conv2 output shape {tuple(preconv.shape)}")
+
+ boundaries: dict[int, tuple[Any, ...]] = {0: initial_state}
+ state = initial_state
+ pixel_lengths: list[int] = []
+ cursor = 0
+ final_state: tuple[Any, ...] | None = None
+ with torch.no_grad():
+ for frame_index in range(latent_frames):
+ if frame_index and frame_index % CHUNK_LATENT_FRAMES == 0:
+ boundary = _state_to_cpu(state, formal=formal)
+ _validate_state(
+ boundary,
+ slots=slots,
+ frames_consumed=frame_index,
+ batch=batch,
+ height=height,
+ width=width,
+ formal=formal,
+ device_type="cpu",
+ )
+ boundaries[frame_index] = boundary
+ pixels, state = _transition(
+ model=model,
+ preconv_frame=preconv[:, :, frame_index : frame_index + 1],
+ state=state,
+ slots=slots,
+ )
+ _validate_state(
+ state,
+ slots=slots,
+ frames_consumed=frame_index + 1,
+ batch=batch,
+ height=height,
+ width=width,
+ formal=formal,
+ device_type=z_detached.device.type,
+ )
+ length = int(pixels.shape[1])
+ if tuple(pixels.shape[2:]) != tuple(g_video01.shape[2:]):
+ raise RuntimeError("Wan transition pixel shape differs from cotangent")
+ if formal:
+ start, end = _formal_pixel_span(frame_index)
+ if cursor != start or length != end - start:
+ raise RuntimeError(
+ f"frame {frame_index} emitted {length} at {cursor}, "
+ f"expected [{start},{end})"
+ )
+ if expected_pixels01 is not None:
+ expected = expected_pixels01[:, cursor : cursor + length]
+ if not torch.equal(pixels, expected):
+ difference = (pixels.float() - expected.float()).abs()
+ raise RuntimeError(
+ "functional Wan decode differs from official values; "
+ f"frame={frame_index}, unequal="
+ f"{int(torch.count_nonzero(pixels != expected).item())}, "
+ f"max_abs={float(difference.max().item())}"
+ )
+ pixel_lengths.append(length)
+ cursor += length
+ final_state = state
+ del pixels
+
+ if cursor != g_video01.shape[1]:
+ raise ValueError(
+ f"Wan emitted {cursor} pixel frames, cotangent has {g_video01.shape[1]}"
+ )
+ if formal:
+ if pixel_lengths != [1] + [4] * 20 or cursor != 81:
+ raise RuntimeError(
+ f"formal output schedule changed: {pixel_lengths}, total={cursor}"
+ )
+ if sorted(boundaries) != [0, 4, 8, 12, 16, 20]:
+ raise RuntimeError(f"formal CPU boundaries changed: {sorted(boundaries)}")
+ if final_state is None:
+ raise RuntimeError("Wan decoder emitted no transitions")
+ if formal and any(not torch.is_tensor(value) for value in final_state):
+ raise RuntimeError("all 32 final formal cache states must be tensors")
+
+ # There is no terminal-state loss. Each reverse transition creates the
+ # exact future-state cotangent consumed by the preceding transition.
+ future_cotangent: tuple[Any, ...] = (None,) * slots
+ del final_state, state
+ grad_preconv = torch.empty_like(preconv)
+ offsets = [0]
+ for length in pixel_lengths:
+ offsets.append(offsets[-1] + length)
+ chunk_starts = list(range(0, latent_frames, CHUNK_LATENT_FRAMES))
+ if formal and list(reversed(chunk_starts)) != [20, 16, 12, 8, 4, 0]:
+ raise RuntimeError(f"formal reverse chunks changed: {chunk_starts}")
+
+ for chunk_start in reversed(chunk_starts):
+ chunk_end = min(chunk_start + CHUNK_LATENT_FRAMES, latent_frames)
+ if chunk_start not in boundaries:
+ raise RuntimeError(f"missing CPU boundary at frame {chunk_start}")
+ replay_state = _state_to_device(
+ boundaries[chunk_start], device=z_detached.device
+ )
+ states_before: list[tuple[Any, ...]] = []
+ with torch.no_grad():
+ for frame_index in range(chunk_start, chunk_end):
+ states_before.append(replay_state)
+ replay_pixels, replay_state = _transition(
+ model=model,
+ preconv_frame=preconv[:, :, frame_index : frame_index + 1],
+ state=replay_state,
+ slots=slots,
+ )
+ if replay_pixels.shape[1] != pixel_lengths[frame_index]:
+ raise RuntimeError("chunk replay changed its pixel length")
+ del replay_pixels
+ del replay_state
+
+ for frame_index in reversed(range(chunk_start, chunk_end)):
+ input_state = _state_as_leaves(states_before.pop())
+ preconv_leaf = (
+ preconv[:, :, frame_index : frame_index + 1]
+ .detach()
+ .requires_grad_(True)
+ )
+ pixels, output_state = _transition(
+ model=model,
+ preconv_frame=preconv_leaf,
+ state=input_state,
+ slots=slots,
+ )
+ start, end = offsets[frame_index], offsets[frame_index + 1]
+ pixel_cotangent = g_video01[:, start:end].detach().float()
+ if tuple(pixel_cotangent.shape) != tuple(pixels.shape):
+ raise RuntimeError("reverse pixel cotangent shape mismatch")
+ if expected_pixels01 is not None:
+ expected = expected_pixels01[:, start:end]
+ if not torch.equal(pixels.detach(), expected):
+ difference = (pixels.detach().float() - expected.float()).abs()
+ raise RuntimeError(
+ "gradient Wan replay differs from forward values; "
+ f"frame={frame_index}, unequal="
+ f"{int(torch.count_nonzero(pixels.detach() != expected).item())}, "
+ f"max_abs={float(difference.max().item())}"
+ )
+ outputs = [pixels]
+ output_cotangents = [pixel_cotangent]
+ for index, (value, cotangent) in enumerate(
+ zip(output_state, future_cotangent)
+ ):
+ if cotangent is None:
+ continue
+ if not torch.is_tensor(value):
+ raise RuntimeError(
+ f"output state {index} has a cotangent but is not a tensor"
+ )
+ if tuple(value.shape) != tuple(cotangent.shape):
+ raise RuntimeError(f"state cotangent mismatch in slot {index}")
+ if value.dtype != cotangent.dtype:
+ raise RuntimeError(f"state cotangent dtype mismatch in slot {index}")
+ outputs.append(value)
+ output_cotangents.append(cotangent.detach())
+
+ positions = [
+ index for index, value in enumerate(input_state) if torch.is_tensor(value)
+ ]
+ state_leaves = [input_state[index] for index in positions]
+ gradients = torch.autograd.grad(
+ outputs,
+ [preconv_leaf, *state_leaves],
+ grad_outputs=output_cotangents,
+ retain_graph=False,
+ create_graph=False,
+ allow_unused=True,
+ )
+ grad_x = gradients[0]
+ if grad_x is None or not torch.isfinite(grad_x).all():
+ raise RuntimeError(f"transition {frame_index} has invalid conv2 gradient")
+ grad_preconv[:, :, frame_index : frame_index + 1].copy_(grad_x.detach())
+ incoming: list[Any] = [None] * slots
+ for index, gradient in zip(positions, gradients[1:]):
+ if gradient is None:
+ raise RuntimeError(
+ f"cache slot {index} lost its recurrent gradient"
+ )
+ if not torch.isfinite(gradient).all():
+ raise RuntimeError(f"cache slot {index} has non-finite cotangent")
+ incoming[index] = gradient.detach()
+ future_cotangent = tuple(incoming)
+ value = cotangent = gradient = None
+ del (
+ gradients,
+ grad_x,
+ outputs,
+ output_cotangents,
+ pixels,
+ output_state,
+ preconv_leaf,
+ input_state,
+ positions,
+ state_leaves,
+ pixel_cotangent,
+ incoming,
+ )
+ if states_before:
+ raise AssertionError("reverse chunk left unreplayed states")
+
+ # Finish the exact chain rule through one full-shape conv2+scale replay.
+ latent_leaf = z_detached.detach().clone().requires_grad_(True)
+ with _amp(z_detached.device.type):
+ replay_preconv = _full_preconv(vae, latent_leaf)
+ if not torch.equal(replay_preconv.detach(), preconv):
+ difference = (replay_preconv.detach().float() - preconv.float()).abs()
+ raise RuntimeError(
+ f"full-shape conv2 replay changed values; max_abs={float(difference.max())}"
+ )
+ (gradient,) = torch.autograd.grad(
+ replay_preconv,
+ latent_leaf,
+ grad_outputs=grad_preconv.detach(),
+ retain_graph=False,
+ create_graph=False,
+ )
+ if not torch.isfinite(gradient).all():
+ raise RuntimeError("exact recurrent Wan VJP produced non-finite gradients")
+ if _model_cache_signature(model) != global_signature:
+ raise RuntimeError("functional Wan VJP mutated model._feat_map/_conv_idx")
+ return gradient.detach()
diff --git a/14B/code/video_fd/samplers/__init__.py b/14B/code/video_fd/samplers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3cf2151b369867c739c7f7023634b5cc8ee09e1a
--- /dev/null
+++ b/14B/code/video_fd/samplers/__init__.py
@@ -0,0 +1,7 @@
+from video_fd.samplers.stride4_clip_sampler import (
+ ClipSpec,
+ Stride4ClipSampler,
+ stride4_candidate_indices,
+)
+
+__all__ = ["ClipSpec", "Stride4ClipSampler", "stride4_candidate_indices"]
diff --git a/14B/code/video_fd/samplers/stride4_clip_sampler.py b/14B/code/video_fd/samplers/stride4_clip_sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..27460b605954fb64eb8726281bf8cdf9718bc054
--- /dev/null
+++ b/14B/code/video_fd/samplers/stride4_clip_sampler.py
@@ -0,0 +1,200 @@
+"""Stride-4 latent-aligned contiguous-clip sampler.
+
+From a ``num_pixel_frames``-frame, ``source_fps``-fps video, select a phase
+``o ~ Uniform{0,1,2,3}`` and take every ``stride``-th frame
+(``I = {o + stride*k}``), giving an effective ``source_fps/stride`` fps stream;
+then take a random contiguous ``clip_frames``-frame window of that stream.
+
+For the canonical setting (81 frames @ 16 fps, stride 4, clip 16) this yields 16
+frames at an effective 4 fps. The *same* sampler is used for target-statistic
+precomputation, generated training, and evaluation (random phase + random valid
+start), with an optional deterministic mode for tests/debugging.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import torch
+
+
+def stride4_candidate_indices(phase: int, stride: int, num_pixel_frames: int) -> List[int]:
+ """All frame indices ``phase + stride*k`` within ``[0, num_pixel_frames-1]``."""
+ last = num_pixel_frames - 1
+ return [phase + stride * k for k in range(num_pixel_frames) if phase + stride * k <= last]
+
+
+@dataclass
+class ClipSpec:
+ """One sampled clip: a set of frame indices applied across the batch.
+
+ ``crop_id`` / ``encoder_name`` are optional bookkeeping fields set by the
+ caller (the spatial crop is applied by the encoder adapter, not here).
+ """
+
+ phase: int
+ start: int
+ frame_indices: List[int]
+ stride: int = 4
+ clip_frames: int = 16
+ crop_id: Optional[str] = None
+ encoder_name: Optional[str] = None
+
+ def to_dict(self) -> dict:
+ return {
+ "phase": int(self.phase),
+ "start": int(self.start),
+ "frame_indices": list(self.frame_indices),
+ "stride": int(self.stride),
+ "clip_frames": int(self.clip_frames),
+ "crop_id": self.crop_id,
+ "encoder_name": self.encoder_name,
+ }
+
+ def index_tensor(self, device=None) -> torch.Tensor:
+ return torch.tensor(self.frame_indices, dtype=torch.long, device=device)
+
+
+class Stride4ClipSampler:
+ """Sampler for stride-4 contiguous clips.
+
+ Parameters
+ ----------
+ num_pixel_frames, source_fps, stride, clip_frames:
+ temporal geometry (defaults: 81, 16, 4, 16 -> 16 frames @ 4 fps).
+ random_phase / random_start:
+ if ``False`` and a ``deterministic`` value is given, use it.
+ deterministic / phase / start:
+ debug mode -- always return ``frame_indices`` for the fixed phase/start.
+ """
+
+ def __init__(
+ self,
+ num_pixel_frames: int = 81,
+ source_fps: int = 16,
+ stride: int = 4,
+ clip_frames: int = 16,
+ random_phase: bool = True,
+ random_start: bool = True,
+ deterministic: bool = False,
+ phase: int = 0,
+ start: int = 0,
+ ):
+ if clip_frames <= 0:
+ raise ValueError("clip_frames must be positive")
+ self.num_pixel_frames = int(num_pixel_frames)
+ self.source_fps = int(source_fps)
+ self.stride = int(stride)
+ self.clip_frames = int(clip_frames)
+ self.random_phase = bool(random_phase)
+ self.random_start = bool(random_start)
+ self.deterministic = bool(deterministic)
+ self.fixed_phase = int(phase)
+ self.fixed_start = int(start)
+
+ # validate feasibility: the shortest phase stream must still hold a clip
+ min_candidates = min(
+ len(stride4_candidate_indices(p, self.stride, self.num_pixel_frames))
+ for p in range(self.stride)
+ )
+ if min_candidates < self.clip_frames:
+ raise ValueError(
+ f"stride={self.stride}, clip_frames={self.clip_frames} infeasible: "
+ f"some phase yields only {min_candidates} candidate frames "
+ f"(< {self.clip_frames}) for {self.num_pixel_frames} pixel frames"
+ )
+
+ @property
+ def effective_fps(self) -> float:
+ return self.source_fps / self.stride
+
+ def _sample_one(self, generator: Optional[torch.Generator]) -> ClipSpec:
+ if self.deterministic:
+ phase = self.fixed_phase
+ start = self.fixed_start
+ else:
+ if self.random_phase:
+ phase = int(torch.randint(0, self.stride, (), generator=generator).item())
+ else:
+ phase = self.fixed_phase
+ candidate = stride4_candidate_indices(phase, self.stride, self.num_pixel_frames)
+ max_start = len(candidate) - self.clip_frames
+ if self.random_start and max_start > 0:
+ start = int(torch.randint(0, max_start + 1, (), generator=generator).item())
+ else:
+ start = min(self.fixed_start, max_start) if not self.random_start else 0
+ candidate = stride4_candidate_indices(phase, self.stride, self.num_pixel_frames)
+ max_start = len(candidate) - self.clip_frames
+ start = max(0, min(start, max_start))
+ frame_indices = candidate[start: start + self.clip_frames]
+ return ClipSpec(
+ phase=phase,
+ start=start,
+ frame_indices=frame_indices,
+ stride=self.stride,
+ clip_frames=self.clip_frames,
+ )
+
+ def sample(
+ self,
+ batch_size: int = 1,
+ num_clips: int = 1,
+ generator: Optional[torch.Generator] = None,
+ ) -> List[ClipSpec]:
+ """Return ``num_clips`` clip specs (each applied across the whole batch).
+
+ ``batch_size`` is accepted for API symmetry; the same frame indices are
+ applied to all batch items of a given clip (see ``gather_frames``).
+ """
+ return [self._sample_one(generator) for _ in range(max(1, num_clips))]
+
+ def config_dict(self) -> dict:
+ """Sampler config to embed in every statistics file (spec section 1.2)."""
+ return {
+ "source_fps": self.source_fps,
+ "num_pixel_frames": self.num_pixel_frames,
+ "stride": self.stride,
+ "effective_fps": self.effective_fps,
+ "clip_frames": self.clip_frames,
+ "phase_distribution": "uniform_0_1_2_3" if self.stride == 4 else f"uniform_0_{self.stride - 1}",
+ "start_distribution": "uniform_valid_contiguous_start",
+ "deterministic": self.deterministic,
+ }
+
+ @classmethod
+ def from_config(cls, cfg: dict) -> "Stride4ClipSampler":
+ return cls(
+ num_pixel_frames=int(cfg.get("num_pixel_frames", 81)),
+ source_fps=int(cfg.get("source_fps", 16)),
+ stride=int(cfg.get("stride", 4)),
+ clip_frames=int(cfg.get("clip_frames", 16)),
+ random_phase=bool(cfg.get("random_phase", True)),
+ random_start=bool(cfg.get("random_start", True)),
+ deterministic=bool(cfg.get("deterministic", cfg.get("deterministic_debug", False))),
+ phase=int(cfg.get("phase", 0)),
+ start=int(cfg.get("start", 0)),
+ )
+
+
+def gather_frames(video: torch.Tensor, frame_indices, time_dim: int = 1) -> torch.Tensor:
+ """Select ``frame_indices`` along ``time_dim`` of ``video`` (differentiable)."""
+ if not torch.is_tensor(frame_indices):
+ frame_indices = torch.tensor(list(frame_indices), dtype=torch.long, device=video.device)
+ else:
+ frame_indices = frame_indices.to(device=video.device, dtype=torch.long)
+ return video.index_select(time_dim, frame_indices)
+
+
+def scatter_add_frames(target: torch.Tensor, frame_indices, values: torch.Tensor, time_dim: int = 1) -> torch.Tensor:
+ """Accumulate ``values`` into ``target`` at ``frame_indices`` along ``time_dim``.
+
+ Uses ``index_add`` so that a pixel frame selected by multiple clips / crops /
+ encoders correctly accumulates gradients (spec section 11.3).
+ """
+ if not torch.is_tensor(frame_indices):
+ frame_indices = torch.tensor(list(frame_indices), dtype=torch.long, device=target.device)
+ else:
+ frame_indices = frame_indices.to(device=target.device, dtype=torch.long)
+ target.index_add_(time_dim, frame_indices, values.to(target.dtype))
+ return target
diff --git a/14B/code/wan/__init__.py b/14B/code/wan/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..764073680ebcb829e03f1f9a207a6248e0da84a6
--- /dev/null
+++ b/14B/code/wan/__init__.py
@@ -0,0 +1 @@
+"""Wan backbone modules used by distribution matching."""
diff --git a/14B/code/wan/modules/__init__.py b/14B/code/wan/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f8935bbb45ab4e3f349d203b673102f7cfc07553
--- /dev/null
+++ b/14B/code/wan/modules/__init__.py
@@ -0,0 +1,16 @@
+from .attention import flash_attention
+from .model import WanModel
+from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
+from .tokenizers import HuggingfaceTokenizer
+from .vae import WanVAE
+
+__all__ = [
+ 'WanVAE',
+ 'WanModel',
+ 'T5Model',
+ 'T5Encoder',
+ 'T5Decoder',
+ 'T5EncoderModel',
+ 'HuggingfaceTokenizer',
+ 'flash_attention',
+]
diff --git a/14B/code/wan/modules/attention.py b/14B/code/wan/modules/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a64f256328cb980dcb749302d685d5c4b657ab1
--- /dev/null
+++ b/14B/code/wan/modules/attention.py
@@ -0,0 +1,218 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+
+try:
+ import flash_attn_interface
+
+ def is_hopper_gpu():
+ if not torch.cuda.is_available():
+ return False
+ device_name = torch.cuda.get_device_name(0).lower()
+ return "h100" in device_name or "hopper" in device_name
+ FLASH_ATTN_3_AVAILABLE = is_hopper_gpu()
+except ModuleNotFoundError:
+ FLASH_ATTN_3_AVAILABLE = False
+
+try:
+ import flash_attn
+ FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_2_AVAILABLE = False
+
+# FLASH_ATTN_3_AVAILABLE = False
+
+import warnings
+
+__all__ = [
+ 'flash_attention',
+ 'attention',
+]
+
+
+def flash_attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ version=None,
+):
+ """
+ q: [B, Lq, Nq, C1].
+ k: [B, Lk, Nk, C1].
+ v: [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+ q_lens: [B].
+ k_lens: [B].
+ dropout_p: float. Dropout probability.
+ softmax_scale: float. The scaling of QK^T before applying softmax.
+ causal: bool. Whether to apply causal attention mask.
+ window_size: (left right). If not (-1, -1), apply sliding window local attention.
+ deterministic: bool. If True, slightly slower and uses more memory.
+ dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+ """
+ half_dtypes = (torch.float16, torch.bfloat16)
+ assert dtype in half_dtypes
+ assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+ # params
+ b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+ def half(x):
+ return x if x.dtype in half_dtypes else x.to(dtype)
+
+ # preprocess query
+ if q_lens is None:
+ q = half(q.flatten(0, 1))
+ q_lens = torch.tensor(
+ [lq] * b, dtype=torch.int32).to(
+ device=q.device, non_blocking=True)
+ else:
+ q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+ # preprocess key, value
+ if k_lens is None:
+ k = half(k.flatten(0, 1))
+ v = half(v.flatten(0, 1))
+ k_lens = torch.tensor(
+ [lk] * b, dtype=torch.int32).to(
+ device=k.device, non_blocking=True)
+ else:
+ k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+ v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+ q = q.to(v.dtype)
+ k = k.to(v.dtype)
+
+ if q_scale is not None:
+ q = q * q_scale
+
+ if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+ warnings.warn(
+ 'Flash attention 3 is not available, use flash attention 2 instead.'
+ )
+
+ # apply attention
+ if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+ # Note: dropout_p, window_size are not supported in FA3 now.
+ x = flash_attn_interface.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ deterministic=deterministic)[0].unflatten(0, (b, lq))
+ else:
+ assert FLASH_ATTN_2_AVAILABLE
+ x = flash_attn.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic).unflatten(0, (b, lq))
+
+ # output
+ return x.type(out_dtype)
+
+
+def attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ fa_version=None,
+):
+ if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+ return flash_attention(
+ q=q,
+ k=k,
+ v=v,
+ q_lens=q_lens,
+ k_lens=k_lens,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ q_scale=q_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic,
+ dtype=dtype,
+ version=fa_version,
+ )
+ else:
+ if window_size != (-1, -1):
+ raise RuntimeError(
+ "PyTorch SDPA fallback does not support Wan sliding windows"
+ )
+ out_dtype = q.dtype
+ batch_size, max_q_len = q.shape[:2]
+ outputs = []
+ for batch_index in range(batch_size):
+ query_len = (
+ max_q_len
+ if q_lens is None
+ else int(q_lens[batch_index].item())
+ )
+ key_len = (
+ k.shape[1]
+ if k_lens is None
+ else int(k_lens[batch_index].item())
+ )
+ query = q[batch_index : batch_index + 1, :query_len]
+ key = k[batch_index : batch_index + 1, :key_len]
+ value = v[batch_index : batch_index + 1, :key_len]
+ query = query.transpose(1, 2).to(dtype)
+ key = key.transpose(1, 2).to(dtype)
+ value = value.transpose(1, 2).to(dtype)
+ if q_scale is not None:
+ query = query * q_scale
+ if query.shape[1] != key.shape[1]:
+ if query.shape[1] % key.shape[1]:
+ raise ValueError(
+ "query heads must be divisible by key/value heads"
+ )
+ repeats = query.shape[1] // key.shape[1]
+ key = key.repeat_interleave(repeats, dim=1)
+ value = value.repeat_interleave(repeats, dim=1)
+ output = torch.nn.functional.scaled_dot_product_attention(
+ query,
+ key,
+ value,
+ dropout_p=dropout_p,
+ is_causal=causal,
+ scale=softmax_scale,
+ ).transpose(1, 2)
+ if query_len < max_q_len:
+ output = torch.nn.functional.pad(
+ output, (0, 0, 0, 0, 0, max_q_len - query_len)
+ )
+ outputs.append(output)
+ return torch.cat(outputs, dim=0).contiguous().to(out_dtype)
diff --git a/14B/code/wan/modules/causal_model.py b/14B/code/wan/modules/causal_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..13bf00b96712386502c9a518d79dfe5e9f45d5e1
--- /dev/null
+++ b/14B/code/wan/modules/causal_model.py
@@ -0,0 +1,1166 @@
+from wan.modules.attention import attention
+from wan.modules.model import (
+ WanRMSNorm,
+ rope_apply,
+ WanLayerNorm,
+ WAN_CROSSATTENTION_CLASSES,
+ rope_params,
+ MLPProj,
+ sinusoidal_embedding_1d
+)
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from torch.nn.attention.flex_attention import BlockMask
+from diffusers.models.modeling_utils import ModelMixin
+import torch.nn as nn
+import torch
+import math
+import torch.distributed as dist
+
+# wan 1.3B model has a weird channel / head configurations and require max-autotune to work with flexattention
+# see https://github.com/pytorch/pytorch/issues/133254
+# change to default for other models
+flex_attention = torch.compile(
+ flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs")
+
+
+def causal_rope_apply(x, grid_sizes, freqs, start_frame=0):
+ n, c = x.size(2), x.size(3) // 2
+
+ # split freqs
+ freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+ # loop over samples
+ output = []
+
+ for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+ seq_len = f * h * w
+
+ # Preserve the pretrained causal Wan numerics. The encoder-specific
+ # memory trade-off must not silently change the student's RoPE.
+ x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
+ seq_len, n, -1, 2))
+ freqs_i = torch.cat([
+ freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1),
+ freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+ freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+ ],
+ dim=-1).reshape(seq_len, 1, -1)
+
+ # apply rotary embedding
+ x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+ x_i = torch.cat([x_i, x[i, seq_len:]])
+
+ # append to collection
+ output.append(x_i)
+ return torch.stack(output).type_as(x)
+
+
+class CausalWanSelfAttention(nn.Module):
+
+ def __init__(self,
+ dim,
+ num_heads,
+ local_attn_size=-1,
+ sink_size=0,
+ qk_norm=True,
+ eps=1e-6):
+ assert dim % num_heads == 0
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.local_attn_size = local_attn_size
+ self.sink_size = sink_size
+ self.qk_norm = qk_norm
+ self.eps = eps
+ self.max_attention_size = 32760 if local_attn_size == -1 else local_attn_size * 1560
+
+ # layers
+ self.q = nn.Linear(dim, dim)
+ self.k = nn.Linear(dim, dim)
+ self.v = nn.Linear(dim, dim)
+ self.o = nn.Linear(dim, dim)
+ self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(
+ self,
+ x,
+ seq_lens,
+ grid_sizes,
+ freqs,
+ block_mask,
+ kv_cache=None,
+ current_start=0,
+ cache_start=None,
+ update_kv_cache=True,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, num_heads, C / num_heads]
+ seq_lens(Tensor): Shape [B]
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ block_mask (BlockMask)
+ """
+ b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+ if cache_start is None:
+ cache_start = current_start
+
+ # query, key, value function
+ def qkv_fn(x):
+ q = self.norm_q(self.q(x)).view(b, s, n, d)
+ k = self.norm_k(self.k(x)).view(b, s, n, d)
+ v = self.v(x).view(b, s, n, d)
+ return q, k, v
+
+ q, k, v = qkv_fn(x)
+
+ if kv_cache is None:
+ # if it is teacher forcing training?
+ is_tf = (s == seq_lens[0].item() * 2)
+ if is_tf:
+ q_chunk = torch.chunk(q, 2, dim=1)
+ k_chunk = torch.chunk(k, 2, dim=1)
+ roped_query = []
+ roped_key = []
+ # rope should be same for clean and noisy parts
+ for ii in range(2):
+ rq = rope_apply(q_chunk[ii], grid_sizes, freqs).type_as(v)
+ rk = rope_apply(k_chunk[ii], grid_sizes, freqs).type_as(v)
+ roped_query.append(rq)
+ roped_key.append(rk)
+
+ roped_query = torch.cat(roped_query, dim=1)
+ roped_key = torch.cat(roped_key, dim=1)
+
+ padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
+ padded_roped_query = torch.cat(
+ [roped_query,
+ torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
+ device=q.device, dtype=v.dtype)],
+ dim=1
+ )
+
+ padded_roped_key = torch.cat(
+ [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],
+ device=k.device, dtype=v.dtype)],
+ dim=1
+ )
+
+ padded_v = torch.cat(
+ [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],
+ device=v.device, dtype=v.dtype)],
+ dim=1
+ )
+
+ x = flex_attention(
+ query=padded_roped_query.transpose(2, 1),
+ key=padded_roped_key.transpose(2, 1),
+ value=padded_v.transpose(2, 1),
+ block_mask=block_mask
+ )[:, :, :-padded_length].transpose(2, 1)
+
+ else:
+ roped_query = rope_apply(q, grid_sizes, freqs).type_as(v)
+ roped_key = rope_apply(k, grid_sizes, freqs).type_as(v)
+
+ padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
+ padded_roped_query = torch.cat(
+ [roped_query,
+ torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
+ device=q.device, dtype=v.dtype)],
+ dim=1
+ )
+
+ padded_roped_key = torch.cat(
+ [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],
+ device=k.device, dtype=v.dtype)],
+ dim=1
+ )
+
+ padded_v = torch.cat(
+ [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],
+ device=v.device, dtype=v.dtype)],
+ dim=1
+ )
+
+ x = flex_attention(
+ query=padded_roped_query.transpose(2, 1),
+ key=padded_roped_key.transpose(2, 1),
+ value=padded_v.transpose(2, 1),
+ block_mask=block_mask
+ )[:, :, :-padded_length].transpose(2, 1)
+ else:
+ frame_seqlen = math.prod(grid_sizes[0][1:]).item()
+ current_start_frame = current_start // frame_seqlen
+ roped_query = causal_rope_apply(
+ q, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)
+ roped_key = causal_rope_apply(
+ k, grid_sizes, freqs, start_frame=current_start_frame).type_as(v)
+
+ current_end = current_start + roped_query.shape[1]
+ sink_tokens = self.sink_size * frame_seqlen
+ if not isinstance(update_kv_cache, bool):
+ raise TypeError("update_kv_cache must be a bool")
+ if not update_kv_cache:
+ # Gradient-enabled Self-Forcing exits must be pure with respect
+ # to their persistent cache. Activation checkpointing
+ # recomputes this forward during backward, after the no-grad
+ # clean pass has committed the same slot. Reading only the
+ # detached clean prefix and concatenating the differentiable
+ # current K/V makes both executions identical and prevents a
+ # checkpoint recomputation from advancing or overwriting cache
+ # state.
+ if self.local_attn_size != -1:
+ raise ValueError(
+ "read-only KV-cache mode requires local_attn_size=-1"
+ )
+ if int(cache_start) != int(current_start):
+ raise ValueError(
+ "read-only KV-cache mode requires cache_start == current_start"
+ )
+ if int(current_start) < 0 or int(current_end) > int(kv_cache["k"].shape[1]):
+ raise ValueError("read-only KV-cache range exceeds cache capacity")
+ global_end = int(kv_cache["global_end_index"].item())
+ local_end = int(kv_cache["local_end_index"].item())
+ if global_end < int(current_start) or local_end < int(current_start):
+ raise ValueError(
+ "read-only KV-cache prefix has not been committed through current_start"
+ )
+ prefix_k = kv_cache["k"][:, :current_start].detach()
+ prefix_v = kv_cache["v"][:, :current_start].detach()
+ attention_k = torch.cat((prefix_k, roped_key), dim=1)
+ attention_v = torch.cat((prefix_v, v), dim=1)
+ x = attention(roped_query, attention_k, attention_v)
+ x = x.flatten(2)
+ return self.o(x)
+ # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache
+ kv_cache_size = kv_cache["k"].shape[1]
+ num_new_tokens = roped_query.shape[1]
+ if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
+ num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
+ # Calculate the number of new tokens added in this step
+ # Shift existing cache content left to discard oldest tokens
+ # Clone the source slice to avoid overlapping memory error
+ num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
+ num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
+ kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+ kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+ kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+ kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+ # Insert the new keys/values at the end
+ local_end_index = kv_cache["local_end_index"].item() + current_end - \
+ kv_cache["global_end_index"].item() - num_evicted_tokens
+ local_start_index = local_end_index - num_new_tokens
+ kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+ kv_cache["v"][:, local_start_index:local_end_index] = v
+ else:
+ # Assign new keys/values directly up to current_end
+ local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
+ local_start_index = local_end_index - num_new_tokens
+ kv_cache["k"][:, local_start_index:local_end_index] = roped_key
+ kv_cache["v"][:, local_start_index:local_end_index] = v
+ # Clone the K/V slices we hand to FlashAttention. FA's autograd
+ # Function saves these tensors verbatim (no internal copy when
+ # the input stride(-1) == 1, which is our case), so without the
+ # clone the saved K/V are VIEWS into the persistent KV-cache
+ # storage. Self-Forcing performs multiple forwards per block --
+ # 4 denoising passes (one of them gradient-enabled) plus a
+ # ``no_grad`` clean re-run -- each of which writes back to the
+ # same cache slot via ``kv_cache["k"][a:b] = roped_key``. Every
+ # write bumps the storage version counter. When backward
+ # eventually runs on the gradient-enabled forward's saved K, the
+ # storage version has advanced past the saved version and
+ # autograd raises ``RuntimeError: one of the variables needed for
+ # gradient computation has been modified by an inplace operation:
+ # [...] is at version N; expected version N-1 instead``.
+ #
+ # ``.clone()`` produces an independent storage whose version is
+ # immune to subsequent cache mutations. Cost: ~24 MB / block /
+ # call at the Wan-1.3B-ish shape (B=1, slice=7800, H=12, D=128,
+ # bf16); only the gradient-enabled forward actually pays for it
+ # since ``.clone()`` is free in ``torch.no_grad()`` (PyTorch
+ # still allocates but doesn't track for backward).
+ #
+ # NOTE: applies to BOTH the all-softmax baseline (which also
+ # contains this latent bug; just rarely triggered because the
+ # baseline uses full-visibility caches and tends not to mix with
+ # alpha-mode gradient flow) and the mixed-attention model. The
+ # clone keeps the autograd graph intact: ``roped_key → IndexPut →
+ # cache_slice_view → clone → attention``. Differentiating
+ # backward through ``clone`` is identity; ``IndexPut.backward``
+ # only needs the upstream gradient, not the cache state.
+ x = attention(
+ roped_query,
+ kv_cache["k"][:, max(0, local_end_index - self.max_attention_size):local_end_index].clone(),
+ kv_cache["v"][:, max(0, local_end_index - self.max_attention_size):local_end_index].clone(),
+ )
+ kv_cache["global_end_index"].fill_(current_end)
+ kv_cache["local_end_index"].fill_(local_end_index)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class CausalWanAttentionBlock(nn.Module):
+
+ def __init__(self,
+ cross_attn_type,
+ dim,
+ ffn_dim,
+ num_heads,
+ local_attn_size=-1,
+ sink_size=0,
+ qk_norm=True,
+ cross_attn_norm=False,
+ eps=1e-6):
+ super().__init__()
+ self.dim = dim
+ self.ffn_dim = ffn_dim
+ self.num_heads = num_heads
+ self.local_attn_size = local_attn_size
+ self.qk_norm = qk_norm
+ self.cross_attn_norm = cross_attn_norm
+ self.eps = eps
+
+ # layers
+ self.norm1 = WanLayerNorm(dim, eps)
+ self.self_attn = CausalWanSelfAttention(dim, num_heads, local_attn_size, sink_size, qk_norm, eps)
+ self.norm3 = WanLayerNorm(
+ dim, eps,
+ elementwise_affine=True) if cross_attn_norm else nn.Identity()
+ self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,
+ num_heads,
+ (-1, -1),
+ qk_norm,
+ eps)
+ self.norm2 = WanLayerNorm(dim, eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+ nn.Linear(ffn_dim, dim))
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+ def forward(
+ self,
+ x,
+ e,
+ seq_lens,
+ grid_sizes,
+ freqs,
+ context,
+ context_lens,
+ block_mask,
+ kv_cache=None,
+ crossattn_cache=None,
+ current_start=0,
+ cache_start=None,
+ update_kv_cache=True,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ e(Tensor): Shape [B, F, 6, C]
+ seq_lens(Tensor): Shape [B], length of each sequence in batch
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
+ # assert e.dtype == torch.float32
+ # with amp.autocast(dtype=torch.float32):
+ e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
+ # assert e[0].dtype == torch.float32
+
+ # self-attention
+ y = self.self_attn(
+ (self.norm1(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten(1, 2),
+ seq_lens,
+ grid_sizes,
+ freqs,
+ block_mask,
+ kv_cache=kv_cache,
+ current_start=current_start,
+ cache_start=cache_start,
+ update_kv_cache=update_kv_cache,
+ )
+
+ # with amp.autocast(dtype=torch.float32):
+ x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2)
+
+ # cross-attention & ffn function
+ def cross_attn_ffn(x, context, context_lens, e, crossattn_cache=None):
+ x = x + self.cross_attn(self.norm3(x), context,
+ context_lens, crossattn_cache=crossattn_cache)
+ y = self.ffn(
+ (self.norm2(x).unflatten(dim=1, sizes=(num_frames,
+ frame_seqlen)) * (1 + e[4]) + e[3]).flatten(1, 2)
+ )
+ # with amp.autocast(dtype=torch.float32):
+ x = x + (y.unflatten(dim=1, sizes=(num_frames,
+ frame_seqlen)) * e[5]).flatten(1, 2)
+ return x
+
+ x = cross_attn_ffn(x, context, context_lens, e, crossattn_cache)
+ return x
+
+
+class CausalHead(nn.Module):
+
+ def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+ super().__init__()
+ self.dim = dim
+ self.out_dim = out_dim
+ self.patch_size = patch_size
+ self.eps = eps
+
+ # layers
+ out_dim = math.prod(patch_size) * out_dim
+ self.norm = WanLayerNorm(dim, eps)
+ self.head = nn.Linear(dim, out_dim)
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+ def forward(self, x, e):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ e(Tensor): Shape [B, F, 1, C]
+ """
+ # assert e.dtype == torch.float32
+ # with amp.autocast(dtype=torch.float32):
+ num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
+ e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
+ x = (self.head(self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]))
+ return x
+
+
+class CausalWanModel(ModelMixin, ConfigMixin):
+ r"""
+ Wan diffusion backbone supporting both text-to-video and image-to-video.
+ """
+
+ ignore_for_config = [
+ 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim'
+ ]
+ _no_split_modules = ['WanAttentionBlock']
+ _supports_gradient_checkpointing = True
+
+ @register_to_config
+ def __init__(self,
+ model_type='t2v',
+ patch_size=(1, 2, 2),
+ text_len=512,
+ in_dim=16,
+ dim=2048,
+ ffn_dim=8192,
+ freq_dim=256,
+ text_dim=4096,
+ out_dim=16,
+ num_heads=16,
+ num_layers=32,
+ local_attn_size=-1,
+ sink_size=0,
+ qk_norm=True,
+ cross_attn_norm=True,
+ eps=1e-6):
+ r"""
+ Initialize the diffusion model backbone.
+
+ Args:
+ model_type (`str`, *optional*, defaults to 't2v'):
+ Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+ patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+ 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+ text_len (`int`, *optional*, defaults to 512):
+ Fixed length for text embeddings
+ in_dim (`int`, *optional*, defaults to 16):
+ Input video channels (C_in)
+ dim (`int`, *optional*, defaults to 2048):
+ Hidden dimension of the transformer
+ ffn_dim (`int`, *optional*, defaults to 8192):
+ Intermediate dimension in feed-forward network
+ freq_dim (`int`, *optional*, defaults to 256):
+ Dimension for sinusoidal time embeddings
+ text_dim (`int`, *optional*, defaults to 4096):
+ Input dimension for text embeddings
+ out_dim (`int`, *optional*, defaults to 16):
+ Output video channels (C_out)
+ num_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads
+ num_layers (`int`, *optional*, defaults to 32):
+ Number of transformer blocks
+ local_attn_size (`int`, *optional*, defaults to -1):
+ Window size for temporal local attention (-1 indicates global attention)
+ sink_size (`int`, *optional*, defaults to 0):
+ Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache
+ qk_norm (`bool`, *optional*, defaults to True):
+ Enable query/key normalization
+ cross_attn_norm (`bool`, *optional*, defaults to False):
+ Enable cross-attention normalization
+ eps (`float`, *optional*, defaults to 1e-6):
+ Epsilon value for normalization layers
+ """
+
+ super().__init__()
+
+ assert model_type in ['t2v', 'i2v']
+ self.model_type = model_type
+
+ self.patch_size = patch_size
+ self.text_len = text_len
+ self.in_dim = in_dim
+ self.dim = dim
+ self.ffn_dim = ffn_dim
+ self.freq_dim = freq_dim
+ self.text_dim = text_dim
+ self.out_dim = out_dim
+ self.num_heads = num_heads
+ self.num_layers = num_layers
+ self.local_attn_size = local_attn_size
+ self.qk_norm = qk_norm
+ self.cross_attn_norm = cross_attn_norm
+ self.eps = eps
+
+ # embeddings
+ self.patch_embedding = nn.Conv3d(
+ in_dim, dim, kernel_size=patch_size, stride=patch_size)
+ self.text_embedding = nn.Sequential(
+ nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
+ nn.Linear(dim, dim))
+
+ self.time_embedding = nn.Sequential(
+ nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+ self.time_projection = nn.Sequential(
+ nn.SiLU(), nn.Linear(dim, dim * 6))
+
+ # blocks
+ cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
+ self.blocks = nn.ModuleList([
+ CausalWanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+ local_attn_size, sink_size, qk_norm, cross_attn_norm, eps)
+ for _ in range(num_layers)
+ ])
+
+ # head
+ self.head = CausalHead(dim, out_dim, patch_size, eps)
+
+ # buffers (don't use register_buffer otherwise dtype will be changed in to())
+ assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+ d = dim // num_heads
+ self.freqs = torch.cat([
+ rope_params(1024, d - 4 * (d // 6)),
+ rope_params(1024, 2 * (d // 6)),
+ rope_params(1024, 2 * (d // 6))
+ ],
+ dim=1)
+
+ if model_type == 'i2v':
+ self.img_emb = MLPProj(1280, dim)
+
+ # initialize weights
+ self.init_weights()
+
+ self.gradient_checkpointing = False
+
+ self.block_mask = None
+
+ self.num_frame_per_block = 1
+ self.independent_first_frame = False
+
+ def _set_gradient_checkpointing(
+ self,
+ module=None,
+ value=False,
+ enable=None,
+ gradient_checkpointing_func=None,
+ ):
+ self.gradient_checkpointing = value if enable is None else enable
+
+ @staticmethod
+ def _prepare_blockwise_causal_attn_mask(
+ device: torch.device | str, num_frames: int = 21,
+ frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1
+ ) -> BlockMask:
+ """
+ we will divide the token sequence into the following format
+ [1 latent frame] [1 latent frame] ... [1 latent frame]
+ We use flexattention to construct the attention mask
+ """
+ total_length = num_frames * frame_seqlen
+
+ # we do right padding to get to a multiple of 128
+ padded_length = math.ceil(total_length / 128) * 128 - total_length
+
+ ends = torch.zeros(total_length + padded_length,
+ device=device, dtype=torch.long)
+
+ # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+ frame_indices = torch.arange(
+ start=0,
+ end=total_length,
+ step=frame_seqlen * num_frame_per_block,
+ device=device
+ )
+
+ for tmp in frame_indices:
+ ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \
+ frame_seqlen * num_frame_per_block
+
+ def attention_mask(b, h, q_idx, kv_idx):
+ if local_attn_size == -1:
+ return (kv_idx < ends[q_idx]) | (q_idx == kv_idx)
+ else:
+ return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx)
+ # return ((kv_idx < total_length) & (q_idx < total_length)) | (q_idx == kv_idx) # bidirectional mask
+
+ block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,
+ KV_LEN=total_length + padded_length, _compile=False, device=device)
+
+ import torch.distributed as dist
+ if not dist.is_initialized() or dist.get_rank() == 0:
+ print(
+ f" cache a block wise causal mask with block size of {num_frame_per_block} frames")
+ print(block_mask)
+
+ # import imageio
+ # import numpy as np
+ # from torch.nn.attention.flex_attention import create_mask
+
+ # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +
+ # padded_length, KV_LEN=total_length + padded_length, device=device)
+ # import cv2
+ # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))
+ # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))
+
+ return block_mask
+
+ @staticmethod
+ def _prepare_teacher_forcing_mask(
+ device: torch.device | str, num_frames: int = 21,
+ frame_seqlen: int = 1560, num_frame_per_block=1
+ ) -> BlockMask:
+ """
+ we will divide the token sequence into the following format
+ [1 latent frame] [1 latent frame] ... [1 latent frame]
+ We use flexattention to construct the attention mask
+ """
+ # debug
+ DEBUG = False
+ if DEBUG:
+ num_frames = 9
+ frame_seqlen = 256
+
+ total_length = num_frames * frame_seqlen * 2
+
+ # we do right padding to get to a multiple of 128
+ padded_length = math.ceil(total_length / 128) * 128 - total_length
+
+ clean_ends = num_frames * frame_seqlen
+ # for clean context frames, we can construct their flex attention mask based on a [start, end] interval
+ context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+ # for noisy frames, we need two intervals to construct the flex attention mask [context_start, context_end] [noisy_start, noisy_end]
+ noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+ noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+ noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+ noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
+
+ # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+ attention_block_size = frame_seqlen * num_frame_per_block
+ frame_indices = torch.arange(
+ start=0,
+ end=num_frames * frame_seqlen,
+ step=attention_block_size,
+ device=device, dtype=torch.long
+ )
+
+ # attention for clean context frames
+ for start in frame_indices:
+ context_ends[start:start + attention_block_size] = start + attention_block_size
+
+ noisy_image_start_list = torch.arange(
+ num_frames * frame_seqlen, total_length,
+ step=attention_block_size,
+ device=device, dtype=torch.long
+ )
+ noisy_image_end_list = noisy_image_start_list + attention_block_size
+
+ # attention for noisy frames
+ for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)):
+ # attend to noisy tokens within the same block
+ noise_noise_starts[start:end] = start
+ noise_noise_ends[start:end] = end
+ # attend to context tokens in previous blocks
+ # noise_context_starts[start:end] = 0
+ noise_context_ends[start:end] = block_index * attention_block_size
+
+ def attention_mask(b, h, q_idx, kv_idx):
+ # first design the mask for clean frames
+ clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx])
+ # then design the mask for noisy frames
+ # noisy frames will attend to all clean preceeding clean frames + itself
+ C1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx])
+ C2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx])
+ noise_mask = (q_idx >= clean_ends) & (C1 | C2)
+
+ eye_mask = q_idx == kv_idx
+ return eye_mask | clean_mask | noise_mask
+
+ block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,
+ KV_LEN=total_length + padded_length, _compile=False, device=device)
+
+ if DEBUG:
+ print(block_mask)
+ import imageio
+ import numpy as np
+ from torch.nn.attention.flex_attention import create_mask
+
+ mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +
+ padded_length, KV_LEN=total_length + padded_length, device=device)
+ import cv2
+ mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))
+ imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))
+
+ return block_mask
+
+ @staticmethod
+ def _prepare_blockwise_causal_attn_mask_i2v(
+ device: torch.device | str, num_frames: int = 21,
+ frame_seqlen: int = 1560, num_frame_per_block=4, local_attn_size=-1
+ ) -> BlockMask:
+ """
+ we will divide the token sequence into the following format
+ [1 latent frame] [N latent frame] ... [N latent frame]
+ The first frame is separated out to support I2V generation
+ We use flexattention to construct the attention mask
+ """
+ total_length = num_frames * frame_seqlen
+
+ # we do right padding to get to a multiple of 128
+ padded_length = math.ceil(total_length / 128) * 128 - total_length
+
+ ends = torch.zeros(total_length + padded_length,
+ device=device, dtype=torch.long)
+
+ # special handling for the first frame
+ ends[:frame_seqlen] = frame_seqlen
+
+ # Block-wise causal mask will attend to all elements that are before the end of the current chunk
+ frame_indices = torch.arange(
+ start=frame_seqlen,
+ end=total_length,
+ step=frame_seqlen * num_frame_per_block,
+ device=device
+ )
+
+ for idx, tmp in enumerate(frame_indices):
+ ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \
+ frame_seqlen * num_frame_per_block
+
+ def attention_mask(b, h, q_idx, kv_idx):
+ if local_attn_size == -1:
+ return (kv_idx < ends[q_idx]) | (q_idx == kv_idx)
+ else:
+ return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | \
+ (q_idx == kv_idx)
+
+ block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length,
+ KV_LEN=total_length + padded_length, _compile=False, device=device)
+
+ if not dist.is_initialized() or dist.get_rank() == 0:
+ print(
+ f" cache a block wise causal mask with block size of {num_frame_per_block} frames")
+ print(block_mask)
+
+ # import imageio
+ # import numpy as np
+ # from torch.nn.attention.flex_attention import create_mask
+
+ # mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length +
+ # padded_length, KV_LEN=total_length + padded_length, device=device)
+ # import cv2
+ # mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024))
+ # imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask))
+
+ return block_mask
+
+ def _forward_inference(
+ self,
+ x,
+ t,
+ context,
+ seq_len,
+ clip_fea=None,
+ y=None,
+ kv_cache: dict = None,
+ crossattn_cache: dict = None,
+ current_start: int = 0,
+ cache_start: int = 0,
+ update_kv_cache: bool = True,
+ ):
+ r"""
+ Run the diffusion model with kv caching.
+ See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details.
+ This function will be run for num_frame times.
+ Process the latent frames one by one (1560 tokens each)
+
+ Args:
+ x (List[Tensor]):
+ List of input video tensors, each with shape [C_in, F, H, W]
+ t (Tensor):
+ Diffusion timesteps tensor of shape [B]
+ context (List[Tensor]):
+ List of text embeddings each with shape [L, C]
+ seq_len (`int`):
+ Maximum sequence length for positional encoding
+ clip_fea (Tensor, *optional*):
+ CLIP image features for image-to-video mode
+ y (List[Tensor], *optional*):
+ Conditional video inputs for image-to-video mode, same shape as x
+
+ Returns:
+ List[Tensor]:
+ List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+ """
+
+ if self.model_type == 'i2v':
+ assert clip_fea is not None and y is not None
+ # params
+ device = self.patch_embedding.weight.device
+ if self.freqs.device != device:
+ self.freqs = self.freqs.to(device)
+
+ if y is not None:
+ x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+ # embeddings
+ x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+ grid_sizes = torch.stack(
+ [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+ x = [u.flatten(2).transpose(1, 2) for u in x]
+ seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+ assert seq_lens.max() <= seq_len
+ x = torch.cat(x)
+ """
+ torch.cat([
+ torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+ dim=1) for u in x
+ ])
+ """
+
+ # time embeddings
+ # with amp.autocast(dtype=torch.float32):
+ e = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
+ e0 = self.time_projection(e).unflatten(
+ 1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
+ # assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+ # context
+ context_lens = None
+ context = self.text_embedding(
+ torch.stack([
+ torch.cat(
+ [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+ for u in context
+ ]))
+
+ if clip_fea is not None:
+ context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+ context = torch.concat([context_clip, context], dim=1)
+
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens,
+ block_mask=self.block_mask,
+ update_kv_cache=update_kv_cache,
+ )
+
+ def create_custom_forward(module, block_index):
+ def custom_forward(*inputs, **kwargs):
+ full_cpu_offload = bool(getattr(
+ self,
+ "gradient_checkpointing_activation_cpu_offload",
+ False,
+ ))
+ cpu_offload_stride = int(getattr(
+ self,
+ "gradient_checkpointing_activation_cpu_offload_block_stride",
+ 0,
+ ))
+ selective_cpu_offload = (
+ cpu_offload_stride > 0
+ and block_index % cpu_offload_stride == 0
+ )
+ if full_cpu_offload or selective_cpu_offload:
+ with torch.autograd.graph.save_on_cpu(
+ pin_memory=False, device_type=inputs[0].device.type):
+ return module(*inputs, **kwargs)
+ return module(*inputs, **kwargs)
+ return custom_forward
+
+ for block_index, block in enumerate(self.blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ kwargs.update(
+ {
+ "kv_cache": kv_cache[block_index],
+ "current_start": current_start,
+ "cache_start": cache_start
+ }
+ )
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block, block_index),
+ x, **kwargs,
+ use_reentrant=False,
+ )
+ else:
+ kwargs.update(
+ {
+ "kv_cache": kv_cache[block_index],
+ "crossattn_cache": crossattn_cache[block_index],
+ "current_start": current_start,
+ "cache_start": cache_start
+ }
+ )
+ x = block(x, **kwargs)
+
+ # head
+ x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+ return torch.stack(x)
+
+ def _forward_train(
+ self,
+ x,
+ t,
+ context,
+ seq_len,
+ clean_x=None,
+ aug_t=None,
+ clip_fea=None,
+ y=None,
+ ):
+ r"""
+ Forward pass through the diffusion model
+
+ Args:
+ x (List[Tensor]):
+ List of input video tensors, each with shape [C_in, F, H, W]
+ t (Tensor):
+ Diffusion timesteps tensor of shape [B]
+ context (List[Tensor]):
+ List of text embeddings each with shape [L, C]
+ seq_len (`int`):
+ Maximum sequence length for positional encoding
+ clip_fea (Tensor, *optional*):
+ CLIP image features for image-to-video mode
+ y (List[Tensor], *optional*):
+ Conditional video inputs for image-to-video mode, same shape as x
+
+ Returns:
+ List[Tensor]:
+ List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+ """
+ if self.model_type == 'i2v':
+ assert clip_fea is not None and y is not None
+ # params
+ device = self.patch_embedding.weight.device
+ if self.freqs.device != device:
+ self.freqs = self.freqs.to(device)
+
+ # Construct blockwise causal attn mask
+ if self.block_mask is None:
+ if clean_x is not None:
+ if self.independent_first_frame:
+ raise NotImplementedError()
+ else:
+ self.block_mask = self._prepare_teacher_forcing_mask(
+ device, num_frames=x.shape[2],
+ frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),
+ num_frame_per_block=self.num_frame_per_block
+ )
+ else:
+ if self.independent_first_frame:
+ self.block_mask = self._prepare_blockwise_causal_attn_mask_i2v(
+ device, num_frames=x.shape[2],
+ frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),
+ num_frame_per_block=self.num_frame_per_block,
+ local_attn_size=self.local_attn_size
+ )
+ else:
+ self.block_mask = self._prepare_blockwise_causal_attn_mask(
+ device, num_frames=x.shape[2],
+ frame_seqlen=x.shape[-2] * x.shape[-1] // (self.patch_size[1] * self.patch_size[2]),
+ num_frame_per_block=self.num_frame_per_block,
+ local_attn_size=self.local_attn_size
+ )
+
+ if y is not None:
+ x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+ # embeddings
+ x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+
+ grid_sizes = torch.stack(
+ [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+ x = [u.flatten(2).transpose(1, 2) for u in x]
+
+ seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+ assert seq_lens.max() <= seq_len
+ x = torch.cat([
+ torch.cat([u, u.new_zeros(1, seq_lens[0] - u.size(1), u.size(2))],
+ dim=1) for u in x
+ ])
+
+ # time embeddings
+ # with amp.autocast(dtype=torch.float32):
+ e = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
+ e0 = self.time_projection(e).unflatten(
+ 1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
+ # assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+ # context
+ context_lens = None
+ context = self.text_embedding(
+ torch.stack([
+ torch.cat(
+ [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+ for u in context
+ ]))
+
+ if clip_fea is not None:
+ context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+ context = torch.concat([context_clip, context], dim=1)
+
+ if clean_x is not None:
+ clean_x = [self.patch_embedding(u.unsqueeze(0)) for u in clean_x]
+ clean_x = [u.flatten(2).transpose(1, 2) for u in clean_x]
+
+ seq_lens_clean = torch.tensor([u.size(1) for u in clean_x], dtype=torch.long)
+ assert seq_lens_clean.max() <= seq_len
+ clean_x = torch.cat([
+ torch.cat([u, u.new_zeros(1, seq_lens_clean[0] - u.size(1), u.size(2))], dim=1) for u in clean_x
+ ])
+
+ x = torch.cat([clean_x, x], dim=1)
+ if aug_t is None:
+ aug_t = torch.zeros_like(t)
+ e_clean = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, aug_t.flatten()).type_as(x))
+ e0_clean = self.time_projection(e_clean).unflatten(
+ 1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
+ e0 = torch.cat([e0_clean, e0], dim=1)
+
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens,
+ block_mask=self.block_mask)
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs, **kwargs):
+ if bool(getattr(
+ self,
+ "gradient_checkpointing_activation_cpu_offload",
+ False,
+ )):
+ with torch.autograd.graph.save_on_cpu(
+ pin_memory=False, device_type=inputs[0].device.type):
+ return module(*inputs, **kwargs)
+ return module(*inputs, **kwargs)
+ return custom_forward
+
+ for block in self.blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x, **kwargs,
+ use_reentrant=False,
+ )
+ else:
+ x = block(x, **kwargs)
+
+ if clean_x is not None:
+ x = x[:, x.shape[1] // 2:]
+
+ # head
+ x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+ return torch.stack(x)
+
+ def forward(
+ self,
+ *args,
+ **kwargs
+ ):
+ if kwargs.get('kv_cache', None) is not None:
+ return self._forward_inference(*args, **kwargs)
+ else:
+ return self._forward_train(*args, **kwargs)
+
+ def unpatchify(self, x, grid_sizes):
+ r"""
+ Reconstruct video tensors from patch embeddings.
+
+ Args:
+ x (List[Tensor]):
+ List of patchified features, each with shape [L, C_out * prod(patch_size)]
+ grid_sizes (Tensor):
+ Original spatial-temporal grid dimensions before patching,
+ shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+ Returns:
+ List[Tensor]:
+ Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+ """
+
+ c = self.out_dim
+ out = []
+ for u, v in zip(x, grid_sizes.tolist()):
+ u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+ u = torch.einsum('fhwpqrc->cfphqwr', u)
+ u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+ out.append(u)
+ return out
+
+ def init_weights(self):
+ r"""
+ Initialize model parameters using Xavier initialization.
+ """
+
+ # basic init
+ for m in self.modules():
+ if isinstance(m, nn.Linear):
+ nn.init.xavier_uniform_(m.weight)
+ if m.bias is not None:
+ nn.init.zeros_(m.bias)
+
+ # init embeddings
+ nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+ for m in self.text_embedding.modules():
+ if isinstance(m, nn.Linear):
+ nn.init.normal_(m.weight, std=.02)
+ for m in self.time_embedding.modules():
+ if isinstance(m, nn.Linear):
+ nn.init.normal_(m.weight, std=.02)
+
+ # init output layer
+ nn.init.zeros_(self.head.head.weight)
diff --git a/14B/code/wan/modules/model.py b/14B/code/wan/modules/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..4329f79988cf1a5ac69f5a3eb7e32e246afa0d22
--- /dev/null
+++ b/14B/code/wan/modules/model.py
@@ -0,0 +1,1103 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import math
+
+import torch
+import torch.nn as nn
+import torch.utils.checkpoint
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from einops import repeat
+
+from .attention import attention, flash_attention
+
+__all__ = ['WanModel']
+
+
+def sinusoidal_embedding_1d(dim, position):
+ # preprocess
+ assert dim % 2 == 0
+ half = dim // 2
+ position = position.type(torch.float64)
+
+ # calculation
+ sinusoid = torch.outer(
+ position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
+ x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+ return x
+
+
+# @amp.autocast(enabled=False)
+def rope_params(max_seq_len, dim, theta=10000):
+ assert dim % 2 == 0
+ freqs = torch.outer(
+ torch.arange(max_seq_len),
+ 1.0 / torch.pow(theta,
+ torch.arange(0, dim, 2).to(torch.float64).div(dim)))
+ freqs = torch.polar(torch.ones_like(freqs), freqs)
+ return freqs
+
+
+# @amp.autocast(enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+ n, c = x.size(2), x.size(3) // 2
+
+ # split freqs
+ freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+ # loop over samples
+ output = []
+ for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+ seq_len = f * h * w
+
+ # Match the original Wan implementation exactly. RoPE is evaluated in
+ # fp64/complex128 and cast back only after the rotation; changing this
+ # globally also changes the causal student's pretrained forward.
+ x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
+ seq_len, n, -1, 2))
+ freqs_i = torch.cat([
+ freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+ freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+ freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+ ],
+ dim=-1).reshape(seq_len, 1, -1)
+
+ # apply rotary embedding
+ x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+ x_i = torch.cat([x_i, x[i, seq_len:]])
+
+ # append to collection
+ output.append(x_i)
+ return torch.stack(output).type_as(x)
+
+
+class WanRMSNorm(nn.Module):
+
+ def __init__(self, dim, eps=1e-5):
+ super().__init__()
+ self.dim = dim
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def forward(self, x):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ return self._norm(x.float()).type_as(x) * self.weight
+
+ def _norm(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+ def reset_parameters(self):
+ # Required by FSDP's meta-device materialisation path; matches __init__.
+ nn.init.ones_(self.weight)
+
+
+class WanLayerNorm(nn.LayerNorm):
+
+ def __init__(self, dim, eps=1e-6, elementwise_affine=False):
+ super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+ def forward(self, x):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ return super().forward(x).type_as(x)
+
+
+class WanSelfAttention(nn.Module):
+
+ def __init__(self,
+ dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6):
+ assert dim % num_heads == 0
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.window_size = window_size
+ self.qk_norm = qk_norm
+ self.eps = eps
+
+ # layers
+ self.q = nn.Linear(dim, dim)
+ self.k = nn.Linear(dim, dim)
+ self.v = nn.Linear(dim, dim)
+ self.o = nn.Linear(dim, dim)
+ self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def project_kv(self, x, grid_sizes, freqs, token_indices=None):
+ """Project the exact K/V used by attention, optionally retaining tokens.
+
+ This is also used by the last encoder block's KV-only fast path, where Q,
+ attention output, cross-attention and FFN are unnecessary.
+ """
+ b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+ k = self.norm_k(self.k(x)).view(b, s, n, d)
+ v = self.v(x).view(b, s, n, d)
+ k = rope_apply(k, grid_sizes, freqs)
+ if token_indices is not None:
+ indices = token_indices.to(k.device)
+ k = k.index_select(1, indices)
+ v = v.index_select(1, indices)
+ return k, v
+
+ def forward(
+ self,
+ x,
+ seq_lens,
+ grid_sizes,
+ freqs,
+ return_kv=False,
+ kv_token_indices=None,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, num_heads, C / num_heads]
+ seq_lens(Tensor): Shape [B]
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+ # query, key, value function
+ def qkv_fn(x):
+ q = self.norm_q(self.q(x)).view(b, s, n, d)
+ k = self.norm_k(self.k(x)).view(b, s, n, d)
+ v = self.v(x).view(b, s, n, d)
+ return q, k, v
+
+ q, k, v = qkv_fn(x)
+ # These are the tensors that are actually consumed by attention. Keeping
+ # K after RoPE (rather than the pre-RoPE projection) makes the extracted
+ # feature identical to a real self-attention KV cache.
+ k_attn = rope_apply(k, grid_sizes, freqs)
+
+ x = flash_attention(
+ q=rope_apply(q, grid_sizes, freqs),
+ k=k_attn,
+ v=v,
+ k_lens=seq_lens,
+ window_size=self.window_size)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ if return_kv:
+ # Attention needs full K/V, but downstream MMD does not. Slice here,
+ # before returning to the block's cross-attention/FFN, so full K/V do
+ # not stay alive for the remainder of a checkpointed block.
+ if kv_token_indices is not None:
+ indices = kv_token_indices.to(k_attn.device)
+ k_attn = k_attn.index_select(1, indices)
+ v = v.index_select(1, indices)
+ return x, k_attn, v
+ return x
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+
+ def forward(self, x, context, context_lens, crossattn_cache=None):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding.
+ """
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x)).view(b, -1, n, d)
+
+ if crossattn_cache is not None:
+ if not crossattn_cache["is_init"]:
+ crossattn_cache["is_init"] = True
+ k = self.norm_k(self.k(context)).view(b, -1, n, d)
+ v = self.v(context).view(b, -1, n, d)
+ crossattn_cache["k"] = k
+ crossattn_cache["v"] = v
+ else:
+ k = crossattn_cache["k"]
+ v = crossattn_cache["v"]
+ else:
+ k = self.norm_k(self.k(context)).view(b, -1, n, d)
+ v = self.v(context).view(b, -1, n, d)
+
+ # compute attention
+ x = attention(q, k, v, k_lens=context_lens)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class WanGanCrossAttention(WanSelfAttention):
+
+ def forward(self, x, context, crossattn_cache=None):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding.
+ """
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ qq = self.norm_q(self.q(context)).view(b, 1, -1, d)
+
+ kk = self.norm_k(self.k(x)).view(b, -1, n, d)
+ vv = self.v(x).view(b, -1, n, d)
+
+ # compute attention
+ x = flash_attention(qq, kk, vv)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class WanI2VCrossAttention(WanSelfAttention):
+
+ def __init__(self,
+ dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6):
+ super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+ self.k_img = nn.Linear(dim, dim)
+ self.v_img = nn.Linear(dim, dim)
+ # self.alpha = nn.Parameter(torch.zeros((1, )))
+ self.norm_k_img = WanRMSNorm(
+ dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, context, context_lens):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+ context_img = context[:, :257]
+ context = context[:, 257:]
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x)).view(b, -1, n, d)
+ k = self.norm_k(self.k(context)).view(b, -1, n, d)
+ v = self.v(context).view(b, -1, n, d)
+ k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)
+ v_img = self.v_img(context_img).view(b, -1, n, d)
+ img_x = flash_attention(q, k_img, v_img, k_lens=None)
+ # compute attention
+ x = flash_attention(q, k, v, k_lens=context_lens)
+
+ # output
+ x = x.flatten(2)
+ img_x = img_x.flatten(2)
+ x = x + img_x
+ x = self.o(x)
+ return x
+
+
+WAN_CROSSATTENTION_CLASSES = {
+ 't2v_cross_attn': WanT2VCrossAttention,
+ 'i2v_cross_attn': WanI2VCrossAttention,
+}
+
+
+class WanAttentionBlock(nn.Module):
+
+ def __init__(self,
+ cross_attn_type,
+ dim,
+ ffn_dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=False,
+ eps=1e-6):
+ super().__init__()
+ self.dim = dim
+ self.ffn_dim = ffn_dim
+ self.num_heads = num_heads
+ self.window_size = window_size
+ self.qk_norm = qk_norm
+ self.cross_attn_norm = cross_attn_norm
+ self.eps = eps
+
+ # layers
+ self.norm1 = WanLayerNorm(dim, eps)
+ self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,
+ eps)
+ self.norm3 = WanLayerNorm(
+ dim, eps,
+ elementwise_affine=True) if cross_attn_norm else nn.Identity()
+ self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim,
+ num_heads,
+ (-1, -1),
+ qk_norm,
+ eps)
+ self.norm2 = WanLayerNorm(dim, eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+ nn.Linear(ffn_dim, dim))
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+ def forward(
+ self,
+ x,
+ e,
+ seq_lens,
+ grid_sizes,
+ freqs,
+ context,
+ context_lens,
+ return_kv=False,
+ kv_token_indices=None,
+ kv_only=False,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ e(Tensor): Shape [B, 6, C]
+ seq_lens(Tensor): Shape [B], length of each sequence in batch
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ # assert e.dtype == torch.float32
+ # with amp.autocast(dtype=torch.float32):
+ e = (self.modulation + e).chunk(6, dim=1)
+ # assert e[0].dtype == torch.float32
+
+ modulated_x = self.norm1(x) * (1 + e[1]) + e[0]
+ if kv_only:
+ if not return_kv:
+ raise ValueError("kv_only requires return_kv=True")
+ return self.self_attn.project_kv(
+ modulated_x, grid_sizes, freqs, token_indices=kv_token_indices)
+
+ # self-attention
+ attn_out = self.self_attn(
+ modulated_x, seq_lens, grid_sizes, freqs,
+ return_kv=return_kv, kv_token_indices=kv_token_indices)
+ if return_kv:
+ y, k, v = attn_out
+ else:
+ y = attn_out
+ # with amp.autocast(dtype=torch.float32):
+ x = x + y * e[2]
+
+ # cross-attention & ffn function
+ def cross_attn_ffn(x, context, context_lens, e):
+ x = x + self.cross_attn(self.norm3(x), context, context_lens)
+ y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3])
+ # with amp.autocast(dtype=torch.float32):
+ x = x + y * e[5]
+ return x
+
+ x = cross_attn_ffn(x, context, context_lens, e)
+ if return_kv:
+ return x, k, v
+ return x
+
+
+class GanAttentionBlock(nn.Module):
+
+ def __init__(self,
+ dim=1536,
+ ffn_dim=8192,
+ num_heads=12,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=True,
+ eps=1e-6):
+ super().__init__()
+ self.dim = dim
+ self.ffn_dim = ffn_dim
+ self.num_heads = num_heads
+ self.window_size = window_size
+ self.qk_norm = qk_norm
+ self.cross_attn_norm = cross_attn_norm
+ self.eps = eps
+
+ # layers
+ # self.norm1 = WanLayerNorm(dim, eps)
+ # self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,
+ # eps)
+ self.norm3 = WanLayerNorm(
+ dim, eps,
+ elementwise_affine=True) if cross_attn_norm else nn.Identity()
+
+ self.norm2 = WanLayerNorm(dim, eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+ nn.Linear(ffn_dim, dim))
+
+ self.cross_attn = WanGanCrossAttention(dim, num_heads,
+ (-1, -1),
+ qk_norm,
+ eps)
+
+ # modulation
+ # self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+ def forward(
+ self,
+ x,
+ context,
+ # seq_lens,
+ # grid_sizes,
+ # freqs,
+ # context,
+ # context_lens,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ e(Tensor): Shape [B, 6, C]
+ seq_lens(Tensor): Shape [B], length of each sequence in batch
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ # assert e.dtype == torch.float32
+ # with amp.autocast(dtype=torch.float32):
+ # e = (self.modulation + e).chunk(6, dim=1)
+ # assert e[0].dtype == torch.float32
+
+ # # self-attention
+ # y = self.self_attn(
+ # self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes,
+ # freqs)
+ # # with amp.autocast(dtype=torch.float32):
+ # x = x + y * e[2]
+
+ # cross-attention & ffn function
+ def cross_attn_ffn(x, context):
+ token = context + self.cross_attn(self.norm3(x), context)
+ y = self.ffn(self.norm2(token)) + token # * (1 + e[4]) + e[3])
+ # with amp.autocast(dtype=torch.float32):
+ # x = x + y * e[5]
+ return y
+
+ x = cross_attn_ffn(x, context)
+ return x
+
+
+class Head(nn.Module):
+
+ def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+ super().__init__()
+ self.dim = dim
+ self.out_dim = out_dim
+ self.patch_size = patch_size
+ self.eps = eps
+
+ # layers
+ out_dim = math.prod(patch_size) * out_dim
+ self.norm = WanLayerNorm(dim, eps)
+ self.head = nn.Linear(dim, out_dim)
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+ def forward(self, x, e):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ e(Tensor): Shape [B, C]
+ """
+ # assert e.dtype == torch.float32
+ # with amp.autocast(dtype=torch.float32):
+ e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+ x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+ return x
+
+
+class MLPProj(torch.nn.Module):
+
+ def __init__(self, in_dim, out_dim):
+ super().__init__()
+
+ self.proj = torch.nn.Sequential(
+ torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),
+ torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),
+ torch.nn.LayerNorm(out_dim))
+
+ def forward(self, image_embeds):
+ clip_extra_context_tokens = self.proj(image_embeds)
+ return clip_extra_context_tokens
+
+
+class RegisterTokens(nn.Module):
+ def __init__(self, num_registers: int, dim: int):
+ super().__init__()
+ self.register_tokens = nn.Parameter(torch.randn(num_registers, dim) * 0.02)
+ self.rms_norm = WanRMSNorm(dim, eps=1e-6)
+
+ def forward(self):
+ return self.rms_norm(self.register_tokens)
+
+ def reset_parameters(self):
+ nn.init.normal_(self.register_tokens, std=0.02)
+
+
+class WanModel(ModelMixin, ConfigMixin):
+ r"""
+ Wan diffusion backbone supporting both text-to-video and image-to-video.
+ """
+
+ ignore_for_config = [
+ 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+ ]
+ _no_split_modules = ['WanAttentionBlock']
+ _supports_gradient_checkpointing = True
+
+ @register_to_config
+ def __init__(self,
+ model_type='t2v',
+ patch_size=(1, 2, 2),
+ text_len=512,
+ in_dim=16,
+ dim=2048,
+ ffn_dim=8192,
+ freq_dim=256,
+ text_dim=4096,
+ out_dim=16,
+ num_heads=16,
+ num_layers=32,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=True,
+ eps=1e-6):
+ r"""
+ Initialize the diffusion model backbone.
+
+ Args:
+ model_type (`str`, *optional*, defaults to 't2v'):
+ Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+ patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+ 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+ text_len (`int`, *optional*, defaults to 512):
+ Fixed length for text embeddings
+ in_dim (`int`, *optional*, defaults to 16):
+ Input video channels (C_in)
+ dim (`int`, *optional*, defaults to 2048):
+ Hidden dimension of the transformer
+ ffn_dim (`int`, *optional*, defaults to 8192):
+ Intermediate dimension in feed-forward network
+ freq_dim (`int`, *optional*, defaults to 256):
+ Dimension for sinusoidal time embeddings
+ text_dim (`int`, *optional*, defaults to 4096):
+ Input dimension for text embeddings
+ out_dim (`int`, *optional*, defaults to 16):
+ Output video channels (C_out)
+ num_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads
+ num_layers (`int`, *optional*, defaults to 32):
+ Number of transformer blocks
+ window_size (`tuple`, *optional*, defaults to (-1, -1)):
+ Window size for local attention (-1 indicates global attention)
+ qk_norm (`bool`, *optional*, defaults to True):
+ Enable query/key normalization
+ cross_attn_norm (`bool`, *optional*, defaults to False):
+ Enable cross-attention normalization
+ eps (`float`, *optional*, defaults to 1e-6):
+ Epsilon value for normalization layers
+ """
+
+ super().__init__()
+
+ assert model_type in ['t2v', 'i2v']
+ self.model_type = model_type
+
+ self.patch_size = patch_size
+ self.text_len = text_len
+ self.in_dim = in_dim
+ self.dim = dim
+ self.ffn_dim = ffn_dim
+ self.freq_dim = freq_dim
+ self.text_dim = text_dim
+ self.out_dim = out_dim
+ self.num_heads = num_heads
+ self.num_layers = num_layers
+ self.window_size = window_size
+ self.qk_norm = qk_norm
+ self.cross_attn_norm = cross_attn_norm
+ self.eps = eps
+ self.local_attn_size = 21
+
+ # embeddings
+ self.patch_embedding = nn.Conv3d(
+ in_dim, dim, kernel_size=patch_size, stride=patch_size)
+ self.text_embedding = nn.Sequential(
+ nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
+ nn.Linear(dim, dim))
+
+ self.time_embedding = nn.Sequential(
+ nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+ self.time_projection = nn.Sequential(
+ nn.SiLU(), nn.Linear(dim, dim * 6))
+
+ # blocks
+ cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
+ self.blocks = nn.ModuleList([
+ WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+ window_size, qk_norm, cross_attn_norm, eps)
+ for _ in range(num_layers)
+ ])
+
+ # head
+ self.head = Head(dim, out_dim, patch_size, eps)
+
+ # buffers (don't use register_buffer otherwise dtype will be changed in to())
+ assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+ d = dim // num_heads
+ self.freqs = torch.cat([
+ rope_params(1024, d - 4 * (d // 6)),
+ rope_params(1024, 2 * (d // 6)),
+ rope_params(1024, 2 * (d // 6))
+ ],
+ dim=1)
+
+ if model_type == 'i2v':
+ self.img_emb = MLPProj(1280, dim)
+
+ # initialize weights
+ self.init_weights()
+
+ self.gradient_checkpointing = False
+
+ def _set_gradient_checkpointing(
+ self,
+ module=None,
+ value=False,
+ enable=None,
+ gradient_checkpointing_func=None,
+ ):
+ self.gradient_checkpointing = value if enable is None else enable
+
+ def forward(
+ self,
+ *args,
+ **kwargs
+ ):
+ # if kwargs.get('classify_mode', False) is True:
+ # kwargs.pop('classify_mode')
+ # return self._forward_classify(*args, **kwargs)
+ # else:
+ return self._forward(*args, **kwargs)
+
+ def _forward(
+ self,
+ x,
+ t,
+ context,
+ seq_len,
+ classify_mode=False,
+ concat_time_embeddings=False,
+ register_tokens=None,
+ cls_pred_branch=None,
+ gan_ca_blocks=None,
+ clip_fea=None,
+ y=None,
+ kv_feature_layers=None,
+ kv_token_indices=None,
+ latent_feature_layer=None,
+ latent_token_indices=None,
+ ):
+ r"""
+ Forward pass through the diffusion model
+
+ Args:
+ x (List[Tensor]):
+ List of input video tensors, each with shape [C_in, F, H, W]
+ t (Tensor):
+ Diffusion timesteps tensor of shape [B]
+ context (List[Tensor]):
+ List of text embeddings each with shape [L, C]
+ seq_len (`int`):
+ Maximum sequence length for positional encoding
+ clip_fea (Tensor, *optional*):
+ CLIP image features for image-to-video mode
+ y (List[Tensor], *optional*):
+ Conditional video inputs for image-to-video mode, same shape as x
+
+ Returns:
+ List[Tensor]:
+ List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+ """
+ if self.model_type == 'i2v':
+ assert clip_fea is not None and y is not None
+ # params
+ device = self.patch_embedding.weight.device
+ if self.freqs.device != device:
+ self.freqs = self.freqs.to(device)
+
+ if y is not None:
+ x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+ # embeddings
+ x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+ grid_sizes = torch.stack(
+ [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+ x = [u.flatten(2).transpose(1, 2) for u in x]
+ seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+ assert seq_lens.max() <= seq_len
+ if kv_feature_layers is not None and latent_feature_layer is not None:
+ raise ValueError(
+ "KV and post-block latent feature extraction are mutually exclusive")
+ if kv_token_indices is not None and kv_feature_layers is None:
+ raise ValueError(
+ "kv_token_indices requires kv_feature_layers")
+ if latent_token_indices is not None and latent_feature_layer is None:
+ raise ValueError(
+ "latent_token_indices requires latent_feature_layer")
+
+ def validate_token_indices(name, indices):
+ if indices is None:
+ return
+ if not isinstance(indices, torch.Tensor):
+ raise ValueError(f"{name} must be a torch.Tensor")
+ if indices.ndim != 1 or indices.dtype != torch.long:
+ raise ValueError(f"{name} must be a 1-D torch.long tensor")
+ if indices.numel() == 0:
+ raise ValueError(f"{name} cannot be empty")
+ min_index = int(indices.min().item())
+ max_index = int(indices.max().item())
+ if min_index < 0 or max_index >= int(seq_lens.min().item()):
+ raise ValueError(
+ f"{name} range [{min_index}, {max_index}] exceeds "
+ f"the shortest valid token sequence ({int(seq_lens.min().item())})")
+
+ if kv_feature_layers is not None:
+ validate_token_indices("kv_token_indices", kv_token_indices)
+ if latent_feature_layer is not None:
+ if (
+ isinstance(latent_feature_layer, bool)
+ or not isinstance(latent_feature_layer, int)
+ ):
+ raise ValueError("latent_feature_layer must be an integer block index")
+ if latent_feature_layer < 0 or latent_feature_layer >= len(self.blocks):
+ raise ValueError(
+ f"latent_feature_layer={latent_feature_layer} outside "
+ f"[0, {len(self.blocks) - 1}]")
+ validate_token_indices("latent_token_indices", latent_token_indices)
+ x = torch.cat([
+ torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+ dim=1) for u in x
+ ])
+
+ # time embeddings
+ # with amp.autocast(dtype=torch.float32):
+ e = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, t).type_as(x))
+ e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+ # assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+ # context
+ context_lens = None
+ context = self.text_embedding(
+ torch.stack([
+ torch.cat(
+ [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+ for u in context
+ ]))
+
+ if clip_fea is not None:
+ context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+ context = torch.concat([context_clip, context], dim=1)
+
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens)
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs, **kwargs):
+ return module(*inputs, **kwargs)
+ return custom_forward
+
+ # Frozen-Wan post-block latent encoder path. Unlike the final KV-only
+ # shortcut below, the requested block runs in full (self-attention,
+ # cross-attention, FFN and all residual updates). The diffusion head
+ # and all later blocks are intentionally skipped.
+ if latent_feature_layer is not None:
+ for ii, block in enumerate(self.blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x, **kwargs,
+ use_reentrant=False,
+ )
+ else:
+ x = block(x, **kwargs)
+ if ii == latent_feature_layer:
+ if x.ndim != 3 or x.shape[-1] != self.dim:
+ raise RuntimeError(
+ "post-block latent must have shape [batch, tokens, dim]")
+ if latent_token_indices is not None:
+ indices = latent_token_indices.to(device=x.device)
+ x = x.index_select(1, indices)
+ return x
+ raise RuntimeError(
+ "latent feature block was not reached despite validated index")
+
+ # Frozen-Wan feature-encoder path. Only the prefix ending at the last
+ # requested block is evaluated. K/V are concatenated along the feature
+ # dimension and remain a [B, tokens, features] matrix; no global pooling
+ # or vector reshape is applied.
+ if kv_feature_layers is not None:
+ capture = sorted({int(i) for i in kv_feature_layers})
+ if not capture:
+ raise ValueError("kv_feature_layers must contain at least one block index")
+ if capture[0] < 0 or capture[-1] >= len(self.blocks):
+ raise ValueError(
+ f"kv_feature_layers={capture} outside [0, {len(self.blocks) - 1}]")
+ captured = []
+ capture_set = set(capture)
+ last_capture = capture[-1]
+ for ii, block in enumerate(self.blocks):
+ if ii > last_capture:
+ break
+ block_kwargs = dict(kwargs)
+ block_kwargs["return_kv"] = ii in capture_set
+ block_kwargs["kv_token_indices"] = (
+ kv_token_indices if ii in capture_set else None)
+ # The last requested feature only needs this block's modulated
+ # K/V projections. No later block consumes its full output.
+ block_kwargs["kv_only"] = ii == last_capture
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ out = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x, **block_kwargs,
+ use_reentrant=False,
+ )
+ else:
+ out = block(x, **block_kwargs)
+ if ii == last_capture:
+ k, v = out
+ kv = torch.cat([k.flatten(2), v.flatten(2)], dim=-1)
+ captured.append(kv)
+ break
+ if ii in capture_set:
+ x, k, v = out
+ # [B,L,H,Dh] -> [B,L,2*D].
+ kv = torch.cat([k.flatten(2), v.flatten(2)], dim=-1)
+ captured.append(kv)
+ else:
+ x = out
+ return torch.cat(captured, dim=-1)
+
+ # TODO: Tune the number of blocks for feature extraction
+ final_x = None
+ if classify_mode:
+ assert register_tokens is not None
+ assert gan_ca_blocks is not None
+ assert cls_pred_branch is not None
+
+ final_x = []
+ registers = repeat(register_tokens(), "n d -> b n d", b=x.shape[0])
+ # x = torch.cat([registers, x], dim=1)
+
+ gan_idx = 0
+ for ii, block in enumerate(self.blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x, **kwargs,
+ use_reentrant=False,
+ )
+ else:
+ x = block(x, **kwargs)
+
+ if classify_mode and ii in [13, 21, 29]:
+ gan_token = registers[:, gan_idx: gan_idx + 1]
+ final_x.append(gan_ca_blocks[gan_idx](x, gan_token))
+ gan_idx += 1
+
+ if classify_mode:
+ final_x = torch.cat(final_x, dim=1)
+ if concat_time_embeddings:
+ final_x = cls_pred_branch(torch.cat([final_x, 10 * e[:, None, :]], dim=1).view(final_x.shape[0], -1))
+ else:
+ final_x = cls_pred_branch(final_x.view(final_x.shape[0], -1))
+
+ # head
+ x = self.head(x, e)
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+
+ if classify_mode:
+ return torch.stack(x), final_x
+
+ return torch.stack(x)
+
+ def _forward_classify(
+ self,
+ x,
+ t,
+ context,
+ seq_len,
+ register_tokens,
+ cls_pred_branch,
+ clip_fea=None,
+ y=None,
+ ):
+ r"""
+ Feature extraction through the diffusion model
+
+ Args:
+ x (List[Tensor]):
+ List of input video tensors, each with shape [C_in, F, H, W]
+ t (Tensor):
+ Diffusion timesteps tensor of shape [B]
+ context (List[Tensor]):
+ List of text embeddings each with shape [L, C]
+ seq_len (`int`):
+ Maximum sequence length for positional encoding
+ clip_fea (Tensor, *optional*):
+ CLIP image features for image-to-video mode
+ y (List[Tensor], *optional*):
+ Conditional video inputs for image-to-video mode, same shape as x
+
+ Returns:
+ List[Tensor]:
+ List of video features with original input shapes [C_block, F, H / 8, W / 8]
+ """
+ if self.model_type == 'i2v':
+ assert clip_fea is not None and y is not None
+ # params
+ device = self.patch_embedding.weight.device
+ if self.freqs.device != device:
+ self.freqs = self.freqs.to(device)
+
+ if y is not None:
+ x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+ # embeddings
+ x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+ grid_sizes = torch.stack(
+ [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+ x = [u.flatten(2).transpose(1, 2) for u in x]
+ seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+ assert seq_lens.max() <= seq_len
+ x = torch.cat([
+ torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+ dim=1) for u in x
+ ])
+
+ # time embeddings
+ # with amp.autocast(dtype=torch.float32):
+ e = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, t).type_as(x))
+ e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+ # assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+ # context
+ context_lens = None
+ context = self.text_embedding(
+ torch.stack([
+ torch.cat(
+ [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+ for u in context
+ ]))
+
+ if clip_fea is not None:
+ context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+ context = torch.concat([context_clip, context], dim=1)
+
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens)
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs, **kwargs):
+ return module(*inputs, **kwargs)
+ return custom_forward
+
+ # TODO: Tune the number of blocks for feature extraction
+ for block in self.blocks[:16]:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x, **kwargs,
+ use_reentrant=False,
+ )
+ else:
+ x = block(x, **kwargs)
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes, c=self.dim // 4)
+ return torch.stack(x)
+
+ def unpatchify(self, x, grid_sizes, c=None):
+ r"""
+ Reconstruct video tensors from patch embeddings.
+
+ Args:
+ x (List[Tensor]):
+ List of patchified features, each with shape [L, C_out * prod(patch_size)]
+ grid_sizes (Tensor):
+ Original spatial-temporal grid dimensions before patching,
+ shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+ Returns:
+ List[Tensor]:
+ Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+ """
+
+ c = self.out_dim if c is None else c
+ out = []
+ for u, v in zip(x, grid_sizes.tolist()):
+ u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+ u = torch.einsum('fhwpqrc->cfphqwr', u)
+ u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+ out.append(u)
+ return out
+
+ def init_weights(self):
+ r"""
+ Initialize model parameters using Xavier initialization.
+ """
+
+ # basic init
+ for m in self.modules():
+ if isinstance(m, nn.Linear):
+ nn.init.xavier_uniform_(m.weight)
+ if m.bias is not None:
+ nn.init.zeros_(m.bias)
+
+ # init embeddings
+ nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+ for m in self.text_embedding.modules():
+ if isinstance(m, nn.Linear):
+ nn.init.normal_(m.weight, std=.02)
+ for m in self.time_embedding.modules():
+ if isinstance(m, nn.Linear):
+ nn.init.normal_(m.weight, std=.02)
+
+ # init output layer
+ nn.init.zeros_(self.head.head.weight)
diff --git a/14B/code/wan/modules/t5.py b/14B/code/wan/modules/t5.py
new file mode 100644
index 0000000000000000000000000000000000000000..9438874681ea14ab1f4368869b474ce7c8dbcd80
--- /dev/null
+++ b/14B/code/wan/modules/t5.py
@@ -0,0 +1,515 @@
+# Modified from transformers.models.t5.modeling_t5
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .tokenizers import HuggingfaceTokenizer
+
+__all__ = [
+ 'T5Model',
+ 'T5Encoder',
+ 'T5Decoder',
+ 'T5EncoderModel',
+]
+
+
+def fp16_clamp(x):
+ if x.dtype == torch.float16 and torch.isinf(x).any():
+ clamp = torch.finfo(x.dtype).max - 1000
+ x = torch.clamp(x, min=-clamp, max=clamp)
+ return x
+
+
+def init_weights(m):
+ if isinstance(m, T5LayerNorm):
+ nn.init.ones_(m.weight)
+ elif isinstance(m, T5Model):
+ nn.init.normal_(m.token_embedding.weight, std=1.0)
+ elif isinstance(m, T5FeedForward):
+ nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)
+ nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)
+ nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)
+ elif isinstance(m, T5Attention):
+ nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5)
+ nn.init.normal_(m.k.weight, std=m.dim**-0.5)
+ nn.init.normal_(m.v.weight, std=m.dim**-0.5)
+ nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5)
+ elif isinstance(m, T5RelativeEmbedding):
+ nn.init.normal_(
+ m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5)
+
+
+class GELU(nn.Module):
+
+ def forward(self, x):
+ return 0.5 * x * (1.0 + torch.tanh(
+ math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+
+ def __init__(self, dim, eps=1e-6):
+ super(T5LayerNorm, self).__init__()
+ self.dim = dim
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def forward(self, x):
+ x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) +
+ self.eps)
+ if self.weight.dtype in [torch.float16, torch.bfloat16]:
+ x = x.type_as(self.weight)
+ return self.weight * x
+
+
+class T5Attention(nn.Module):
+
+ def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
+ assert dim_attn % num_heads == 0
+ super(T5Attention, self).__init__()
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.num_heads = num_heads
+ self.head_dim = dim_attn // num_heads
+
+ # layers
+ self.q = nn.Linear(dim, dim_attn, bias=False)
+ self.k = nn.Linear(dim, dim_attn, bias=False)
+ self.v = nn.Linear(dim, dim_attn, bias=False)
+ self.o = nn.Linear(dim_attn, dim, bias=False)
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, x, context=None, mask=None, pos_bias=None):
+ """
+ x: [B, L1, C].
+ context: [B, L2, C] or None.
+ mask: [B, L2] or [B, L1, L2] or None.
+ """
+ # check inputs
+ context = x if context is None else context
+ b, n, c = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.q(x).view(b, -1, n, c)
+ k = self.k(context).view(b, -1, n, c)
+ v = self.v(context).view(b, -1, n, c)
+
+ # attention bias
+ attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+ if pos_bias is not None:
+ attn_bias += pos_bias
+ if mask is not None:
+ assert mask.ndim in [2, 3]
+ mask = mask.view(b, 1, 1,
+ -1) if mask.ndim == 2 else mask.unsqueeze(1)
+ attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+
+ # compute attention (T5 does not use scaling)
+ attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias
+ attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+ x = torch.einsum('bnij,bjnc->binc', attn, v)
+
+ # output
+ x = x.reshape(b, -1, n * c)
+ x = self.o(x)
+ x = self.dropout(x)
+ return x
+
+
+class T5FeedForward(nn.Module):
+
+ def __init__(self, dim, dim_ffn, dropout=0.1):
+ super(T5FeedForward, self).__init__()
+ self.dim = dim
+ self.dim_ffn = dim_ffn
+
+ # layers
+ self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+ self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+ self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, x):
+ x = self.fc1(x) * self.gate(x)
+ x = self.dropout(x)
+ x = self.fc2(x)
+ x = self.dropout(x)
+ return x
+
+
+class T5SelfAttention(nn.Module):
+
+ def __init__(self,
+ dim,
+ dim_attn,
+ dim_ffn,
+ num_heads,
+ num_buckets,
+ shared_pos=True,
+ dropout=0.1):
+ super(T5SelfAttention, self).__init__()
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.dim_ffn = dim_ffn
+ self.num_heads = num_heads
+ self.num_buckets = num_buckets
+ self.shared_pos = shared_pos
+
+ # layers
+ self.norm1 = T5LayerNorm(dim)
+ self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+ self.norm2 = T5LayerNorm(dim)
+ self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+ self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+ num_buckets, num_heads, bidirectional=True)
+
+ def forward(self, x, mask=None, pos_bias=None):
+ e = pos_bias if self.shared_pos else self.pos_embedding(
+ x.size(1), x.size(1))
+ x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+ x = fp16_clamp(x + self.ffn(self.norm2(x)))
+ return x
+
+
+class T5CrossAttention(nn.Module):
+
+ def __init__(self,
+ dim,
+ dim_attn,
+ dim_ffn,
+ num_heads,
+ num_buckets,
+ shared_pos=True,
+ dropout=0.1):
+ super(T5CrossAttention, self).__init__()
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.dim_ffn = dim_ffn
+ self.num_heads = num_heads
+ self.num_buckets = num_buckets
+ self.shared_pos = shared_pos
+
+ # layers
+ self.norm1 = T5LayerNorm(dim)
+ self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+ self.norm2 = T5LayerNorm(dim)
+ self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+ self.norm3 = T5LayerNorm(dim)
+ self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+ self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+ num_buckets, num_heads, bidirectional=False)
+
+ def forward(self,
+ x,
+ mask=None,
+ encoder_states=None,
+ encoder_mask=None,
+ pos_bias=None):
+ e = pos_bias if self.shared_pos else self.pos_embedding(
+ x.size(1), x.size(1))
+ x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e))
+ x = fp16_clamp(x + self.cross_attn(
+ self.norm2(x), context=encoder_states, mask=encoder_mask))
+ x = fp16_clamp(x + self.ffn(self.norm3(x)))
+ return x
+
+
+class T5RelativeEmbedding(nn.Module):
+
+ def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
+ super(T5RelativeEmbedding, self).__init__()
+ self.num_buckets = num_buckets
+ self.num_heads = num_heads
+ self.bidirectional = bidirectional
+ self.max_dist = max_dist
+
+ # layers
+ self.embedding = nn.Embedding(num_buckets, num_heads)
+
+ def forward(self, lq, lk):
+ device = self.embedding.weight.device
+ # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \
+ # torch.arange(lq).unsqueeze(1).to(device)
+ rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \
+ torch.arange(lq, device=device).unsqueeze(1)
+ rel_pos = self._relative_position_bucket(rel_pos)
+ rel_pos_embeds = self.embedding(rel_pos)
+ rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(
+ 0) # [1, N, Lq, Lk]
+ return rel_pos_embeds.contiguous()
+
+ def _relative_position_bucket(self, rel_pos):
+ # preprocess
+ if self.bidirectional:
+ num_buckets = self.num_buckets // 2
+ rel_buckets = (rel_pos > 0).long() * num_buckets
+ rel_pos = torch.abs(rel_pos)
+ else:
+ num_buckets = self.num_buckets
+ rel_buckets = 0
+ rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+
+ # embeddings for small and large positions
+ max_exact = num_buckets // 2
+ rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) /
+ math.log(self.max_dist / max_exact) *
+ (num_buckets - max_exact)).long()
+ rel_pos_large = torch.min(
+ rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+ rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+ return rel_buckets
+
+
+class T5Encoder(nn.Module):
+
+ def __init__(self,
+ vocab,
+ dim,
+ dim_attn,
+ dim_ffn,
+ num_heads,
+ num_layers,
+ num_buckets,
+ shared_pos=True,
+ dropout=0.1):
+ super(T5Encoder, self).__init__()
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.dim_ffn = dim_ffn
+ self.num_heads = num_heads
+ self.num_layers = num_layers
+ self.num_buckets = num_buckets
+ self.shared_pos = shared_pos
+
+ # layers
+ self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+ else nn.Embedding(vocab, dim)
+ self.pos_embedding = T5RelativeEmbedding(
+ num_buckets, num_heads, bidirectional=True) if shared_pos else None
+ self.dropout = nn.Dropout(dropout)
+ self.blocks = nn.ModuleList([
+ T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+ shared_pos, dropout) for _ in range(num_layers)
+ ])
+ self.norm = T5LayerNorm(dim)
+
+ # initialize weights
+ self.apply(init_weights)
+
+ def forward(self, ids, mask=None):
+ x = self.token_embedding(ids)
+ x = self.dropout(x)
+ e = self.pos_embedding(x.size(1),
+ x.size(1)) if self.shared_pos else None
+ for block in self.blocks:
+ x = block(x, mask, pos_bias=e)
+ x = self.norm(x)
+ x = self.dropout(x)
+ return x
+
+
+class T5Decoder(nn.Module):
+
+ def __init__(self,
+ vocab,
+ dim,
+ dim_attn,
+ dim_ffn,
+ num_heads,
+ num_layers,
+ num_buckets,
+ shared_pos=True,
+ dropout=0.1):
+ super(T5Decoder, self).__init__()
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.dim_ffn = dim_ffn
+ self.num_heads = num_heads
+ self.num_layers = num_layers
+ self.num_buckets = num_buckets
+ self.shared_pos = shared_pos
+
+ # layers
+ self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+ else nn.Embedding(vocab, dim)
+ self.pos_embedding = T5RelativeEmbedding(
+ num_buckets, num_heads, bidirectional=False) if shared_pos else None
+ self.dropout = nn.Dropout(dropout)
+ self.blocks = nn.ModuleList([
+ T5CrossAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+ shared_pos, dropout) for _ in range(num_layers)
+ ])
+ self.norm = T5LayerNorm(dim)
+
+ # initialize weights
+ self.apply(init_weights)
+
+ def forward(self, ids, mask=None, encoder_states=None, encoder_mask=None):
+ b, s = ids.size()
+
+ # causal mask
+ if mask is None:
+ mask = torch.tril(torch.ones(1, s, s).to(ids.device))
+ elif mask.ndim == 2:
+ mask = torch.tril(mask.unsqueeze(1).expand(-1, s, -1))
+
+ # layers
+ x = self.token_embedding(ids)
+ x = self.dropout(x)
+ e = self.pos_embedding(x.size(1),
+ x.size(1)) if self.shared_pos else None
+ for block in self.blocks:
+ x = block(x, mask, encoder_states, encoder_mask, pos_bias=e)
+ x = self.norm(x)
+ x = self.dropout(x)
+ return x
+
+
+class T5Model(nn.Module):
+
+ def __init__(self,
+ vocab_size,
+ dim,
+ dim_attn,
+ dim_ffn,
+ num_heads,
+ encoder_layers,
+ decoder_layers,
+ num_buckets,
+ shared_pos=True,
+ dropout=0.1):
+ super(T5Model, self).__init__()
+ self.vocab_size = vocab_size
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.dim_ffn = dim_ffn
+ self.num_heads = num_heads
+ self.encoder_layers = encoder_layers
+ self.decoder_layers = decoder_layers
+ self.num_buckets = num_buckets
+
+ # layers
+ self.token_embedding = nn.Embedding(vocab_size, dim)
+ self.encoder = T5Encoder(self.token_embedding, dim, dim_attn, dim_ffn,
+ num_heads, encoder_layers, num_buckets,
+ shared_pos, dropout)
+ self.decoder = T5Decoder(self.token_embedding, dim, dim_attn, dim_ffn,
+ num_heads, decoder_layers, num_buckets,
+ shared_pos, dropout)
+ self.head = nn.Linear(dim, vocab_size, bias=False)
+
+ # initialize weights
+ self.apply(init_weights)
+
+ def forward(self, encoder_ids, encoder_mask, decoder_ids, decoder_mask):
+ x = self.encoder(encoder_ids, encoder_mask)
+ x = self.decoder(decoder_ids, decoder_mask, x, encoder_mask)
+ x = self.head(x)
+ return x
+
+
+def _t5(name,
+ encoder_only=False,
+ decoder_only=False,
+ return_tokenizer=False,
+ tokenizer_kwargs={},
+ dtype=torch.float32,
+ device='cpu',
+ **kwargs):
+ # sanity check
+ assert not (encoder_only and decoder_only)
+
+ # params
+ if encoder_only:
+ model_cls = T5Encoder
+ kwargs['vocab'] = kwargs.pop('vocab_size')
+ kwargs['num_layers'] = kwargs.pop('encoder_layers')
+ _ = kwargs.pop('decoder_layers')
+ elif decoder_only:
+ model_cls = T5Decoder
+ kwargs['vocab'] = kwargs.pop('vocab_size')
+ kwargs['num_layers'] = kwargs.pop('decoder_layers')
+ _ = kwargs.pop('encoder_layers')
+ else:
+ model_cls = T5Model
+
+ # init model
+ with torch.device(device):
+ model = model_cls(**kwargs)
+
+ # set device
+ model = model.to(dtype=dtype, device=device)
+
+ # init tokenizer
+ if return_tokenizer:
+ from .tokenizers import HuggingfaceTokenizer
+ tokenizer_options = dict(tokenizer_kwargs)
+ tokenizer_path = tokenizer_options.pop("name", f"assets/models/tokenizers/{name}")
+ tokenizer = HuggingfaceTokenizer(tokenizer_path, **tokenizer_options)
+ return model, tokenizer
+ else:
+ return model
+
+
+def umt5_xxl(**kwargs):
+ cfg = dict(
+ vocab_size=256384,
+ dim=4096,
+ dim_attn=4096,
+ dim_ffn=10240,
+ num_heads=64,
+ encoder_layers=24,
+ decoder_layers=24,
+ num_buckets=32,
+ shared_pos=False,
+ dropout=0.1)
+ cfg.update(**kwargs)
+ return _t5('umt5-xxl', **cfg)
+
+
+class T5EncoderModel:
+
+ def __init__(
+ self,
+ text_len,
+ dtype=torch.bfloat16,
+ device=torch.cuda.current_device(),
+ checkpoint_path=None,
+ tokenizer_path=None,
+ shard_fn=None,
+ ):
+ self.text_len = text_len
+ self.dtype = dtype
+ self.device = device
+ self.checkpoint_path = checkpoint_path
+ self.tokenizer_path = tokenizer_path
+
+ # init model
+ model = umt5_xxl(
+ encoder_only=True,
+ return_tokenizer=False,
+ dtype=dtype,
+ device=device).eval().requires_grad_(False)
+ logging.info(f'loading {checkpoint_path}')
+ model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
+ self.model = model
+ if shard_fn is not None:
+ self.model = shard_fn(self.model, sync_module_states=False)
+ else:
+ self.model.to(self.device)
+ # init tokenizer
+ self.tokenizer = HuggingfaceTokenizer(
+ name=tokenizer_path, seq_len=text_len, clean='whitespace')
+
+ def __call__(self, texts, device):
+ ids, mask = self.tokenizer(
+ texts, return_mask=True, add_special_tokens=True)
+ ids = ids.to(device)
+ mask = mask.to(device)
+ seq_lens = mask.gt(0).sum(dim=1).long()
+ context = self.model(ids, mask)
+ return [u[:v] for u, v in zip(context, seq_lens)]
diff --git a/14B/code/wan/modules/tokenizers.py b/14B/code/wan/modules/tokenizers.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c717e6293b5cd4252dd23b5531faa1627c1248b
--- /dev/null
+++ b/14B/code/wan/modules/tokenizers.py
@@ -0,0 +1,88 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import html
+import string
+from pathlib import Path
+
+import ftfy
+import regex as re
+from transformers import AutoTokenizer
+
+__all__ = ['HuggingfaceTokenizer']
+
+
+def basic_clean(text):
+ text = ftfy.fix_text(text)
+ text = html.unescape(html.unescape(text))
+ return text.strip()
+
+
+def whitespace_clean(text):
+ text = re.sub(r'\s+', ' ', text)
+ text = text.strip()
+ return text
+
+
+def canonicalize(text, keep_punctuation_exact_string=None):
+ text = text.replace('_', ' ')
+ if keep_punctuation_exact_string:
+ text = keep_punctuation_exact_string.join(
+ part.translate(str.maketrans('', '', string.punctuation))
+ for part in text.split(keep_punctuation_exact_string))
+ else:
+ text = text.translate(str.maketrans('', '', string.punctuation))
+ text = text.lower()
+ text = re.sub(r'\s+', ' ', text)
+ return text.strip()
+
+
+class HuggingfaceTokenizer:
+
+ def __init__(self, name, seq_len=None, clean=None, **kwargs):
+ assert clean in (None, 'whitespace', 'lower', 'canonicalize')
+ self.name = name
+ self.seq_len = seq_len
+ self.clean = clean
+
+ # Load the tokenizer from an existing local model bundle.
+ model_dir = Path(name).expanduser()
+ if not model_dir.is_dir():
+ raise FileNotFoundError(f"Prepare the local tokenizer directory: {model_dir}")
+ name = str(model_dir.resolve())
+ kwargs["local_files_only"] = True
+ self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs)
+ self.vocab_size = self.tokenizer.vocab_size
+
+ def __call__(self, sequence, **kwargs):
+ return_mask = kwargs.pop('return_mask', False)
+
+ # arguments
+ _kwargs = {'return_tensors': 'pt'}
+ if self.seq_len is not None:
+ _kwargs.update({
+ 'padding': 'max_length',
+ 'truncation': True,
+ 'max_length': self.seq_len
+ })
+ _kwargs.update(**kwargs)
+
+ # tokenization
+ if isinstance(sequence, str):
+ sequence = [sequence]
+ if self.clean:
+ sequence = [self._clean(u) for u in sequence]
+ ids = self.tokenizer(sequence, **_kwargs)
+
+ # output
+ if return_mask:
+ return ids.input_ids, ids.attention_mask
+ else:
+ return ids.input_ids
+
+ def _clean(self, text):
+ if self.clean == 'whitespace':
+ text = whitespace_clean(basic_clean(text))
+ elif self.clean == 'lower':
+ text = whitespace_clean(basic_clean(text)).lower()
+ elif self.clean == 'canonicalize':
+ text = canonicalize(basic_clean(text))
+ return text
diff --git a/14B/code/wan/modules/vae.py b/14B/code/wan/modules/vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..c50dea913c32eccf971fd528bb15b3173ea5f9b9
--- /dev/null
+++ b/14B/code/wan/modules/vae.py
@@ -0,0 +1,683 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+__all__ = [
+ 'WanVAE',
+]
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+ """
+ Causal 3d convolusion.
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self._padding = (self.padding[2], self.padding[2], self.padding[1],
+ self.padding[1], 2 * self.padding[0], 0)
+ self.padding = (0, 0, 0)
+
+ def forward(self, x, cache_x=None):
+ padding = list(self._padding)
+ if cache_x is not None and self._padding[4] > 0:
+ cache_x = cache_x.to(x.device)
+ x = torch.cat([cache_x, x], dim=2)
+ padding[4] -= cache_x.shape[2]
+ x = F.pad(x, padding)
+
+ return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+ def __init__(self, dim, channel_first=True, images=True, bias=False):
+ super().__init__()
+ broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+ shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+ self.channel_first = channel_first
+ self.scale = dim**0.5
+ self.gamma = nn.Parameter(torch.ones(shape))
+ self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.
+
+ def forward(self, x):
+ return F.normalize(
+ x, dim=(1 if self.channel_first else
+ -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+
+ def forward(self, x):
+ """
+ Fix bfloat16 support for nearest neighbor interpolation.
+ """
+ return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+ def __init__(self, dim, mode):
+ assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',
+ 'downsample3d')
+ super().__init__()
+ self.dim = dim
+ self.mode = mode
+
+ # layers
+ if mode == 'upsample2d':
+ self.resample = nn.Sequential(
+ Upsample(scale_factor=(2., 2.), mode='nearest'),
+ nn.Conv2d(dim, dim // 2, 3, padding=1))
+ elif mode == 'upsample3d':
+ self.resample = nn.Sequential(
+ Upsample(scale_factor=(2., 2.), mode='nearest'),
+ nn.Conv2d(dim, dim // 2, 3, padding=1))
+ self.time_conv = CausalConv3d(
+ dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+
+ elif mode == 'downsample2d':
+ self.resample = nn.Sequential(
+ nn.ZeroPad2d((0, 1, 0, 1)),
+ nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+ elif mode == 'downsample3d':
+ self.resample = nn.Sequential(
+ nn.ZeroPad2d((0, 1, 0, 1)),
+ nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+ self.time_conv = CausalConv3d(
+ dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+
+ else:
+ self.resample = nn.Identity()
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ b, c, t, h, w = x.size()
+ if self.mode == 'upsample3d':
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ if feat_cache[idx] is None:
+ feat_cache[idx] = 'Rep'
+ feat_idx[0] += 1
+ else:
+
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[
+ idx] is not None and feat_cache[idx] != 'Rep':
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ if cache_x.shape[2] < 2 and feat_cache[
+ idx] is not None and feat_cache[idx] == 'Rep':
+ cache_x = torch.cat([
+ torch.zeros_like(cache_x).to(cache_x.device),
+ cache_x
+ ],
+ dim=2)
+ if feat_cache[idx] == 'Rep':
+ x = self.time_conv(x)
+ else:
+ x = self.time_conv(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+
+ x = x.reshape(b, 2, c, t, h, w)
+ x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+ 3)
+ x = x.reshape(b, c, t * 2, h, w)
+ t = x.shape[2]
+ x = rearrange(x, 'b c t h w -> (b t) c h w')
+ x = self.resample(x)
+ x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
+
+ if self.mode == 'downsample3d':
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ if feat_cache[idx] is None:
+ feat_cache[idx] = x.clone()
+ feat_idx[0] += 1
+ else:
+
+ cache_x = x[:, :, -1:, :, :].clone()
+ # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep':
+ # # cache last frame of last two chunk
+ # cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+
+ x = self.time_conv(
+ torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ return x
+
+ def init_weight(self, conv):
+ conv_weight = conv.weight
+ nn.init.zeros_(conv_weight)
+ c1, c2, t, h, w = conv_weight.size()
+ one_matrix = torch.eye(c1, c2)
+ init_matrix = one_matrix
+ nn.init.zeros_(conv_weight)
+ # conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5
+ conv_weight.data[:, :, 1, 0, 0] = init_matrix # * 0.5
+ conv.weight.data.copy_(conv_weight)
+ nn.init.zeros_(conv.bias.data)
+
+ def init_weight2(self, conv):
+ conv_weight = conv.weight.data
+ nn.init.zeros_(conv_weight)
+ c1, c2, t, h, w = conv_weight.size()
+ init_matrix = torch.eye(c1 // 2, c2)
+ # init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)
+ conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+ conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+ conv.weight.data.copy_(conv_weight)
+ nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+ def __init__(self, in_dim, out_dim, dropout=0.0):
+ super().__init__()
+ self.in_dim = in_dim
+ self.out_dim = out_dim
+
+ # layers
+ self.residual = nn.Sequential(
+ RMS_norm(in_dim, images=False), nn.SiLU(),
+ CausalConv3d(in_dim, out_dim, 3, padding=1),
+ RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout),
+ CausalConv3d(out_dim, out_dim, 3, padding=1))
+ self.shortcut = CausalConv3d(in_dim, out_dim, 1) \
+ if in_dim != out_dim else nn.Identity()
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ h = self.shortcut(x)
+ for layer in self.residual:
+ if isinstance(layer, CausalConv3d) and feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = layer(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = layer(x)
+ return x + h
+
+
+class AttentionBlock(nn.Module):
+ """
+ Causal self-attention with a single head.
+ """
+
+ def __init__(self, dim):
+ super().__init__()
+ self.dim = dim
+
+ # layers
+ self.norm = RMS_norm(dim)
+ self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+ self.proj = nn.Conv2d(dim, dim, 1)
+
+ # zero out the last layer params
+ nn.init.zeros_(self.proj.weight)
+
+ def forward(self, x):
+ identity = x
+ b, c, t, h, w = x.size()
+ x = rearrange(x, 'b c t h w -> (b t) c h w')
+ x = self.norm(x)
+ # compute query, key, value
+ q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3,
+ -1).permute(0, 1, 3,
+ 2).contiguous().chunk(
+ 3, dim=-1)
+
+ # apply attention
+ x = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ )
+ x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+ # output
+ x = self.proj(x)
+ x = rearrange(x, '(b t) c h w-> b c t h w', t=t)
+ return x + identity
+
+
+class Encoder3d(nn.Module):
+
+ def __init__(self,
+ dim=128,
+ z_dim=4,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_downsample=[True, True, False],
+ dropout=0.0):
+ super().__init__()
+ self.dim = dim
+ self.z_dim = z_dim
+ self.dim_mult = dim_mult
+ self.num_res_blocks = num_res_blocks
+ self.attn_scales = attn_scales
+ self.temperal_downsample = temperal_downsample
+
+ # dimensions
+ dims = [dim * u for u in [1] + dim_mult]
+ scale = 1.0
+
+ # init block
+ self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
+
+ # downsample blocks
+ downsamples = []
+ for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+ # residual (+attention) blocks
+ for _ in range(num_res_blocks):
+ downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+ if scale in attn_scales:
+ downsamples.append(AttentionBlock(out_dim))
+ in_dim = out_dim
+
+ # downsample block
+ if i != len(dim_mult) - 1:
+ mode = 'downsample3d' if temperal_downsample[
+ i] else 'downsample2d'
+ downsamples.append(Resample(out_dim, mode=mode))
+ scale /= 2.0
+ self.downsamples = nn.Sequential(*downsamples)
+
+ # middle blocks
+ self.middle = nn.Sequential(
+ ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim),
+ ResidualBlock(out_dim, out_dim, dropout))
+
+ # output blocks
+ self.head = nn.Sequential(
+ RMS_norm(out_dim, images=False), nn.SiLU(),
+ CausalConv3d(out_dim, z_dim, 3, padding=1))
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = self.conv1(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = self.conv1(x)
+
+ # downsamples
+ for layer in self.downsamples:
+ if feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ # middle
+ for layer in self.middle:
+ if isinstance(layer, ResidualBlock) and feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ # head
+ for layer in self.head:
+ if isinstance(layer, CausalConv3d) and feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = layer(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = layer(x)
+ return x
+
+
+class Decoder3d(nn.Module):
+
+ def __init__(self,
+ dim=128,
+ z_dim=4,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_upsample=[False, True, True],
+ dropout=0.0):
+ super().__init__()
+ self.dim = dim
+ self.z_dim = z_dim
+ self.dim_mult = dim_mult
+ self.num_res_blocks = num_res_blocks
+ self.attn_scales = attn_scales
+ self.temperal_upsample = temperal_upsample
+
+ # dimensions
+ dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+ scale = 1.0 / 2**(len(dim_mult) - 2)
+
+ # init block
+ self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+ # middle blocks
+ self.middle = nn.Sequential(
+ ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]),
+ ResidualBlock(dims[0], dims[0], dropout))
+
+ # upsample blocks
+ upsamples = []
+ for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+ # residual (+attention) blocks
+ if i == 1 or i == 2 or i == 3:
+ in_dim = in_dim // 2
+ for _ in range(num_res_blocks + 1):
+ upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+ if scale in attn_scales:
+ upsamples.append(AttentionBlock(out_dim))
+ in_dim = out_dim
+
+ # upsample block
+ if i != len(dim_mult) - 1:
+ mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'
+ upsamples.append(Resample(out_dim, mode=mode))
+ scale *= 2.0
+ self.upsamples = nn.Sequential(*upsamples)
+
+ # output blocks
+ self.head = nn.Sequential(
+ RMS_norm(out_dim, images=False), nn.SiLU(),
+ CausalConv3d(out_dim, 3, 3, padding=1))
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ # conv1
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = self.conv1(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = self.conv1(x)
+
+ # middle
+ for layer in self.middle:
+ if isinstance(layer, ResidualBlock) and feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ # upsamples
+ for layer in self.upsamples:
+ if feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ # head
+ for layer in self.head:
+ if isinstance(layer, CausalConv3d) and feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = layer(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = layer(x)
+ return x
+
+
+def count_conv3d(model):
+ count = 0
+ for m in model.modules():
+ if isinstance(m, CausalConv3d):
+ count += 1
+ return count
+
+
+class WanVAE_(nn.Module):
+
+ def __init__(self,
+ dim=128,
+ z_dim=4,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_downsample=[True, True, False],
+ dropout=0.0):
+ super().__init__()
+ self.dim = dim
+ self.z_dim = z_dim
+ self.dim_mult = dim_mult
+ self.num_res_blocks = num_res_blocks
+ self.attn_scales = attn_scales
+ self.temperal_downsample = temperal_downsample
+ self.temperal_upsample = temperal_downsample[::-1]
+
+ # modules
+ self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
+ attn_scales, self.temperal_downsample, dropout)
+ self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+ self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+ self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,
+ attn_scales, self.temperal_upsample, dropout)
+ self.clear_cache()
+
+ def forward(self, x):
+ mu, log_var = self.encode(x)
+ z = self.reparameterize(mu, log_var)
+ x_recon = self.decode(z)
+ return x_recon, mu, log_var
+
+ def encode(self, x, scale):
+ self.clear_cache()
+ # cache
+ t = x.shape[2]
+ iter_ = 1 + (t - 1) // 4
+ # 对encode输入的x，按时间拆分为1、4、4、4....
+ for i in range(iter_):
+ self._enc_conv_idx = [0]
+ if i == 0:
+ out = self.encoder(
+ x[:, :, :1, :, :],
+ feat_cache=self._enc_feat_map,
+ feat_idx=self._enc_conv_idx)
+ else:
+ out_ = self.encoder(
+ x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+ feat_cache=self._enc_feat_map,
+ feat_idx=self._enc_conv_idx)
+ out = torch.cat([out, out_], 2)
+ mu, log_var = self.conv1(out).chunk(2, dim=1)
+ if isinstance(scale[0], torch.Tensor):
+ mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+ 1, self.z_dim, 1, 1, 1)
+ else:
+ mu = (mu - scale[0]) * scale[1]
+ self.clear_cache()
+ return mu
+
+ def decode(self, z, scale):
+ self.clear_cache()
+ # z: [b,c,t,h,w]
+ if isinstance(scale[0], torch.Tensor):
+ z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+ 1, self.z_dim, 1, 1, 1)
+ else:
+ z = z / scale[1] + scale[0]
+ iter_ = z.shape[2]
+ x = self.conv2(z)
+ for i in range(iter_):
+ self._conv_idx = [0]
+ if i == 0:
+ out = self.decoder(
+ x[:, :, i:i + 1, :, :],
+ feat_cache=self._feat_map,
+ feat_idx=self._conv_idx)
+ else:
+ out_ = self.decoder(
+ x[:, :, i:i + 1, :, :],
+ feat_cache=self._feat_map,
+ feat_idx=self._conv_idx)
+ out = torch.cat([out, out_], 2)
+ self.clear_cache()
+ return out
+
+ def cached_decode(self, z, scale):
+ # z: [b,c,t,h,w]
+ if isinstance(scale[0], torch.Tensor):
+ z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+ 1, self.z_dim, 1, 1, 1)
+ else:
+ z = z / scale[1] + scale[0]
+ iter_ = z.shape[2]
+ x = self.conv2(z)
+ for i in range(iter_):
+ self._conv_idx = [0]
+ if i == 0:
+ out = self.decoder(
+ x[:, :, i:i + 1, :, :],
+ feat_cache=self._feat_map,
+ feat_idx=self._conv_idx)
+ else:
+ out_ = self.decoder(
+ x[:, :, i:i + 1, :, :],
+ feat_cache=self._feat_map,
+ feat_idx=self._conv_idx)
+ out = torch.cat([out, out_], 2)
+ return out
+
+ def sample(self, imgs, deterministic=False):
+ mu, log_var = self.encode(imgs)
+ if deterministic:
+ return mu
+ std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+ return mu + std * torch.randn_like(std)
+
+ def clear_cache(self):
+ self._conv_num = count_conv3d(self.decoder)
+ self._conv_idx = [0]
+ self._feat_map = [None] * self._conv_num
+ # cache encode
+ self._enc_conv_num = count_conv3d(self.encoder)
+ self._enc_conv_idx = [0]
+ self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs):
+ """
+ Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.
+ """
+ # params
+ cfg = dict(
+ dim=96,
+ z_dim=z_dim,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_downsample=[False, True, True],
+ dropout=0.0)
+ cfg.update(**kwargs)
+
+ # init model
+ with torch.device('meta'):
+ model = WanVAE_(**cfg)
+
+ # load checkpoint
+ logging.info(f'loading {pretrained_path}')
+ model.load_state_dict(
+ torch.load(pretrained_path, map_location=device), assign=True)
+
+ return model
+
+
+class WanVAE:
+
+ def __init__(self,
+ z_dim=16,
+ vae_pth='cache/vae_step_411000.pth',
+ dtype=torch.float,
+ device="cuda"):
+ self.dtype = dtype
+ self.device = device
+
+ mean = [
+ -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
+ 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
+ ]
+ std = [
+ 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
+ 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
+ ]
+ self.mean = torch.tensor(mean, dtype=dtype, device=device)
+ self.std = torch.tensor(std, dtype=dtype, device=device)
+ self.scale = [self.mean, 1.0 / self.std]
+
+ # init model
+ self.model = _video_vae(
+ pretrained_path=vae_pth,
+ z_dim=z_dim,
+ ).eval().requires_grad_(False).to(device)
+
+ def encode(self, videos):
+ """
+ videos: A list of videos each with shape [C, T, H, W].
+ """
+ with amp.autocast(dtype=self.dtype):
+ return [
+ self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0)
+ for u in videos
+ ]
+
+ def decode(self, zs):
+ with amp.autocast(dtype=self.dtype):
+ return [
+ self.model.decode(u.unsqueeze(0),
+ self.scale).float().clamp_(-1, 1).squeeze(0)
+ for u in zs
+ ]
diff --git a/14B/configs/inference.yaml b/14B/configs/inference.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..323fc2acb6887bd25c820f151403221b0a5a655e
--- /dev/null
+++ b/14B/configs/inference.yaml
@@ -0,0 +1,32 @@
+# Human/VLM-evaluated step-80 generator: four trained stages + terminal t=0.
+# The auxiliary seed controls context-cache noise, independently of the video seed.
+attention_type: softmax
+causal: true
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true
+num_train_timestep: 1000
+timestep_shift: 5.0
+denoising_loss_type: flow
+mixed_precision: true
+independent_first_frame: false
+same_step_across_blocks: true
+context_noise: 0
+num_frame_per_block: 3
+latent_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+guidance_scale: 1.0
+negative_prompt: ''
+model_kwargs:
+ model_name: Wan2.1-T2V-14B
+ timestep_shift: 5.0
+inference_terminal_zero_cleanup: true
+inference_train_matched_context_cache: true
+inference_aux_noise_seed: 12345
diff --git a/14B/configs/train.yaml b/14B/configs/train.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..511c829500768931046ba3bc60b62ee2e0219f49
--- /dev/null
+++ b/14B/configs/train.yaml
@@ -0,0 +1,312 @@
+# Public names follow the paper; code/utils/config.py accepts both these and legacy keys.
+# Three-encoder recipe used by the human and VLM evaluations. See ../docs/evaluated-settings.md.
+
+# Step 80 uses the first 80 updates of the original 200-step cosine schedule.
+
+# Initialization and checkpoint compatibility. Only the generator is optimized.
+initial_generator_checkpoint: /path/to/video-generation-assets/checkpoints/causal-ode-init.safetensors
+initial_generator_checkpoint_sha256: 01a21da77adf0e1e803d74ae73ac9c2bc397b1a05a74cbf89709a653282f7ab6
+resume_checkpoint: null
+
+# Distributed execution: wrap units and precision; score-model settings below are legacy/inactive.
+generator_fsdp_wrap_strategy: transformer
+real_score_fsdp_wrap_strategy: size
+fake_score_fsdp_wrap_strategy: size
+text_encoder_fsdp_wrap_strategy: size
+uses_score_models: false
+uses_real_score: false
+uses_fake_score: false
+
+# Autoregressive rollout: four candidate denoising times, with stochastic training exit steps.
+causal: true
+denoising_step_list:
+- 1000
+- 750
+- 500
+- 250
+warp_denoising_step: true
+timestep_shift: 5.0
+guidance_scale: 3.0
+mixed_precision: true
+
+# Training seed 0 requests a randomly chosen seed; use an explicit positive value for repeatable input draws.
+seed: 0
+
+# Paper Sec. 3.1: match generated videos in frozen representation spaces, not score-model outputs.
+training_objective: distribution_matching
+training_loop: distribution_matching
+feature_space: frozen_video_encoders
+mmd_estimator: hybrid_nystrom_monte_carlo
+objective_weight: 1.0
+
+# Each encoder loss M uses M / (stop_gradient(M) + epsilon); encoder weights follow this scaling.
+normalize_encoder_losses: true
+encoder_loss_normalization_epsilon: 0.0001
+
+# Reference features are read from disk. "Teacher" in legacy keys means reference data, not a live network.
+reference_feature_source: matched_encoder_features
+reference_feature_manifest: /path/to/video-generation-assets/teacher_bank/manifest.json
+reference_feature_manifest_sha256: c755ccc39ddddca7c9551ce4657c3e95f202453a0943711f0f9cd90b47d106e4
+reference_feature_content_hash: 6b8f249733c39d865a8a9d4b1963779d5578c6a6d68aa5d96002a4b706c766e9
+
+# Reference video/prompt correspondence and immutable asset identities.
+reference_video_manifest: /path/to/video-generation-assets/source/matched.jsonl
+matched_manifest_sha256: a047d08b22757f979f40787af816fa76bf31c68742242250ffbb6177cb815a02
+reference_video_count: 8000
+matched_source_csv_sha256: 5dd5e93f25d7fec7adfee10cfdf4ba85851f40624c3274cceef068ec1f0c36e9
+matched_sample_identity_sha256: 43c752141b38dadb8791ca3e16507b2bc2c0deae3345c5ca609e03ae82a3a021
+matched_sampling_algorithm: sha256_seed_step_position_mod_n_v1
+matched_sampling_with_replacement: true
+
+# Paper Sec. 3.2: combine a persistent Nyström summary with the current reference minibatch.
+triple_token_sampling: uniform_fixed
+triple_nystrom_attraction_ema_enabled: false
+
+# Nyström:Monte Carlo = 2:1, corresponding to paper mixture alpha = 1/3.
+nystrom_reference_weight: 2.0
+triple_nystrom_current_teacher_attraction_enabled: true
+monte_carlo_reference_weight: 1.0
+triple_nystrom_repulsion_fifo_size: 0
+
+# The legacy triple_nystrom_* reference fields belong to V-JEPA 2; other encoders have separate packs.
+triple_nystrom_reference_path: /path/to/video-generation-assets/references/vjepa2.pt
+triple_nystrom_reference_sha256: caa965b372e8dd75ba3757adf9caefa7c9b0bbc758d39c9d566a0056b2c19e20
+triple_nystrom_landmarks: 1024
+triple_nystrom_kernel_chunk_size: 32
+triple_nystrom_seed: 0
+
+# All three frozen encoders contribute; normalize each branch, then combine with weights 10:1:1.
+active_feature_encoders:
+- dinov3
+- vjepa2
+- videomae_v2
+vjepa_loss_weight: 1.0
+videomae_loss_weight: 1.0
+encoder_loss_divisor: 3.0
+
+# DINOv3: 20 sampled frames, five-frame temporal means, 1024 patch tokens.
+dino_loss_weight: 10.0
+dinov3_repo_path: /path/to/video-generation-assets/models/dinov3-src
+dinov3_weight_path: /path/to/video-generation-assets/models/dinov3_vitl16_lvd1689m.pth
+dinov3_source_commit: 6876159a11b4df116f30f667f8c9888617df0751
+dinov3_checkpoint_sha256: b361ce4cccaae083fccaa7a732bd9abb143b5a3736175152f28f331b600f8eec
+dinov3_dense_geometry: temporal_mean5_full256
+dinov3_dense_source_num_frames: 81
+dinov3_dense_frame_indices:
+- 0
+- 4
+- 8
+- 13
+- 17
+- 21
+- 25
+- 29
+- 34
+- 38
+- 42
+- 46
+- 51
+- 55
+- 59
+- 63
+- 67
+- 72
+- 76
+- 80
+dinov3_dense_input_size: 256
+dinov3_dense_patch_size: 16
+dinov3_dense_spatial_axis_indices:
+- 0
+- 1
+- 2
+- 3
+- 4
+- 5
+- 6
+- 7
+- 8
+- 9
+- 10
+- 11
+- 12
+- 13
+- 14
+- 15
+dinov3_dense_temporal_groups: 4
+dinov3_dense_temporal_pool_size: 5
+dinov3_dense_forward_chunk_frames: 4
+dinov3_dense_gradient_checkpointing: false
+dinov3_nystrom_reference_path: /path/to/video-generation-assets/references/dinov3.pt
+dinov3_nystrom_reference_sha256: bc488a15ea499cc55eaaf201401eecfe37f9a92ccbe43d6e8d3e0fc37415f011
+dinov3_nystrom_landmarks: 256
+dinov3_nystrom_kernel_chunk_size: 32
+dinov3_nystrom_seed: 0
+dinov3_feature_metadata_preserve_runtime_paths: true
+
+# V-JEPA 2: last 64 of 81 frames, block index 15 (the 16th block), 1024 selected tokens.
+vjepa2_model_id: /path/to/video-generation-assets/models/vjepa2-vitl-fpc64-256
+vjepa2_reference_model_id: vjepa2-vitl-fpc64-256
+vjepa2_model_sha256: 25466aef85727d16546c6cf8c99f12fcfad9cbca8225d45f23685e2e025b786b
+vjepa2_dense_layers:
+- 15
+vjepa2_dense_layer: 15
+vjepa2_dense_num_frames: 64
+vjepa2_dense_source_num_frames: 81
+vjepa2_dense_frame_sampling: last_contiguous
+vjepa2_student_source_num_frames: 81
+vjepa2_student_frame_sampling: last_contiguous
+vjepa2_dense_num_tokens: 1024
+vjepa2_dense_resize_short_side: 292
+vjepa2_dense_crop_size: 256
+vjepa2_dense_gradient_checkpointing: true
+
+# VideoMAEv2: skip frame 0, split the remaining 80 frames into five 16-frame segments.
+videomae_v2_model_id: /path/to/video-generation-assets/models/VideoMAEv2-Large
+videomae_v2_model_revision: 9981a9c8f77118c421e5228e1b219468a4b0238d
+videomae_v2_download_manifest_content_hash: 5e362cfe6fdb1bec9e3e195e5ef43cdfebbdfaec2625b1f60550efd05ead0d7d
+videomae_v2_local_files_only: true
+videomae_v2_dense_layers:
+- 15
+videomae_v2_dense_layer: 15
+videomae_v2_encoder_count: 5
+videomae_v2_share_encoder_weights: true
+videomae_v2_skip_leading_frames: 1
+videomae_v2_segment_num_frames: 16
+videomae_v2_source_num_frames: 81
+
+# 256 tokens per segment, hence 1280 tokens for the complete video.
+videomae_v2_dense_num_tokens: 256
+videomae_v2_dense_resize_short_side: 224
+videomae_v2_dense_crop_size: 224
+videomae_v2_dense_gradient_checkpointing: true
+videomae_v2_nystrom_reference_path: /path/to/video-generation-assets/references/videomae_v2.pt
+videomae_v2_nystrom_reference_sha256: 17284c8eba334b091ab9668621ed9acccb045911e2acac51fd2cf3e4f4d56a36
+videomae_v2_nystrom_landmarks: 256
+videomae_v2_nystrom_kernel_chunk_size: 32
+videomae_v2_nystrom_seed: 0
+videomae_v2_nystrom_attraction_weight: 2.0
+videomae_v2_nystrom_current_teacher_attraction_weight: 1.0
+
+# Paper Sec. 3.3: separate distribution batch B, differentiation batch b, and reference minibatch B_ref.
+per_device_microbatch_size: 2
+fresh_video_logical_micro_batch_size: 16
+fresh_video_physical_replay_batch_size: 1
+distribution_samples_per_rank: 32
+
+# Replay the configured per-rank suffix; this release does not draw the paper pseudocode's uniform subset.
+differentiation_samples_per_rank: 16
+fresh_video_partial_gradient: true
+differentiate_full_distribution_loss: true
+
+# Apply B/b once on replay gradients, separately from FSDP averaging compensation.
+differentiation_gradient_scale: 2.0
+distribution_batch_size: 256
+expected_world_size: 8
+reference_minibatch_size: 128
+
+# Recorded estimator geometry: [B, b, B_ref]. Released self-kernel mean includes diagonal pairs and divides by B^2.
+partial_mmd_geometry:
+- 256
+- 128
+- 128
+partial_mmd_normalization_identity: full_student_population_denominators_v1
+partial_mmd_kss_denominator: 65536
+partial_mmd_fixed_attraction_denominator: 256
+partial_mmd_live_cross_denominator: 32768
+partial_mmd_teacher_self_denominator: 16384
+partial_mmd_gradient_compensation: 2.0
+fresh_video_replay_atol: 1.0e-05
+fresh_video_replay_rtol: 0.0001
+
+# Parameter EMA is disabled; inference uses the trained generator weights directly.
+ema_weight: 0.0
+sharding_strategy: full
+fsdp_backward_prefetch: post
+optimizer_foreach: false
+optimizer_fused: false
+optimizer_state_cpu_offload_during_backward: true
+
+# Generator optimizer. Learning-rate schedule parameters below remain unchanged.
+generator_learning_rate: 2.5e-06
+
+# 14B: retain the original 200-update cosine curve while stopping at update 80.
+lr_schedule: cosine
+lr_final: 5.0e-07
+lr_schedule_total_steps: 200
+beta1: 0.0
+beta2: 0.999
+weight_decay: 0.0
+max_grad_norm_generator: 10.0
+negative_prompt: ''
+dfake_gen_update_ratio: 1
+
+# Latent tensor [batch, time, channels, height, width]; 21 latent frames decode to 81 RGB frames at 480 x 832.
+latent_video_shape:
+- 1
+- 21
+- 16
+- 60
+- 104
+denoising_loss_type: flow
+
+# Activation recomputation saves GPU memory; it is unrelated to saving model checkpoints.
+gradient_checkpointing: true
+generator_replay_activation_cpu_offload: true
+generator_replay_activation_cpu_offload_block_stride: 0
+generator_replay_frozen_encoder_cpu_offload: true
+num_frame_per_block: 3
+load_raw_video: false
+model_kwargs:
+ model_name: Wan2.1-T2V-14B
+ timestep_shift: 5.0
+
+# Checkpoint/logging policy. Full states include optimizer and per-rank RNG state.
+log_iters: 10
+save_full_state: true
+save_full_state_final_only: false
+checkpoint_keep_last: 0
+checkpoint_model_min_free_gib: 12
+checkpoint_full_min_free_gib: 32
+allow_resume_world_size_change: false
+
+# Optional in-training evaluation is disabled; use scripts/infer.sh for generation.
+eval_interval: 0
+eval_seed: 12345
+eval_num_prompts: 0
+eval_prompt_path: prompts/test_8.txt
+num_workers: 0
+stdout_log_iters: 1
+
+# Optimizer updates, not epochs or diffusion timesteps. Do not shorten the configured cosine horizon.
+training_updates: 80
+num_train_timestep: 1000
+
+# Clean cached context; blocks share the sampled training exit step. Frozen encoders still pass input gradients.
+context_noise: 0
+independent_first_frame: false
+same_step_across_blocks: true
+last_step_only: false
+i2v: false
+num_training_frames: 21
+gc_interval: 25
+
+# Opaque feature/data identities are retained so existing reference packs and checkpoints still match.
+feature_protocol_identity: distribution_matching_14b_three_encoder_step80
+matched_video_root: /path/to/video-generation-assets/videos
+dataset_identity:
+ name: wan21_14b_reference_8000
+ sample_count: 8000
+ pixel_geometry:
+ - 81
+ - 480
+ - 832
+ latent_geometry:
+ - 21
+ - 60
+ - 104
+ caption_type: source_wan_dataset_prompt
+ source_manifest_sha256: a047d08b22757f979f40787af816fa76bf31c68742242250ffbb6177cb815a02
+ teacher_bank_content_hash: 6b8f249733c39d865a8a9d4b1963779d5578c6a6d68aa5d96002a4b706c766e9
+ reference_manifest_sha256: bc6041e517b6aa37d149c33cc9e5da69693e9efca515387e6bdb16a015307c4d
+ checkpoint_policy: every_10_keep_all
+fresh_video_progress_logs: true
diff --git a/14B/docs/assets.md b/14B/docs/assets.md
new file mode 100644
index 0000000000000000000000000000000000000000..1045dae6a7e06510938856c63d7104416a011fee
--- /dev/null
+++ b/14B/docs/assets.md
@@ -0,0 +1,73 @@
+# 14B Asset Preparation
+
+The evaluated model is the **three-encoder step-80 non-EMA generator**, trained on **8,000 reference videos**. Its asset identities are recorded in [`configs/train.yaml`](../configs/train.yaml). The former two-encoder / 8,852-video bank is a different experiment and cannot reproduce this model.
+
+Prepare matching model and reference assets locally; this anonymous release does not embed hosted resource locations. Use the matching three-encoder checkpoint and reference assets; filenames or the step number alone do not identify the experiment. This source repository does not include model weights or reference tensors. A final distilled model is not interchangeable with the **causal ODE initialization** needed for training.
+
+## Directory layout
+
+```text
+ASSET_ROOT/
+├── models/
+│ ├── Wan2.1-T2V-14B/ # Base generator, T5, tokenizer, and VAE
+│ ├── dinov3-src/ # Official DINOv3 source at the pinned commit
+│ ├── dinov3_vitl16_lvd1689m.pth
+│ ├── vjepa2-vitl-fpc64-256/
+│ └── VideoMAEv2-Large/ # Includes download_manifest.json
+├── checkpoints/
+│ └── causal-ode-init.safetensors
+├── teacher_bank/
+│ ├── manifest.json
+│ └── ... # Preserve all shard-relative paths
+├── source/
+│ ├── matched.jsonl
+│ └── matched.meta.json
+├── references/
+│ ├── dinov3.pt
+│ ├── vjepa2.pt
+│ └── videomae_v2.pt
+└── videos/ # Feature extraction only; not decoded during bank-based training
+```
+
+Preserve original serialized manifests, reference packs, and shards. The three reference files may be renamed as above without changing their bytes. `prepare_config.py` reads DINOv3 and VideoMAEv2 metadata identities from the bank and keeps them separate from local encoder load paths. These local settings are excluded from Git.
+
+## Frozen encoders
+
+Wan2.1-T2V-14B includes the shared T5 weights, VAE weights, and tokenizer. Place its files in the directory described in the model README. Before offline training, prepare the other encoders:
+
+```bash
+git clone https://github.com/facebookresearch/dinov3.git "$ASSET_ROOT/models/dinov3-src"
+git -C "$ASSET_ROOT/models/dinov3-src" checkout 6876159a11b4df116f30f667f8c9888617df0751
+# Obtain dinov3_vitl16_lvd1689m.pth through the official DINOv3 model access process.
+# Place the matching model files at the local paths described below.
+PYTHONPATH="$PWD/code" python code/tools/download_videomae_v2_large.py \
+ --source-dir /path/to/local/VideoMAEv2-Large \
+ --output-dir "$ASSET_ROOT/models/VideoMAEv2-Large"
+```
+
+The exact encoder hashes are in the training config; a newer checkpoint is not interchangeable with a frozen reference bank. VideoMAEv2 uses revision `9981a9c8f77118c421e5228e1b219468a4b0238d`.
+
+| Encoder | Input and representation | Nyström landmarks | RBF bandwidth |
+|---|---|---:|---:|
+| DINOv3 ViT-L/16 | 20 sampled frames resized to 256²; final normalized patch features; mean each group of five frames; 1,024 × 1,024 matrix | 256 | 341.2794189453125 |
+| V-JEPA 2 ViT-L | Frames 17–80; short side 292, center crop 256²; block index 15; 1,024 × 1,024 matrix | 1,024 | 3025.305908203125 |
+| VideoMAEv2-Large | Frames 1–80 split into five 16-frame clips; short side/crop 224; block index 15; 256 tokens per clip, 1,280 × 1,024 per video | 256 | 3204.123291015625 |
+
+VideoMAEv2 shares one frozen encoder across the five clips. All branches use ImageNet pixel normalization. The video-level kernel is `exp(-||X-Y||_F² / (2 sigma²))`, without dividing the distance by the number of tokens or channels. The fixed Nyström references use ridge `1e-6` and 20 k-means iterations. The loss normalization epsilon is separately `1e-4`. Full metadata is in [`reference-settings.json`](../provenance/reference-settings.json).
+
+## Rebuilding a reference bank
+
+Rebuilding requires the original captioned videos and matched manifest. Re-extraction may change serialized hashes; another video collection defines a new experiment. The included three-encoder tools are:
+
+```bash
+PYTHONPATH="$PWD/code" python code/tools/extract_triple_matched_teacher_bank.py --help
+PYTHONPATH="$PWD/code" python code/tools/build_triple_nystrom_references.py --help
+```
+
+Use `--expected-count 8000 --shard-rows 32` for extraction. Pass all three pinned encoder locations and identities, the source manifest and CSV hashes, and the video root. For reference construction use `--expected-sample-count 8000 --expected-shard-size 32 --kmeans-iterations 20 --ridge 1e-6`, together with the resulting bank identities. The builder fixes the landmark counts to DINOv3 256, V-JEPA 2 1,024, and VideoMAEv2 256.
+
+For a rebuilt bank, update a local copy of the training config with its new manifest/content hashes and all three reference filenames/hashes before preparing local paths. The optional `*_no_dinov3_*` tools are two-encoder ablation utilities and are not part of this recipe.
+
+## Local snapshot preparation
+
+All model loaders require existing local directories. Hosted model repository URLs, repository IDs and automatic downloads are omitted from this source release. The legacy-named `download_videomae_v2_large.py` stages an existing local snapshot using `--source-dir` and an empty `--output-dir`; it does not download files. If the source includes a valid pinned `download_manifest.json`, that manifest is preserved for compatibility with frozen reference banks. A snapshot without a manifest gets a new local manifest; rebuild matching reference assets and update their configuration identities before training. `prepare_config.py` reads model identities from an available local feature bank into private local configuration files.
diff --git a/14B/docs/evaluated-settings.md b/14B/docs/evaluated-settings.md
new file mode 100644
index 0000000000000000000000000000000000000000..a2985580b4371c4691f89d7a63c2abe007056eb4
--- /dev/null
+++ b/14B/docs/evaluated-settings.md
@@ -0,0 +1,41 @@
+# Evaluated 14B step-80 settings
+
+The released configuration corresponds to the **three-encoder step-80 non-EMA model used in the human and VLM evaluations**. The earlier two-encoder / 8,852-reference configuration described a different experiment.
+
+| Setting | Value |
+|---|---|
+| Generator | Wan2.1-T2V-14B, causal softmax attention |
+| Initialization | Causal ODE; SHA-256 `01a21da77adf0e1e803d74ae73ac9c2bc397b1a05a74cbf89709a653282f7ab6` |
+| Frozen encoders | DINOv3 + V-JEPA 2 + VideoMAEv2 |
+| Loss weights / divisor | **10 : 1 : 1 / 3**, after detached per-encoder normalization |
+| Reference videos | **8,000** |
+| Distribution / differentiation / reference minibatch | **256 / 128 / 128** |
+| GPUs / per-rank generated / per-rank differentiated | **8 / 32 / 16** |
+| Physical generation / logical feature / physical replay batch | **2 / 16 / 1** |
+| AdamW | LR `2.5e-6`, betas `(0, 0.999)`, epsilon `1e-8`, weight decay `0`, gradient clipping `10` |
+| Schedule | Cosine to `5e-7` over **200** updates; take checkpoint **80** |
+| Nyström / Monte Carlo mixture | **2 : 1** |
+| Training video / causal chunk | 81 RGB frames at 480 × 832; 21 latent frames; 3 latent frames per chunk |
+| Inference weights | `checkpoint_model_000080/model.pt`, regular `generator`, no EMA or CFG |
+
+The original training config SHA-256 is `fba9c9abd26f9b6b08ed5450c97d01a250da2747c7325db469a2e2f204f694ab`. Public asset paths and descriptive identities are relocated; the three reference-pack hashes and numerical training settings are retained. The public recipe stops at update 80 without shortening the original 200-update learning-rate curve. `seed: 0` chooses an effective seed at startup; exact historical training replay additionally requires its effective seed and RNG state.
+
+## Inference
+
+The pipeline and sample-batch noise order follow the previous step-80 inference run:
+
+- Four trained stages `[1000, 750, 500, 250]`, warped to approximately `[1000, 937.5, 833.3333, 625]`, then a terminal `t=0` cleanup.
+- Training-matched context-cache writes, with auxiliary noise seed `12345`. Scheduler `t=0` has a small nonzero noise level; cache input is noised separately from the saved clean output.
+- **Seed 0, five samples per prompt**, continuous random stream. Draw all five samples' initial/transition noises together, then execute one video at a time.
+- 81 output frames, 832 × 480, 16 FPS.
+
+```bash
+CHECKPOINT=/path/to/checkpoint_model_000080/model.pt \
+PROMPTS=/path/to/prompts.txt \
+SEED=0 NUM_SAMPLES=5 OUTPUT_DIR=outputs/step80 \
+bash scripts/infer.sh
+```
+
+Set `NUM_SAMPLES=1` for a smaller demo. This changes the random-number sequence. Exact historical video reproduction requires the same prompts, ordering and shard assignment in addition to the same checkpoint. The original VBench source used four contiguous shards with 237 / 237 / 236 / 236 prompt rows, restarting seed 0 for each shard. This package provides generation, not human/VLM scoring infrastructure.
+
+See [asset preparation](assets.md) for encoder geometry and reference bandwidths, and [validation](validation.md) for checks and limits.
diff --git a/14B/docs/validation.md b/14B/docs/validation.md
new file mode 100644
index 0000000000000000000000000000000000000000..3b3527fb3892b155952ea26ce435b899b939ccce
--- /dev/null
+++ b/14B/docs/validation.md
@@ -0,0 +1,27 @@
+# Validation
+
+## Evaluated checkpoint
+
+The restored recipe is the **three-encoder, 8,000-reference-video run**. The original checkpoint header records step **80**, **1,095 generator parameter tensors**, and all three branches: DINOv3, V-JEPA 2, and VideoMAEv2. Its bank identity matches the released training config. The historical human-evaluation generation records and the VLM source-video run point to this same checkpoint.
+
+The original step-80 log reports loss **3.854**, generator gradient norm **10.98**, three encoder VJPs, batch geometry **256 / 128 / 128**, zero reported feature replay error for all three encoders, and a completed optimizer update. These are historical measurements, not a new run of this package. See [`original-step80.json`](../results/original-step80.json).
+
+## Checks of the corrected release
+
+- **51 CPU tests passed; 1 CUDA-only test skipped.** Coverage includes three-encoder configuration, bank loading, MMD gradients, replay/offload, the 200-update cosine curve, DINOv3 metadata relocation, terminal cleanup and context-cache inference, and five-sample noise ordering.
+- All **184 fields** from the original training configuration were accounted for. Asset paths and descriptive identities are public equivalents; the public stopping point is 80 rather than 200, while the cosine horizon remains 200. The three active encoders are specified explicitly.
+- The public sampler was compared with the actual historical evaluation sampler using a small recording generator and the real scheduler. Outputs and every generator input/timestep matched exactly for two consecutive prompts, both with ordinary random draws and with supplied noise tensors.
+- The five-sample adapter draws the same 21 main transition tensors and 14 auxiliary tensors per prompt as the original 81-frame inference adapter, before single-video execution. CPU tests verify noise slices and final RNG state.
+- DINOv3 metadata relocation reproduces the complete actual frozen-bank metadata and digest without loading models from the original paths.
+- The single-encoder MMD kernel, training rollout, and encoder VJP implementation match the original three-encoder snapshot's executable AST after removing documentation and restoring readability aliases.
+- Python parsing, shell syntax, and diff whitespace checks passed.
+
+Path-free evidence is in [`release-validation.json`](../results/release-validation.json). Original source-fingerprint manifests are omitted from this anonymous release.
+
+**Full 14B GPU training, GPU video generation, and VBench were not rerun for this correction.** CPU sampler comparisons use a small test generator, not the 14B network, and do not establish identical generated videos or a new benchmark score.
+
+Run the CPU suite from `14B/`:
+
+```bash
+python scripts/test_cpu.py
+```
diff --git a/14B/provenance/reference-settings.json b/14B/provenance/reference-settings.json
new file mode 100644
index 0000000000000000000000000000000000000000..132695668ad31e5ae96667fcc4eea8164ceaae94
--- /dev/null
+++ b/14B/provenance/reference-settings.json
@@ -0,0 +1,98 @@
+{
+ "variant": "14B three encoders, checkpoint 80",
+ "reference_count": 8000,
+ "weights": [
+ 10.0,
+ 1.0,
+ 1.0
+ ],
+ "divisor": 3.0,
+ "batch_geometry": [
+ 256,
+ 128,
+ 128
+ ],
+ "branches": {
+ "DINOv3": {
+ "sigma": 341.2794189453125,
+ "landmark_shape": [
+ 256,
+ 1024,
+ 1024
+ ],
+ "nystrom": {
+ "bandwidth_source": "median_frobenius_distance",
+ "generated_gram": "biased_exact_including_diagonal",
+ "kernel": "gaussian_rbf",
+ "kernel_metric": "frobenius",
+ "kmeans": "deterministic_farthest_first_lloyd",
+ "kmeans_iterations": 20,
+ "median_max_samples": 2048,
+ "num_landmarks": 256,
+ "reference_count": 8000,
+ "ridge": 1e-06,
+ "sample_axis": "video",
+ "sample_shape": [
+ 1024,
+ 1024
+ ],
+ "seed": 0
+ },
+ "content_hash": "67851f7828ee0b9e12643242ff93f2b2b9b7c56059e51b9b2f267972da15eba9"
+ },
+ "V-JEPA 2": {
+ "sigma": 3025.305908203125,
+ "landmark_shape": [
+ 1024,
+ 1024,
+ 1024
+ ],
+ "nystrom": {
+ "bandwidth_source": "median_frobenius_distance",
+ "generated_gram": "biased_exact_including_diagonal",
+ "kernel": "gaussian_rbf",
+ "kernel_metric": "frobenius",
+ "kmeans": "deterministic_farthest_first_lloyd",
+ "kmeans_iterations": 20,
+ "median_max_samples": 2048,
+ "num_landmarks": 1024,
+ "reference_count": 8000,
+ "ridge": 1e-06,
+ "sample_axis": "video",
+ "sample_shape": [
+ 1024,
+ 1024
+ ],
+ "seed": 0
+ },
+ "content_hash": "d588e77dd6ea5d5611f5aeaac2c7f9a8c4e4fc36940209df32118a91e22d168b"
+ },
+ "VideoMAEv2": {
+ "sigma": 3204.123291015625,
+ "landmark_shape": [
+ 256,
+ 1280,
+ 1024
+ ],
+ "nystrom": {
+ "bandwidth_source": "median_frobenius_distance",
+ "generated_gram": "biased_exact_including_diagonal",
+ "kernel": "gaussian_rbf",
+ "kernel_metric": "frobenius",
+ "kmeans": "deterministic_farthest_first_lloyd",
+ "kmeans_iterations": 20,
+ "median_max_samples": 2048,
+ "num_landmarks": 256,
+ "reference_count": 8000,
+ "ridge": 1e-06,
+ "sample_axis": "video",
+ "sample_shape": [
+ 1280,
+ 1024
+ ],
+ "seed": 0
+ },
+ "content_hash": "88d55200ebb71d20ea75c70358eb8e5e68511c030546c5f832b53e591dd5a9b2"
+ }
+ }
+}
diff --git a/14B/requirements.txt b/14B/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..0dd2c6ae5d0deb7722dcd03fd1ac0f31a380629f
--- /dev/null
+++ b/14B/requirements.txt
@@ -0,0 +1,29 @@
+# Training and inference dependencies. Install torch/torchvision first.
+numpy==1.26.4
+diffusers==0.31.0
+transformers==4.53.3
+tokenizers==0.21.4
+timm==1.0.28
+accelerate==1.14.0
+tqdm==4.70.0
+imageio==2.37.4
+imageio-ffmpeg==0.6.0
+easydict==1.13
+ftfy==6.3.1
+omegaconf>=2.3.0,<2.4
+einops==0.8.1
+av==13.1.0
+opencv-python
+open-clip-torch
+lmdb==2.3.0
+matplotlib
+sentencepiece==0.2.2
+scikit-image
+huggingface-hub==0.36.2
+safetensors==0.8.0
+PyYAML==6.0.3
+pandas
+scipy
+Pillow==12.2.0
+requests==2.34.2
+packaging==26.3
diff --git a/14B/results/original-step80.json b/14B/results/original-step80.json
new file mode 100644
index 0000000000000000000000000000000000000000..3dd1126dca25183f920d541110379a9f0a074642
--- /dev/null
+++ b/14B/results/original-step80.json
@@ -0,0 +1,24 @@
+{
+ "scope": "Historical three-encoder training log; not a rerun of this package. Values retain the precision printed by the original logger.",
+ "checkpoint_step": 80,
+ "active_encoders": [
+ "dinov3",
+ "vjepa2",
+ "videomae_v2"
+ ],
+ "reference_count": 8000,
+ "use_ema": false,
+ "metrics": {
+ "generator_loss": 3.854,
+ "generator_grad_norm": 10.98,
+ "learning_rate": 1.818e-06,
+ "fresh_video_optimizer_step_skipped": 0.0,
+ "triple_dense_global_video_samples": 256.0,
+ "triple_dense_gradient_global_video_samples": 128.0,
+ "triple_dense_current_teacher_global_samples": 128.0,
+ "triple_dense_encoder_vjp_count": 3.0,
+ "dinov3_dense_replay_max_abs": 0.0,
+ "vjepa2_dense_replay_max_abs": 0.0,
+ "videomae_v2_dense_replay_max_abs": 0.0
+ }
+}
diff --git a/14B/results/release-validation.json b/14B/results/release-validation.json
new file mode 100644
index 0000000000000000000000000000000000000000..cefb4748042872717bdf29575c45523fe1bcdf3a
--- /dev/null
+++ b/14B/results/release-validation.json
@@ -0,0 +1,24 @@
+{
+ "checkpoint_step": 80,
+ "generator_parameter_tensors": 1095,
+ "active_encoders": [
+ "dinov3",
+ "vjepa2",
+ "videomae_v2"
+ ],
+ "reference_count": 8000,
+ "teacher_bank_content_hash": "6b8f249733c39d865a8a9d4b1963779d5578c6a6d68aa5d96002a4b706c766e9",
+ "human_motion_and_camera_protocols_same_checkpoint": true,
+ "vlm_source_vbench_same_checkpoint": true,
+ "cpu_sampler_matches_evaluated_source": {
+ "continuous_two_prompts": true,
+ "paired_two_prompts": true,
+ "scope": "Small recording generator and real scheduler; no 14B model weights were executed."
+ },
+ "actual_dinov3_metadata_relocation_exact": true,
+ "cpu_tests": {
+ "passed": 51,
+ "skipped_cuda": 1
+ },
+ "five_sample_noise_order_cpu_test": true
+}
diff --git a/14B/scripts/common.sh b/14B/scripts/common.sh
new file mode 100644
index 0000000000000000000000000000000000000000..6daf65f0cf949a83e28777ec9309042fcd494060
--- /dev/null
+++ b/14B/scripts/common.sh
@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+set -euo pipefail
+REPRO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+source "$REPRO_ROOT/configs/local.env"
+PYTHON_BIN="$(command -v "${PYTHON_BIN:-python}")"
+CONFIG="$REPRO_ROOT/configs/local.yaml"
+export SELF_FORCING_ASSET_ROOT="$VIDEO_GENERATION_ASSET_ROOT"
+export SELF_FORCING_MODELS_DIR="$VIDEO_GENERATION_ASSET_ROOT/models"
+export SELF_FORCING_CKPT_DIR="$VIDEO_GENERATION_ASSET_ROOT/checkpoints"
+export PYTHONPATH="$REPRO_ROOT/code${PYTHONPATH:+:$PYTHONPATH}"
+export PYTHONUNBUFFERED=1 PYTHONDONTWRITEBYTECODE=1
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-12}" TOKENIZERS_PARALLELISM=false
+export HF_HOME="${HF_HOME:-$REPRO_ROOT/cache/huggingface}"
+export TORCH_HOME="${TORCH_HOME:-$REPRO_ROOT/cache/torch}"
+export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-$REPRO_ROOT/cache/triton}"
+export HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+export TORCH_NCCL_ASYNC_ERROR_HANDLING=1 NCCL_DEBUG=WARN
+unset WORLD_SIZE RANK LOCAL_RANK MASTER_ADDR MASTER_PORT
diff --git a/14B/scripts/infer.sh b/14B/scripts/infer.sh
new file mode 100644
index 0000000000000000000000000000000000000000..c93f62ba6866d3aa8183c651848658e2d23880e4
--- /dev/null
+++ b/14B/scripts/infer.sh
@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+set -euo pipefail
+source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
+: "${CHECKPOINT:?Set CHECKPOINT to the step-80 model.pt file}"
+CHECKPOINT="$(realpath "$CHECKPOINT")"
+PROMPTS="$(realpath "${PROMPTS:-$REPRO_ROOT/code/prompts/test_8.txt}")"
+OUTPUT_DIR="$(realpath -m "${OUTPUT_DIR:-$REPRO_ROOT/outputs/demo}")"
+export CUDA_VISIBLE_DEVICES="${GPUS:-0}"
+cd "$REPRO_ROOT/code"
+"$PYTHON_BIN" inference.py --config_path "$REPRO_ROOT/configs/inference.yaml" \
+ --checkpoint_path "$CHECKPOINT" --data_path "$PROMPTS" \
+ --output_folder "$OUTPUT_DIR" --num_output_frames 21 \
+ --num_samples "${NUM_SAMPLES:-5}" --seed "${SEED:-0}" --save_with_index
diff --git a/14B/scripts/prepare_config.py b/14B/scripts/prepare_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..b9a8dfc951096593c3563b26830e4f456bb488f7
--- /dev/null
+++ b/14B/scripts/prepare_config.py
@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""Create local asset paths without changing the published training settings."""
+import argparse
+import hashlib
+import json
+from pathlib import Path
+import shlex
+import yaml
+import sys
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "code"))
+from utils.config import resolve_config_aliases
+PREFIX = "/path/to/video-generation-assets"
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--asset-root", required=True, type=Path)
+ parser.add_argument("--inference-only", action="store_true")
+ parser.add_argument("--resume-checkpoint", type=Path)
+ args = parser.parse_args()
+ asset_root = str(args.asset_root.expanduser().resolve())
+ cfg = resolve_config_aliases(yaml.safe_load((ROOT / "configs/train.yaml").read_text()))
+ cfg = {k: asset_root + v[len(PREFIX):]
+ if isinstance(v, str) and v.startswith(PREFIX + "/") else v
+ for k, v in cfg.items()}
+ if not args.inference_only:
+ bank_path = Path(cfg["triple_teacher_bank_manifest"])
+ data = bank_path.read_bytes()
+ if hashlib.sha256(data).hexdigest() != cfg["triple_teacher_bank_manifest_sha256"]:
+ raise ValueError("Feature bank manifest does not match the published configuration")
+ bank = json.loads(data)
+ # Frozen metadata identities are separate from local model load paths.
+ dino = bank["branches"]["dinov3"]["feature_contract"]
+ cfg["dinov3_reference_paths"] = {k: dino[k] for k in ("repo_path", "weight_path")}
+ for branch in ("vjepa2", "videomae_v2"):
+ record = bank["branches"][branch]
+ identity = record.get("feature_contract", record.get("model_identity", {}))
+ if identity.get("model_id"):
+ cfg[f"{branch}_reference_model_id"] = identity["model_id"]
+ if args.resume_checkpoint:
+ cfg["resume_checkpoint"] = str(args.resume_checkpoint.expanduser().resolve())
+ (ROOT / "configs/local.yaml").write_text(yaml.safe_dump(cfg, sort_keys=False))
+ (ROOT / "configs/local.env").write_text(
+ "export VIDEO_GENERATION_ASSET_ROOT=" + shlex.quote(asset_root) + "\n")
+ print("Local configuration written. Asset paths are excluded from Git.")
+
+if __name__ == "__main__":
+ main()
diff --git a/14B/scripts/test_cpu.py b/14B/scripts/test_cpu.py
new file mode 100644
index 0000000000000000000000000000000000000000..48890b9fee5b859f0e6d77d9cc9fb711d4a2eec1
--- /dev/null
+++ b/14B/scripts/test_cpu.py
@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+"""Run existing CPU tensor tests without initializing the historical T5 CUDA default."""
+import os
+from pathlib import Path
+import sys
+from unittest.mock import patch
+
+ROOT = Path(__file__).resolve().parents[1]
+os.environ['CUDA_VISIBLE_DEVICES'] = ''
+os.environ.setdefault('OMP_NUM_THREADS', '2')
+os.environ['PYTHONDONTWRITEBYTECODE'] = '1'
+sys.dont_write_bytecode = True
+sys.path.insert(0, str(ROOT/'code'))
+os.chdir(ROOT/'code')
+import torch
+import pytest
+
+# Only bypass import-time device discovery. No model weights or CUDA tensors load.
+# Original sources remain untouched; the patch is removed before executing tests.
+with patch('torch.cuda.current_device', return_value=0):
+ from trainer.distillation import Trainer
+ from model.wan_kv_mmd import WanKVMMD
+
+raise SystemExit(pytest.main([
+ '-q', '-p', 'no:cacheprovider',
+ 'tests/test_triple_partial_geometry.py',
+ 'tests/test_readable_config.py',
+ 'tests/test_evaluated_inference.py',
+ 'tests/test_active_encoder_offload.py',
+ 'tests/test_no_dinov3_filtered_bank.py',
+ 'tests/test_no_dinov3_ablation.py',
+ 'tests/test_release_schedule.py',
+ 'tests/test_triple_video_feature_bank.py',
+ 'tests/test_triple_feature_metadata_contract.py',
+ 'tests/test_triple_pixel_cotangent_accumulation.py',
+ 'tests/test_video_nystrom_partial_mmd.py',
+ 'tests/test_multi_representation_video_nystrom_mmd.py',
+] + sys.argv[1:]))
diff --git a/14B/scripts/train.sh b/14B/scripts/train.sh
new file mode 100644
index 0000000000000000000000000000000000000000..15bff30a7078b634c569ed3ab988691cee01ccdd
--- /dev/null
+++ b/14B/scripts/train.sh
@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+set -euo pipefail
+source "$(dirname "${BASH_SOURCE[0]}")/common.sh"
+extra_args=()
+case "${1:-}" in
+ "") ;;
+ --smoke) extra_args=(--max_steps 1) ;;
+ *) echo 'Usage: scripts/train.sh [--smoke]' >&2; exit 2 ;;
+esac
+[[ $# -le 1 ]] || { echo 'Too many arguments' >&2; exit 2; }
+export CUDA_VISIBLE_DEVICES="${GPUS:-${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}}"
+LOGDIR="$(realpath -m "${LOGDIR:-$REPRO_ROOT/runs/three-encoder}")"
+mkdir -p "$LOGDIR"
+cp "$CONFIG" "$LOGDIR/config.yaml"
+cd "$REPRO_ROOT/code"
+"$PYTHON_BIN" -m torch.distributed.run --standalone --nnodes=1 --nproc-per-node=8 train.py \
+ --config_path "$LOGDIR/config.yaml" --logdir "$LOGDIR" \
+ --disable-wandb --no_visualize "${extra_args[@]}" 2>&1 | tee -a "$LOGDIR/train.log"
diff --git a/ANONYMIZATION.txt b/ANONYMIZATION.txt
new file mode 100644
index 0000000000000000000000000000000000000000..39b5b77dcc69abcf5ac2a1624b31293014a61681
--- /dev/null
+++ b/ANONYMIZATION.txt
@@ -0,0 +1,35 @@
+Anonymous source release
+
+Removed author names, affiliations, contact details, personal repository and
+project links, tracking badges, the identifying paper PDF, shared-chat links,
+cluster-specific paths, private collaboration acknowledgments, original source
+fingerprint manifests, and Git metadata/history. Project-specific module names,
+class names and environment variables now use neutral names.
+
+Hosted model repository URLs and hard-coded repository IDs were removed.
+Model and tokenizer loaders require existing local directories. The legacy-named
+VideoMAEv2 preparation utility stages an existing local snapshot with --source-dir;
+it does not download models. Existing valid local manifests can be preserved for
+frozen-bank compatibility. New manifests require matching reference assets and
+configuration identities. Local preparation reads bank identities into private
+generated configuration files; do not include those files or model assets when
+sharing this source package.
+
+Third-party source attribution and license notices are retained. Scientific
+training settings and reported results are retained. This package contains source
+code, configurations, example prompts and documentation, without model weights
+or training data.
+
+Verification passed:
+- Python syntax: 180 files; shell syntax: 6 scripts.
+- Both training recipes passed the existing source/configuration checks.
+- Both local snapshot tools passed staging, occupied-output rejection, missing
+ source rejection, nested-directory rejection and manifest-preservation checks.
+- Model adapters passed local-only loading and missing-path checks with mocks.
+- Both preparation scripts passed path relocation and bank-identity checks.
+- Final content and archive checks found no targeted author/project identifiers,
+ personal links, hosted model repository addresses, Git metadata or paper PDF.
+
+No GPU training or model inference was run for this source anonymization task.
+Previously public source code and retained scientific results can still be matched
+to public material; removing identifying metadata cannot guarantee anonymity.
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,201 @@
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
diff --git a/README.md b/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..457eb0a7a891047cac3d7665d1c1909a905f8de3
--- /dev/null
+++ b/README.md
@@ -0,0 +1,74 @@
+---
+license: apache-2.0
+tags:
+- source-code
+- video-generation
+- pytorch
+---
+
+
+# Anonymous Video Generation
+
+Anonymous source-code repository.
+
+[Model checkpoints](https://huggingface.co/Elastic-forcing/anonymous-video-generation)
+
+[1.3B Training & Inference](1.3B/README.md) · [14B Training & Inference](14B/README.md)
+
+
+
+---
+
+
+Anonymous Video Generation post-trains few-step autoregressive video generators by matching generated and reference videos in frozen representation spaces. It combines a hybrid Nyström–Monte Carlo distribution estimator with gradient replay. Only the generator is optimized; diffusion score models are not needed during post-training.
+Both released recipes use the same three-encoder MMD objective; no auxiliary frame-level distribution objective is enabled or included.
+
+Read the implementation through the [paper-to-code guide](docs/paper-to-code.md), which maps the paper's notation to configuration fields, the training loop, the hybrid MMD objective, and gradient replay.
+
+## Models
+
+Each model directory contains its own implementation, configuration, dependency list, asset guide, and training/inference scripts. Run its commands from that directory.
+
+| Release | Frozen feature encoders | Training | Inference | Guide |
+|---|---|---|---|---|
+| **1.3B** | DINOv3 + V-JEPA 2 + VideoMAEv2 | 4 GPUs, 150 updates | 4 causal steps per chunk | [Training and inference](1.3B/README.md) |
+| **14B** | DINOv3 + V-JEPA 2 + VideoMAEv2 | 8 GPUs, 80 updates | 4 trained steps + terminal zero-step cleanup | [Training and inference](14B/README.md) |
+
+Both configurations use a distribution batch of 256 generated videos, a differentiation batch of 128 videos, and a reference minibatch of 128 videos. They generate 81-frame, 832 × 480 videos at 16 FPS.
+
+The 1.3B three-encoder configuration achieves **84.25 VBench Total** in the paper. The 14B release targets the **step-80 checkpoint**, retaining the first 80 updates of the original **200-step cosine learning-rate schedule**. This is the **three-encoder, 8,000-reference-video model used in the human and VLM evaluations**; see [evaluated settings](14B/docs/evaluated-settings.md).
+
+```text
+anonymous-video-generation-code/
+├── 1.3B/ # Three-encoder training and inference
+│ ├── code/
+│ ├── configs/
+│ ├── scripts/
+│ └── README.md
+├── 14B/ # Three-encoder training and step-80 inference
+│ ├── code/
+│ ├── configs/
+│ ├── scripts/
+│ └── README.md
+└── docs/ # Implementation and validation guides
+```
+
+## Getting Started
+
+```bash
+git clone https://huggingface.co/Elastic-forcing/anonymous-video-generation-code
+cd anonymous-video-generation-code
+
+# Choose one model, then follow its README.
+cd 14B
+# Or: cd 1.3B
+```
+
+See the [1.3B asset guide](1.3B/docs/assets.md) or [14B asset guide](14B/docs/assets.md) for model and reference-data requirements. This release contains source code; trained generator weights, the 14B ODE initialization, and frozen reference assets are stored separately and are not bundled here.
+
+Validation is documented separately for [1.3B](1.3B/docs/validation.md) and [14B](14B/docs/validation.md).
+The cleaned source-layout and CPU objective checks are described in [source validation](docs/readability-changes.md).
+
+## Acknowledgments
+
+This implementation builds on [Self-Forcing](https://github.com/guandeh17/Self-Forcing) and [Wan2.1](https://github.com/Wan-Video/Wan2.1). We also thank the authors of [DINOv3](https://github.com/facebookresearch/dinov3), [V-JEPA 2](https://github.com/facebookresearch/vjepa2), [VideoMAEv2](https://github.com/OpenGVLab/VideoMAEv2), and [VBench](https://github.com/Vchitect/VBench) for making their work available. Bundled source retains its upstream license notices; model weights and datasets follow their respective upstream terms.
diff --git a/docs/first-step-comparison.md b/docs/first-step-comparison.md
new file mode 100644
index 0000000000000000000000000000000000000000..bcc275f2da0bf3ae07a366fecb6cd063e6b6da27
--- /dev/null
+++ b/docs/first-step-comparison.md
@@ -0,0 +1,76 @@
+> Historical validation record: this comparison predates removal of the disabled auxiliary loss branches. Its source identifiers and field counts describe that earlier snapshot. Use scripts/check_readability.py and scripts/check_loss_parity.py for the cleaned release.
+
+# Readability refactor: first optimizer-step comparison
+
+The refactored 1.3B three-encoder code completed one full optimizer update on four H200 GPUs. Initial state, sampled inputs, forward features, loss and feature gradients match the earlier public-source run bitwise. The first recorded numerical difference is in `encoder_and_vae_backward`. Compare against the same-source repeated-run control below.
+
+## Source and run identity
+
+- Refactor base: `7985874ff49de326b8350a66516c50645a604f45`; the GPU comparison used a source snapshot captured before committing the refactor.
+- Tested source snapshot manifest SHA-256: `a806daa28aa2aa75c5adbe1e6bb930dbae822aee07cbfaf830980f6f8111b001`.
+- Previous public trace commit: `985ebec542b1c2bdb7af95a4a7404b3069cc15e7`. Its code tree is byte-identical to the pre-refactor 1.3B code tree: `60b92b4230405641d8d65a048d7b7cf842a653cf`.
+- Original-source control commit: `e8ec87e10c4825dfaaec993643724c5e8a2e28e9`.
+- Base seed **1862643**, with rank seeds **1862643–1862646**, matching the recovered historical effective seed. A legacy setting of `seed: 0` requests seed randomization and is not the effective seed.
+- Distribution / differentiation / reference batches: **256 / 128 / 128**. Four ranks, 64 forward samples and 32 backward samples per rank, microbatch size 1.
+- Encoder weights **10:1:1**, loss divisor **3**, fixed-reference / Monte Carlo reference weights **2:1**.
+- Same physical model/reference assets and four physical GPUs as the earlier runs. Independently queried software versions match.
+- All **173** readable input configuration fields resolve to the previous values in the same order. The new config loader and `--training-updates 1` alias were exercised. The configured cosine schedule remains steps 0–125, initial learning-rate setting 2.5e-6, floor 5e-7.
+- Observation hooks use the renamed methods and otherwise preserve the earlier observer. No production computation was changed for this run.
+
+## Scalar observations
+
+| Quantity | Previous public source | Readable source | Original A | Original B |
+|---|---:|---:|---:|---:|
+| Loss | 3.994718074798584 | 3.994718074798584 | 3.994718074798584 | 3.994718074798584 |
+| Gradient norm before clipping | 48.80884552001953 | 48.82303619384766 | 48.82280349731445 | 48.81880187988281 |
+
+## Full-vector differences
+
+| Quantity | Public → readable | Original A → readable | Original A → original B |
+|---|---:|---:|---:|
+| Latent gradient — relative L2 | 0.507891% | 0.508414% | 0.508167% |
+| Latent gradient — cosine similarity | 0.999987102368 | 0.99998707571 | 0.999987088324 |
+| Latent gradient — maximum absolute difference | 4.95910644531e-05 | 6.103515625e-05 | 6.103515625e-05 |
+| Generator gradient before clipping — relative L2 | 0.299963% | 0.301090% | 0.312999% |
+| Generator gradient before clipping — cosine similarity | 0.999995544683 | 0.999995467272 | 0.999995104525 |
+| Generator gradient before clipping — maximum absolute difference | 0.0146789550781 | 0.00958251953125 | 0.0114288330078 |
+| Generator gradient after clipping — relative L2 | 0.298507% | 0.301089% | 0.312905% |
+| Generator gradient after clipping — cosine similarity | 0.999995544684 | 0.999995467273 | 0.999995104526 |
+| Generator gradient after clipping — maximum absolute difference | 0.00274252891541 | 0.0019583106041 | 0.00238353013992 |
+| Actual Adam parameter update — relative L2 | 7.580502% | 7.667867% | 7.707839% |
+| Actual Adam parameter update — cosine similarity | 0.997126798637 | 0.997060191614 | 0.99702946052 |
+| Actual Adam parameter update — maximum absolute difference | 5.00679016113e-06 | 5.00679016113e-06 | 5.00679016113e-06 |
+
+Relative L2 means `||right - left||₂ / ||left||₂`, accumulated across all four ranks. Adam update differences compare `weights_after - initial_weights`; they are not normalized by the much larger model weight norm. Initial parameters are byte-identical across all variants.
+
+## Exact comparison with the previous public source
+
+| Recorded stage | Bitwise equality |
+|---|---|
+| effective_configuration | Exact |
+| initial_weights | Exact |
+| initial_rng | Exact |
+| forward_inputs_and_features | Exact |
+| mmd_result | Exact |
+| encoder_and_vae_backward | Different (645/2304 tensors) |
+| gradients_before_clip | Different (918/918 tensors) |
+| gradients_after_clip | Different (918/918 tensors) |
+| optimizer_state | Different (1836/2754 tensors) |
+| updated_weights | Different (916/3300 tensors) |
+| post_step_rng | Exact |
+
+All **424 tensor archives** (106 per rank) were compared, including all 256 forward samples and 128 selected backward samples. Equality checks compare raw bytes, including signed zeros. Configuration comparison excludes only output-directory and input-file-location fields. Timings and the final process-completion event are excluded from numerical comparison. All non-timing tensor values and structural fields remain compared.
+
+## Interpretation and scope
+
+The first recorded numerical difference is in `encoder_and_vae_backward`. The earlier same-source repeat has nonidentical encoder/VAE backward, generator gradients and Adam updates. The control demonstrates backward execution variability in this GPU configuration; it does not identify a specific nondeterministic kernel or establish a statistical tolerance from one control pair.
+
+The public-to-readable gradient relative L2 difference is 0.9583 times the measured same-source control difference; for the Adam update it is 0.9835 times. These ratios contextualize the observed discrepancies; they do not define a pass/fail tolerance or guarantee matching long-run trajectories.
+
+The readable-source run completed backward, clipping, Adam, native checkpoint saving, and all four ranks successfully (exit code 0). The previous public run also completed successfully. Original A completed the first update and every compared trace, then failed at checkpoint saving because its audit output directory was absent; original B completed normally.
+
+This comparison concerns reruns under the historical settings, not tensors captured during the original historical first step. It covers 1.3B only and does not evaluate 14B training, all training updates, generated-video quality, or VBench.
+
+## Software environment
+
+`{"cuda": "12.1", "cudnn": 90100, "flash_attn": "2.7.4.post1", "numpy": "1.26.4", "python": "3.10.12", "torch": "2.5.1+cu121", "torchvision": "0.20.1+cu121", "transformers": "4.53.3"}`
diff --git a/docs/paper-to-code.md b/docs/paper-to-code.md
new file mode 100644
index 0000000000000000000000000000000000000000..323c3be95be819ae021220db1b6dec0ffe2f43bd
--- /dev/null
+++ b/docs/paper-to-code.md
@@ -0,0 +1,158 @@
+# Reading the Anonymous Video Generation implementation
+
+Start with the training update, then follow the loss and replay calls. The paper's
+Sections 3.1–3.3 describe the objective and the separation between estimating a
+distribution and differentiating selected rollouts. Appendix C.4 and Algorithm 1
+describe the replay procedure. The notes below distinguish that mathematical
+presentation from the numerical choices retained in this release.
+
+Both releases contain one active distribution objective: the three-encoder MMD described below. The historical code/video_fd directory name is retained because its encoder adapters and replay VJPs are part of that active pipeline; it does not select an additional FD loss.
+
+## Reading order
+
+The two model directories have the same main entry points and both use DINOv3,
+V-JEPA 2 and VideoMAEv2. Start with 1.3B when reading the shared flow; the 14B
+release adds its evaluated step-80 schedule and memory-management details.
+
+| Read | Purpose | 1.3B | 14B |
+|---|---|---|---|
+| Configuration | Batch sizes, encoders, loss weights, initialization and optimizer | [three_encoder.yaml](../1.3B/configs/three_encoder.yaml) | [train.yaml](../14B/configs/train.yaml) |
+| Training entry | Load configuration and start distributed training | [train.py](../1.3B/code/train.py) | [train.py](../14B/code/train.py) |
+| `DistributionMatchingTrainer` | Assemble one complete update: collect rollouts, compute feature gradients, replay and update the generator | [trainer/distribution_matching.py](../1.3B/code/trainer/distribution_matching.py) | [trainer/distribution_matching.py](../14B/code/trainer/distribution_matching.py) |
+| `FeatureGradientReplay` | Gather the distribution batch and compute gradients with respect to selected features | [trainer/feature_gradient_replay.py](../1.3B/code/trainer/feature_gradient_replay.py) | [trainer/feature_gradient_replay.py](../14B/code/trainer/feature_gradient_replay.py) |
+| `MultiEncoderMMD` | Combine the independently scaled encoder losses | [losses/multi_encoder_mmd.py](../1.3B/code/losses/multi_encoder_mmd.py) | [losses/multi_encoder_mmd.py](../14B/code/losses/multi_encoder_mmd.py) |
+| `VideoNystromMMD` | Evaluate generated–generated repulsion and hybrid reference attraction | [losses/video_nystrom_mmd.py](../1.3B/code/losses/video_nystrom_mmd.py) | [losses/video_nystrom_mmd.py](../14B/code/losses/video_nystrom_mmd.py) |
+| `DistributionMatchingModel` | Connect rollout generation, frozen encoders, VAE decoding and generator replay | [model/distribution_matching.py](../1.3B/code/model/distribution_matching.py) | [model/distribution_matching.py](../14B/code/model/distribution_matching.py) |
+
+The previous module paths remain compatibility shims. Read the implementations
+linked above rather than the shims. Low-level network definitions in `wan/` are
+useful after the training flow is clear.
+
+In `DistributionMatchingTrainer`, start at `_run_distribution_matching_step` and
+follow `_collect_rollout_microbatch`. In `DistributionMatchingModel`, the corresponding
+stages are `collect_rollout_features`, `prepare_latent_gradients` and
+`replay_generator_gradients`.
+
+## Three different batch sizes
+
+| Paper symbol | Meaning | Released settings |
+|---|---|---|
+| $B$ | **Distribution batch:** generated videos that participate in the loss | 256 globally; 64 per GPU on 1.3B, 32 per GPU on 14B |
+| $b$ | **Differentiation batch:** generated videos replayed for generator backpropagation | 128 globally; 32 per GPU on 1.3B, 16 per GPU on 14B |
+| $B_{\mathrm{ref}}$ | **Reference minibatch:** stored reference-video features used by the Monte Carlo term | 128 globally |
+
+The readable configuration calls these `distribution_batch_size`,
+`distribution_samples_per_rank`, `differentiation_samples_per_rank` and
+`reference_minibatch_size`. `differentiation_gradient_scale` is $B/b$;
+`encoder_loss_divisor` is the final loss divisor. `utils/config.py` translates
+these names into the compatible runtime fields without changing their values.
+
+A microbatch is the smaller group processed at once; it is not $B$ or $b$.
+All 256 generated videos affect the loss, including the 128 whose generator
+computation is not replayed. Computing independent MMD losses on microbatches
+would omit their cross-microbatch interactions.
+
+Compatibility names can be misleading: `student_reference` means the full
+**generated** population, while `current_teacher` means the **reference
+minibatch**. Reference features are read from a frozen feature bank; this path
+does not run an online diffusion teacher.
+
+## One optimizer update
+
+```text
+Prompts + noise
+ → generate B videos without retained generator graphs
+ → extract frozen-encoder features and retain replay state
+ → gather the complete distribution batch across GPUs
+ → combine persistent Nyström statistics with a reference minibatch
+ → compute feature gradients for the b selected videos
+ → replay encoders: feature gradients → pixel gradients
+ → replay VAE: pixel gradients → latent gradients
+ → replay generator: latent gradients → parameter gradients
+ → accumulate contributions, clip gradients, update once
+```
+
+A **cotangent** is the incoming gradient used in a vector–Jacobian product. For
+example, a feature cotangent is $g_i^e=\partial L/\partial z_i^e$; replay passes it
+back through the corresponding encoder and generator. Frozen encoder and VAE
+parameters are not optimized, but derivatives with respect to their inputs are
+still required. Replay restores the original noise, random choices and segment
+inputs so that the gradient is paired with the intended computation.
+
+Two scaling factors have different purposes. The configured $B/b=2$ compensates
+for replaying only part of the distribution batch. Multiplication by the number
+of GPUs compensates for FSDP averaging parameter gradients across ranks.
+
+## Connecting the loss to the paper
+
+Section 3.1 compares generated and reference distributions in frozen encoder
+spaces. Section 3.2 combines two estimates of the reference kernel mean:
+
+- **Persistent Nyström mean:** landmark kernels multiplied by precomputed
+ coefficients, as in Equation (8) and Appendix C.2, Equations (18)–(21).
+- **Monte Carlo mean:** average kernel similarity to a sampled reference
+ minibatch, as in Equation (6).
+
+The released Nyström:Monte Carlo weights are 2:1, so the mixture in Equation (9)
+has $\alpha_e=1/3$. The tensor saved under the legacy name `alpha` is instead
+the **Nyström coefficient vector**, denoted $\boldsymbol\beta_e$ in Equation
+(19). These are different quantities.
+
+For each encoder, the released implementation evaluates
+
+$$
+M_e=\frac{1}{B^2}\sum_{i,j}k_e(z_i,z_j)
+-\frac{2}{B}\sum_i\left[\frac{2}{3}m^e_{\mathrm{Nys}}(z_i)
++\frac{1}{3}\widehat m^e_{\mathrm{MC}}(z_i)\right]+C_e,
+$$
+
+where $C_e$ is the retained reference–reference contribution. It then applies
+detached per-encoder loss scaling:
+
+$$
+L=\frac{1}{d}\sum_e w_e
+\frac{M_e}{\operatorname{stopgrad}(M_e)+\epsilon}.
+$$
+
+Here $d$ is the explicit aggregation divisor; it is not inferred from the sum
+of the weights. Both 1.3B and the human/VLM-evaluated 14B use
+DINOv3:V-JEPA 2:VideoMAEv2 = 10:1:1. Both configurations use $d=3$ and
+$\epsilon=10^{-4}$.
+
+## Details to preserve when reviewing or refactoring
+
+- **Self-pairs and constants.** Paper Equations (11)/(13) use the off-diagonal
+ denominator $B(B-1)$ and omit reference–reference constants. This release
+ retains self-pairs with denominator $B^2$ and retains the constants. Although
+ a reference constant has no direct generator derivative, it affects the
+ detached normalization denominator above. Changing either convention changes
+ the released numerical objective.
+- **Replay selection.** Algorithm 1 uniformly samples a subset. The released
+ implementation selects a fixed suffix of each rank's generated batch and
+ creates feature-gradient leaves for those selected samples. Their gradients
+ still include interactions with the full distribution batch. The conditional
+ unbiasedness statement in Equation (27) assumes uniform subset sampling;
+ it should not be claimed for this fixed-suffix implementation.
+- **VAE temporal derivatives.** The active model paths call `replay_decoder_vjp`,
+ which carries temporal-cache values but detaches their gradients between latent
+ frames. This preserves decoded values while truncating cross-frame cache
+ derivatives. The 14B tree also provides `replay_decoder_vjp_exact` and
+ `wan_recurrent_decoder_vjp.py`, which propagate future-state cotangents; their
+ presence does not mean the released training path selects them. Appendix C.4
+ distinguishes prescribed gradient truncation from exact replay.
+- **Compatibility.** Resource-pack keys, checkpoint parameter names and legacy
+ configuration fields remain supported. Their historical wording is not an
+ additional training objective or a requirement to instantiate score models.
+
+## Instrumentation and compatibility names
+
+Old imports and direct calls remain available as aliases. New training code
+calls the readable method names, so observers that intercept methods should
+patch `_collect_rollout_microbatch`, `_run_distribution_matching_step`,
+`collect_rollout_features`, `prepare_latent_gradients`, and
+`replay_generator_gradients`. Patching an old method alias does not intercept
+calls dispatched through the new name.
+
+Edit the public YAML when changing experiment settings. `prepare_config.py`
+resolves those readable fields and writes the legacy runtime schema to
+`configs/local.yaml`; this preserves resource and checkpoint identities.
diff --git a/docs/readability-changes.md b/docs/readability-changes.md
new file mode 100644
index 0000000000000000000000000000000000000000..dd52ac9bae41d996261ccad68c71729223718612
--- /dev/null
+++ b/docs/readability-changes.md
@@ -0,0 +1,44 @@
+# Source cleanup and validation
+
+The public 1.3B and 14B releases now expose the training path used by the
+reported three-encoder experiments: DINOv3, V-JEPA 2 and VideoMAEv2 features are
+combined by the hybrid Nyström–Monte Carlo MMD objective, and selected feature
+gradients are replayed to the generator.
+
+Earlier source snapshots also contained optional frame-level and plain
+Fréchet-distance objectives that were disabled in both released recipes. Their
+models, controllers, preprocessing tools, configuration fields and replay
+helpers have been removed so that the public tree follows the evaluated path.
+
+## Names that intentionally remain
+
+The code/video_fd directory remains part of both releases. Despite its
+historical directory name, it contains frozen encoder adapters, temporal
+sampling and decoder/encoder VJP utilities used by the active three-encoder MMD
+pipeline.
+
+Checkpoint methods and payload keys containing fd also remain where changing
+them would break loading existing three-encoder checkpoints. They describe the
+saved distribution-matching state; they do not enable a second loss branch.
+
+## Validation
+
+Run the root checks from the repository root in an environment with the release dependencies installed. They are self-contained and do not require an old Git commit:
+
+ python3 scripts/check_readability.py
+ python3 scripts/check_loss_parity.py --release 1.3B
+ python3 scripts/check_loss_parity.py --release 14B
+ python3 1.3B/scripts/test_cpu.py
+ python3 14B/scripts/test_cpu.py
+
+check_readability.py verifies the cleaned file layout, rejects stale auxiliary
+loss references and configuration keys, parses every Python source, checks
+shell syntax, and confirms the main scientific settings of both public recipes.
+check_loss_parity.py independently composes the three encoder branches and
+checks the active MMD values, metadata, detached normalization and feature
+gradients on CPU.
+
+The first optimizer-step report in first-step-comparison.md records a historical
+comparison made before this cleanup. It remains useful evidence for the active
+training calculation, but it is not a source-identity claim for the cleaned
+tree.
diff --git a/scripts/check_loss_parity.py b/scripts/check_loss_parity.py
new file mode 100644
index 0000000000000000000000000000000000000000..e92076cce37b698e0df2af5cdf00e966721ed065
--- /dev/null
+++ b/scripts/check_loss_parity.py
@@ -0,0 +1,169 @@
+#!/usr/bin/env python3
+"""Validate the active three-encoder MMD composition and gradients on CPU."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+
+sys.dont_write_bytecode = True
+
+
+def assert_same(left: torch.Tensor, right: torch.Tensor, label: str) -> None:
+ torch.testing.assert_close(left, right, rtol=0, atol=0, msg=label)
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--release", choices=("1.3B", "14B"), required=True)
+ args = parser.parse_args()
+
+ root = Path(__file__).resolve().parents[1]
+ sys.path.insert(0, str(root / args.release / "code"))
+
+ from losses.multi_encoder_mmd import (
+ MultiEncoderMMD,
+ normalize_mmd_with_detached_scale,
+ )
+ from losses.video_nystrom_mmd import (
+ VideoNystromMMD,
+ build_video_nystrom_reference,
+ )
+
+ shapes = {
+ "dinov3": (1, 2),
+ "vjepa2": (2, 3),
+ "videomae_v2": (1, 4),
+ }
+
+ def objective() -> MultiEncoderMMD:
+ branches = {}
+ for index, (name, shape) in enumerate(shapes.items()):
+ generator = torch.Generator().manual_seed(9000 + index)
+ reference = torch.randn(12, *shape, generator=generator)
+ pack = build_video_nystrom_reference(
+ reference,
+ num_landmarks=5,
+ metadata={
+ "feature": {"name": name, "matrix_shape": list(shape)},
+ "reference": {"dataset": "active-mmd-check"},
+ },
+ bandwidth=1.7,
+ kmeans_iterations=3,
+ ridge=1.0e-5,
+ chunk_size=16,
+ )
+ branches[name] = VideoNystromMMD(
+ pack,
+ kernel_chunk_size=3,
+ attraction_weight=2.0,
+ current_teacher_attraction_enabled=True,
+ current_teacher_attraction_weight=1.0,
+ )
+ return MultiEncoderMMD(
+ branches,
+ weights={"dinov3": 10.0, "vjepa2": 1.0, "videomae_v2": 1.0},
+ normalization_eps=1.0e-4,
+ aggregation_divisor=3.0,
+ )
+
+ generator = torch.Generator().manual_seed(1862643)
+ for mode in ("full", "partial"):
+ combined = objective()
+ generated = {
+ name: torch.randn(4, *shape, generator=generator)
+ for name, shape in shapes.items()
+ }
+ teacher = {
+ name: torch.randn(
+ 4 if mode == "full" else 2, *shape, generator=generator)
+ for name, shape in shapes.items()
+ }
+
+ if mode == "full":
+ leaves = {
+ name: value.clone().requires_grad_()
+ for name, value in generated.items()
+ }
+ loss, metadata = combined(
+ leaves, current_teacher=teacher, stage=False)
+ else:
+ indices = torch.tensor([0, 2])
+ leaves = {
+ name: value.index_select(0, indices).clone().requires_grad_()
+ for name, value in generated.items()
+ }
+ loss, metadata = combined(
+ gradient_students=leaves,
+ student_reference=generated,
+ current_teacher=teacher,
+ gradient_student_indices=indices,
+ stage=False,
+ )
+
+ expected_terms = []
+ for name, branch in combined.branches.items():
+ if mode == "full":
+ raw, _ = branch(
+ leaves[name], current_teacher=teacher[name], stage=False)
+ else:
+ raw, _ = branch(
+ gradient_students=leaves[name],
+ student_reference=generated[name],
+ current_teacher=teacher[name],
+ gradient_student_indices=indices,
+ stage=False,
+ )
+ normalized, denominator = normalize_mmd_with_detached_scale(
+ raw, combined.normalization_eps)
+ expected_terms.append(combined.branch_weights[name] * normalized)
+ assert_same(
+ metadata["branches"][name]["raw_mmd"],
+ raw.detach(),
+ f"{mode} {name} raw MMD",
+ )
+ assert_same(
+ metadata["branches"][name]["normalization_denominator"],
+ denominator.detach(),
+ f"{mode} {name} normalization denominator",
+ )
+
+ expected = torch.stack(expected_terms).sum() / 3.0
+ assert_same(loss, expected, f"{mode} combined MMD")
+ observed_gradients = torch.autograd.grad(
+ loss, tuple(leaves.values()), retain_graph=True)
+ expected_gradients = torch.autograd.grad(
+ expected, tuple(leaves.values()))
+ for name, observed, reference in zip(
+ leaves, observed_gradients, expected_gradients):
+ assert_same(observed, reference, f"{mode} {name} feature gradient")
+ assert torch.isfinite(observed).all() and observed.norm() > 0
+ print(
+ f"PASS {args.release} {mode}: weighted three-encoder MMD, "
+ "metadata and feature gradients")
+
+ eps = 1.0e-4
+ for value in (-2.0e-4, 0.0, 1.0e-8, 0.1, 7.0):
+ observed_input = torch.tensor(
+ value, dtype=torch.float64, requires_grad=True)
+ expected_input = observed_input.detach().clone().requires_grad_(True)
+ observed, observed_denominator = normalize_mmd_with_detached_scale(
+ observed_input, eps)
+ expected_denominator = expected_input.detach() + eps
+ expected = expected_input / expected_denominator
+ assert_same(observed.detach(), expected.detach(), "normalization value")
+ assert_same(
+ observed_denominator, expected_denominator,
+ "normalization denominator")
+ observed_gradient, = torch.autograd.grad(observed, observed_input)
+ expected_gradient, = torch.autograd.grad(expected, expected_input)
+ assert_same(
+ observed_gradient, expected_gradient, "normalization gradient")
+ print(f"PASS {args.release}: detached MMD normalization behavior")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/check_readability.py b/scripts/check_readability.py
new file mode 100644
index 0000000000000000000000000000000000000000..140a34565febd30fba4d5a2f6cf3aa260ccd7306
--- /dev/null
+++ b/scripts/check_readability.py
@@ -0,0 +1,231 @@
+#!/usr/bin/env python3
+"""Check the cleaned release layout, public recipes, and source syntax.
+
+This check is self-contained. It does not require the historical Git commit
+used by the earlier readability-only comparison.
+"""
+
+from __future__ import annotations
+
+import ast
+import importlib.util
+from pathlib import Path
+import subprocess
+import sys
+
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.dont_write_bytecode = True
+
+RELEASES = {
+ "1.3B": {
+ "recipe": "three_encoder.yaml",
+ "required": (
+ "code/model/generation.py",
+ "code/model/distribution_matching.py",
+ "code/trainer/distribution_matching.py",
+ "code/trainer/feature_gradient_replay.py",
+ "code/losses/multi_encoder_mmd.py",
+ "code/losses/video_nystrom_mmd.py",
+ "code/video_fd/replay/replay_video_encoder_vjp.py",
+ ),
+ "removed": (
+ "code/losses/replay_video_fd.py",
+ "code/losses/video_fd.py",
+ "code/losses/video_nystrom_reference_bank.py",
+ "code/model/fd.py",
+ "code/model/frame_fd_head.py",
+ "code/trainer/fresh_frame_fd_gradcache.py",
+ "code/trainer/two_level_fidel.py",
+ "code/utils/chunked_vae_decode.py",
+ "code/utils/fd_feature_grad.py",
+ "code/utils/video_chunking.py",
+ "code/utils/video_feature_extract.py",
+ ),
+ "expected": {
+ "total_batch_size": 256,
+ "expected_current_teacher_samples": 128,
+ "fresh_video_local_samples": 64,
+ "fresh_video_generator_grad_local_samples": 32,
+ "triple_dinov3_weight": 10.0,
+ "triple_vjepa2_weight": 1.0,
+ "triple_videomae_v2_weight": 1.0,
+ "triple_aggregation_divisor": 3.0,
+ "max_steps": 150,
+ "generator_lr_cosine_end_step": 125,
+ "distribution_loss": "wan_kv_mmd",
+ },
+ },
+ "14B": {
+ "recipe": "train.yaml",
+ "required": (
+ "code/model/generation_base.py",
+ "code/model/distribution_matching.py",
+ "code/trainer/distribution_matching.py",
+ "code/trainer/feature_gradient_replay.py",
+ "code/losses/multi_encoder_mmd.py",
+ "code/losses/video_nystrom_mmd.py",
+ "code/video_fd/replay/replay_video_encoder_vjp.py",
+ ),
+ "removed": (
+ "code/losses/replay_video_fd.py",
+ "code/losses/video_fd.py",
+ "code/losses/video_nystrom_reference_bank.py",
+ "code/model/fd.py",
+ "code/model/frame_fd_head.py",
+ "code/tools/precompute_frame_fd_stats.py",
+ "code/tools/precompute_vjepa2_fidel_reference_bank.py",
+ "code/trainer/fresh_frame_fd_gradcache.py",
+ "code/trainer/two_level_fidel.py",
+ "code/utils/chunked_vae_decode.py",
+ "code/utils/fd_feature_grad.py",
+ "code/utils/video_chunking.py",
+ "code/utils/video_feature_extract.py",
+ ),
+ "expected": {
+ "total_batch_size": 256,
+ "expected_current_teacher_samples": 128,
+ "fresh_video_local_samples": 32,
+ "fresh_video_generator_grad_local_samples": 16,
+ "triple_active_branches": [
+ "dinov3", "vjepa2", "videomae_v2"],
+ "triple_dinov3_weight": 10.0,
+ "triple_vjepa2_weight": 1.0,
+ "triple_videomae_v2_weight": 1.0,
+ "triple_aggregation_divisor": 3.0,
+ "max_steps": 80,
+ "lr_schedule_total_steps": 200,
+ "distribution_loss": "wan_kv_mmd",
+ },
+ },
+}
+
+# These names belonged only to the removed auxiliary objectives. Historical
+# checkpoint names such as fd_state_dict and the active video_fd encoder package
+# deliberately remain allowed.
+REMOVED_TOKENS = (
+ "frame_fd_",
+ "fidel_",
+ "fd_repr_",
+ "fd_guidance_scale",
+ "fresh_frame_fd_gradcache",
+ "losses.replay_video_fd",
+ "losses.video_fd",
+ "from model.fd import",
+ "import model.fd",
+ "model.frame_fd_head",
+ "precompute_frame_fd_stats",
+ "precompute_vjepa2_fidel_reference_bank",
+ "trainer.two_level_fidel",
+ "utils.chunked_vae_decode",
+ "utils.fd_feature_grad",
+ "utils.video_chunking",
+ "utils.video_feature_extract",
+)
+
+REMOVED_CONFIG_KEYS = {
+ "fd_decode_frames",
+ "fd_ema_beta",
+ "fd_exclude_first_frame",
+ "fd_feature_aggregation",
+ "fd_fid_norm_eps",
+ "fd_guidance_scale",
+ "fd_num_frames",
+ "fd_queue_size",
+}
+REMOVED_CONFIG_PREFIXES = ("frame_fd_", "fidel_", "fd_repr_")
+
+
+def load_config_helper(release: str):
+ path = ROOT / release / "code/utils/config.py"
+ spec = importlib.util.spec_from_file_location(
+ f"release_config_{release.replace('.', '_')}", path)
+ module = importlib.util.module_from_spec(spec)
+ assert spec.loader is not None
+ spec.loader.exec_module(module)
+ return module
+
+
+def nested_keys(value):
+ if isinstance(value, dict):
+ for key, child in value.items():
+ yield str(key)
+ yield from nested_keys(child)
+ elif isinstance(value, list):
+ for child in value:
+ yield from nested_keys(child)
+
+
+def check_release(release: str, settings: dict) -> tuple[int, int]:
+ release_root = ROOT / release
+
+ for relative in settings["required"]:
+ assert (release_root / relative).is_file(), (
+ f"missing active release file: {release}/{relative}")
+ for relative in settings["removed"]:
+ assert not (release_root / relative).exists(), (
+ f"removed auxiliary-loss file is still present: "
+ f"{release}/{relative}")
+
+ source_paths = sorted((release_root / "code").rglob("*.py"))
+ for path in source_paths:
+ text = path.read_text()
+ ast.parse(text, filename=str(path.relative_to(ROOT)))
+ for token in REMOVED_TOKENS:
+ assert token not in text, (
+ f"stale auxiliary-loss reference {token!r} in "
+ f"{path.relative_to(ROOT)}")
+
+ shell_paths = sorted((release_root / "scripts").rglob("*.sh"))
+ for path in shell_paths:
+ subprocess.run(["bash", "-n", str(path)], check=True)
+
+ yaml_paths = sorted((release_root / "configs").rglob("*.yaml"))
+ yaml_paths += sorted((release_root / "code/configs").rglob("*.yaml"))
+ for path in yaml_paths:
+ payload = yaml.safe_load(path.read_text())
+ stale = sorted({
+ key for key in nested_keys(payload)
+ if key in REMOVED_CONFIG_KEYS
+ or key.startswith(REMOVED_CONFIG_PREFIXES)
+ })
+ assert not stale, (
+ f"stale auxiliary-loss config keys in {path.relative_to(ROOT)}: "
+ f"{', '.join(stale)}")
+
+ recipe_path = release_root / "configs" / settings["recipe"]
+ public = yaml.safe_load(recipe_path.read_text())
+ resolved = load_config_helper(release).resolve_config_aliases(public)
+ for key, expected in settings["expected"].items():
+ actual = resolved.get(key)
+ assert actual == expected, (
+ f"{recipe_path.relative_to(ROOT)}: {key}={actual!r}, "
+ f"expected {expected!r}")
+
+ print(
+ f"PASS {release}: active three-encoder recipe, cleaned layout, "
+ f"{len(source_paths)} Python files and {len(shell_paths)} shell files")
+ return len(source_paths), len(shell_paths)
+
+
+def main() -> None:
+ python_count = 0
+ shell_count = 0
+ for release, settings in RELEASES.items():
+ parsed, checked_shell = check_release(release, settings)
+ python_count += parsed
+ shell_count += checked_shell
+
+ root_scripts = sorted((ROOT / "scripts").glob("*.py"))
+ for path in root_scripts:
+ ast.parse(path.read_text(), filename=str(path.relative_to(ROOT)))
+ print(
+ f"PASS root scripts: {len(root_scripts)} Python files parse; "
+ f"total release syntax checks {python_count} Python / "
+ f"{shell_count} shell files")
+
+
+if __name__ == "__main__":
+ main()