1107 lines
44 KiB
Python
1107 lines
44 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""CausalWanModel — 40-layer DiT with causal attention and KV cache.
|
|
|
|
Key differences from WanTransformer3DModel:
|
|
- Causal self-attention (new frames only see history)
|
|
- KV cache for streaming inference
|
|
- Action/state token support (appended after video tokens)
|
|
- Extended RoPE with action/state-specific frequencies
|
|
- Inference-only forward with KV cache
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from vllm.distributed import (
|
|
get_tensor_model_parallel_rank,
|
|
get_tensor_model_parallel_world_size,
|
|
tensor_model_parallel_all_reduce,
|
|
)
|
|
from vllm.model_executor.layers.conv import Conv3dLayer
|
|
from vllm.model_executor.layers.linear import (
|
|
ColumnParallelLinear,
|
|
QKVParallelLinear,
|
|
RowParallelLinear,
|
|
)
|
|
from vllm.model_executor.utils import set_weight_attrs
|
|
|
|
from vllm_omni.diffusion.attention.layer import Attention
|
|
from vllm_omni.diffusion.models.dreamzero.action_encoder import (
|
|
CategorySpecificMLP,
|
|
MultiEmbodimentActionEncoder,
|
|
)
|
|
|
|
# AR-Diffusion paged self-attention (in-tree experimental engine). Import at
|
|
# module level so the isinstance check + custom-op call trace cleanly inside
|
|
# the fullgraph-compiled DiT block (an import inside the traced region would
|
|
# graph-break). The model still works without the engine: the payload type is
|
|
# only ever constructed by the AR-Diffusion runner.
|
|
try:
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
|
|
ARDiffusionPagedLayerInputs,
|
|
paged_write_attn,
|
|
)
|
|
except ImportError: # pragma: no cover - experimental package always ships in-tree
|
|
ARDiffusionPagedLayerInputs = None
|
|
paged_write_attn = None
|
|
|
|
# ── RoPE utilities ──────────────────────────────────────────────────
|
|
|
|
|
|
def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor:
|
|
"""Sinusoidal positional embedding for timesteps."""
|
|
if dim % 2 != 0:
|
|
raise ValueError(f"dim must be even, got {dim}.")
|
|
half = dim // 2
|
|
position = position.type(torch.float64)
|
|
sinusoid = torch.outer(
|
|
position,
|
|
torch.pow(10000, -torch.arange(half, dtype=position.dtype, device=position.device).div(half)),
|
|
)
|
|
x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
|
|
return x
|
|
|
|
|
|
def rope_params(max_seq_len: int, dim: int) -> torch.Tensor:
|
|
"""Precompute complex-valued RoPE frequencies (polar form).
|
|
Returns: complex tensor [max_seq_len, dim // 2]
|
|
"""
|
|
if dim % 2 != 0:
|
|
raise ValueError(f"dim must be even, got {dim}.")
|
|
freqs = torch.outer(
|
|
torch.arange(max_seq_len),
|
|
1.0 / torch.pow(10000, torch.arange(0, dim, 2).to(torch.float64).div(dim)),
|
|
)
|
|
freqs = torch.polar(torch.ones_like(freqs), freqs)
|
|
return freqs
|
|
|
|
|
|
def rope_apply(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
|
|
"""Apply RoPE to x using precomputed complex freqs."""
|
|
B, seq_len, n, _ = x.shape
|
|
x = torch.view_as_complex(x.to(torch.float64).reshape(B, seq_len, n, -1, 2))
|
|
freqs = freqs.unsqueeze(0)
|
|
x = torch.view_as_real(x * freqs).flatten(3)
|
|
return x
|
|
|
|
|
|
def rope_action_apply(
|
|
x: torch.Tensor,
|
|
freqs: torch.Tensor,
|
|
freqs_action: torch.Tensor,
|
|
freqs_state: torch.Tensor,
|
|
action_register_length: int | None,
|
|
num_action_per_block: int = 32,
|
|
num_state_per_block: int = 1,
|
|
) -> torch.Tensor:
|
|
"""RoPE with action/state frequency tables for multi-step sequences."""
|
|
B, seq_len, n, _ = x.shape
|
|
x = torch.view_as_complex(x.to(torch.float64).reshape(B, seq_len, n, -1, 2))
|
|
if action_register_length is not None:
|
|
if num_action_per_block is None:
|
|
raise ValueError("num_action_per_block is required when action_register_length is set.")
|
|
if num_state_per_block is None:
|
|
raise ValueError("num_state_per_block is required when action_register_length is set.")
|
|
chunk_size = action_register_length // (num_action_per_block + num_state_per_block)
|
|
freqs_1d_action = freqs_action[: chunk_size * num_action_per_block].view(
|
|
chunk_size * num_action_per_block, 1, -1
|
|
)
|
|
freqs_1d_state = freqs_state[: chunk_size * num_state_per_block].view(chunk_size * num_state_per_block, 1, -1)
|
|
freqs = torch.cat([freqs, freqs_1d_action, freqs_1d_state], dim=0)
|
|
freqs = freqs.unsqueeze(0)
|
|
x = torch.view_as_real(x * freqs).flatten(3)
|
|
return x
|
|
|
|
|
|
def causal_rope_action_apply(
|
|
x: torch.Tensor,
|
|
freqs: torch.Tensor,
|
|
freqs_action: torch.Tensor,
|
|
freqs_state: torch.Tensor,
|
|
action_register_length: int | None,
|
|
num_action_per_block: int,
|
|
num_state_per_block: int,
|
|
action_state_index: int,
|
|
) -> torch.Tensor:
|
|
"""RoPE for single inference step (causal / KV-cache mode)."""
|
|
B, seq_len, n, _ = x.shape
|
|
x = torch.view_as_complex(x.to(torch.float64).reshape(B, seq_len, n, -1, 2))
|
|
if action_register_length is not None:
|
|
expected_length = num_action_per_block + num_state_per_block
|
|
if action_register_length != expected_length:
|
|
raise ValueError(
|
|
f"action_register_length must equal num_action_per_block + num_state_per_block "
|
|
f"({expected_length}), got {action_register_length}."
|
|
)
|
|
freqs_action = freqs_action[
|
|
action_state_index * num_action_per_block : (action_state_index + 1) * num_action_per_block
|
|
]
|
|
freqs_state = freqs_state[
|
|
action_state_index * num_state_per_block : (action_state_index + 1) * num_state_per_block
|
|
]
|
|
freqs_1d = torch.cat([freqs_action, freqs_state], dim=0).view(action_register_length, 1, -1)
|
|
freqs = torch.cat([freqs, freqs_1d], dim=0)
|
|
freqs = freqs.unsqueeze(0)
|
|
x = torch.view_as_real(x * freqs).flatten(3)
|
|
return x
|
|
|
|
|
|
# ── Normalization ───────────────────────────────────────────────────
|
|
|
|
|
|
class WanLayerNorm(nn.LayerNorm):
|
|
"""LayerNorm wrapper used by DreamZero blocks."""
|
|
|
|
def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = False) -> None:
|
|
super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
|
|
|
|
|
|
class DistributedRMSNorm(nn.Module):
|
|
"""RMSNorm that computes global RMS across tensor parallel ranks."""
|
|
|
|
def __init__(self, hidden_size: int, eps: float = 1e-6):
|
|
super().__init__()
|
|
self.eps = eps
|
|
self.weight = nn.Parameter(torch.ones(hidden_size))
|
|
set_weight_attrs(self.weight, {"weight_loader": self.weight_loader})
|
|
|
|
def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
|
|
if param.shape == loaded_weight.shape:
|
|
param.data.copy_(loaded_weight)
|
|
return
|
|
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
if loaded_weight.shape[0] % tp_size != 0:
|
|
raise ValueError(
|
|
f"Cannot shard RMSNorm weight of shape {tuple(loaded_weight.shape)} across tp_size={tp_size}."
|
|
)
|
|
|
|
shard_size = loaded_weight.shape[0] // tp_size
|
|
start_idx = get_tensor_model_parallel_rank() * shard_size
|
|
shard = loaded_weight.narrow(0, start_idx, shard_size)
|
|
if param.shape != shard.shape:
|
|
raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.")
|
|
param.data.copy_(shard)
|
|
|
|
def _local_sum_sq(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
|
|
"""Per-rank float32 activation, local sum-of-squares, and local width."""
|
|
x_float = x.float()
|
|
local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
|
|
return x_float, local_sum_sq, x.shape[-1]
|
|
|
|
def _scale(
|
|
self,
|
|
x_float: torch.Tensor,
|
|
global_sum_sq: torch.Tensor,
|
|
global_count: int,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Apply the (already reduced) RMS and the per-rank weight shard."""
|
|
mean_sq = global_sum_sq / global_count
|
|
return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
x_float, local_sum_sq, local_count = self._local_sum_sq(x)
|
|
|
|
if tp_size > 1:
|
|
# Use vLLM's collective (custom all-reduce / symmetric-mem fast path
|
|
# for small tensors) instead of raw torch.distributed.all_reduce, and
|
|
# take the return value so the custom-AR path (which may return a new
|
|
# buffer) is handled correctly. No .clone() needed.
|
|
global_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq)
|
|
global_count = local_count * tp_size
|
|
else:
|
|
global_sum_sq = local_sum_sq
|
|
global_count = local_count
|
|
|
|
return self._scale(x_float, global_sum_sq, global_count, x)
|
|
|
|
|
|
def fused_qk_rms_norm(
|
|
norm_q: nn.Module,
|
|
norm_k: nn.Module,
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Apply q/k :class:`DistributedRMSNorm` with a SINGLE fused TP all-reduce.
|
|
|
|
WHY: self-attention norms q and k every step; on TP each norm issues its own
|
|
tiny all-reduce of the per-token sum-of-squares. For DreamZero's decode-like
|
|
per-step forward these latency-bound micro-collectives dominate, so we pack
|
|
both sum-of-squares into one tensor and reduce once (2 collectives → 1).
|
|
|
|
NUMERICALLY IDENTICAL to ``norm_q(q), norm_k(k)``: all-reduce is elementwise,
|
|
so packing along the last dim reduces each slice independently with the same
|
|
fp32 accumulation. Requires q and k to share the same shape (true for
|
|
self-attention — both come from the same hidden states). Falls back to
|
|
independent application when either norm is not a DistributedRMSNorm
|
|
(e.g. nn.Identity when qk_norm=False).
|
|
"""
|
|
if not (isinstance(norm_q, DistributedRMSNorm) and isinstance(norm_k, DistributedRMSNorm)):
|
|
return norm_q(q), norm_k(k)
|
|
|
|
# The fused path reduces one packed sum-of-squares and reuses q's token width
|
|
# as the RMS count for BOTH q and k, so q and k must share the same shape.
|
|
# (Self-attention guarantees this: q and k are projected from the same x.)
|
|
assert q.shape == k.shape, "fused_qk_rms_norm requires q and k to have the same shape."
|
|
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
q_float, q_sum_sq, count = norm_q._local_sum_sq(q)
|
|
k_float, k_sum_sq, _ = norm_k._local_sum_sq(k)
|
|
|
|
if tp_size > 1:
|
|
packed = torch.cat([q_sum_sq, k_sum_sq], dim=-1)
|
|
packed = tensor_model_parallel_all_reduce(packed)
|
|
q_sum_sq, k_sum_sq = packed[..., 0:1], packed[..., 1:2]
|
|
count = count * tp_size
|
|
|
|
q_out = norm_q._scale(q_float, q_sum_sq, count, q)
|
|
k_out = norm_k._scale(k_float, k_sum_sq, count, k)
|
|
return q_out, k_out
|
|
|
|
|
|
# ── Projections ─────────────────────────────────────────────────────
|
|
|
|
|
|
class MLPProj(nn.Module):
|
|
"""CLIP feature projection for i2v.
|
|
Uses ColumnParallelLinear + RowParallelLinear (Qwen3_VisionMLP pattern).
|
|
"""
|
|
|
|
def __init__(self, in_dim: int, out_dim: int) -> None:
|
|
super().__init__()
|
|
self.norm1 = nn.LayerNorm(in_dim)
|
|
self.fc1 = ColumnParallelLinear(
|
|
in_dim,
|
|
in_dim,
|
|
bias=True,
|
|
return_bias=False,
|
|
)
|
|
self.act = nn.GELU()
|
|
self.fc2 = RowParallelLinear(
|
|
in_dim,
|
|
out_dim,
|
|
bias=True,
|
|
return_bias=False,
|
|
)
|
|
self.norm2 = nn.LayerNorm(out_dim)
|
|
|
|
def forward(self, image_embeds: torch.Tensor) -> torch.Tensor:
|
|
x = self.norm1(image_embeds)
|
|
x = self.fc1(x)
|
|
x = self.act(x)
|
|
x = self.fc2(x)
|
|
x = self.norm2(x)
|
|
return x
|
|
|
|
|
|
# ── Cross-Attention ─────────────────────────────────────────────────
|
|
# T2V and I2V cross-attention variants
|
|
|
|
|
|
class WanT2VCrossAttention(nn.Module):
|
|
"""Text-to-video cross-attention.
|
|
Uses vllm-omni Attention for FlashAttn backend.
|
|
"""
|
|
|
|
def __init__(self, dim: int, num_heads: int, window_size=(-1, -1), qk_norm: bool = True, eps: float = 1e-6) -> None:
|
|
super().__init__()
|
|
if dim % num_heads != 0:
|
|
raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
|
|
self.dim = dim
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
if num_heads % tp_size != 0:
|
|
raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.")
|
|
self.tp_num_heads = num_heads // tp_size
|
|
self.tp_inner_dim = self.tp_num_heads * self.head_dim
|
|
self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
|
|
self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
|
|
self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
|
|
self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False)
|
|
self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
|
self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
|
self.attn = Attention(
|
|
self.tp_num_heads,
|
|
self.head_dim,
|
|
causal=False,
|
|
softmax_scale=self.head_dim**-0.5,
|
|
skip_sequence_parallel=True,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
context: torch.Tensor,
|
|
context_lens: torch.Tensor | None = None,
|
|
crossattn_cache: dict | None = None,
|
|
) -> torch.Tensor:
|
|
del context_lens
|
|
n, d = self.tp_num_heads, self.head_dim
|
|
q = self.norm_q(self.q(x)).unflatten(2, (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)).unflatten(2, (n, d))
|
|
v = self.v(context).unflatten(2, (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)).unflatten(2, (n, d))
|
|
v = self.v(context).unflatten(2, (n, d))
|
|
x = self.attn(q, k, v)
|
|
x = x.flatten(2)
|
|
x = self.o(x)
|
|
return x
|
|
|
|
|
|
class WanI2VCrossAttention(nn.Module):
|
|
"""Image-to-video cross-attention (splits first 257 image tokens).
|
|
Uses vllm-omni Attention for FlashAttn backend.
|
|
"""
|
|
|
|
def __init__(self, dim: int, num_heads: int, window_size=(-1, -1), qk_norm: bool = True, eps: float = 1e-6) -> None:
|
|
super().__init__()
|
|
if dim % num_heads != 0:
|
|
raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
|
|
self.dim = dim
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
if num_heads % tp_size != 0:
|
|
raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.")
|
|
self.tp_num_heads = num_heads // tp_size
|
|
self.tp_inner_dim = self.tp_num_heads * self.head_dim
|
|
self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
|
|
self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
|
|
self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
|
|
self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False)
|
|
self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
|
self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
|
self.k_img = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
|
|
self.v_img = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False)
|
|
self.norm_k_img = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
|
self.attn = Attention(
|
|
self.tp_num_heads,
|
|
self.head_dim,
|
|
causal=False,
|
|
softmax_scale=self.head_dim**-0.5,
|
|
skip_sequence_parallel=True,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
context: torch.Tensor,
|
|
context_lens: torch.Tensor | None = None,
|
|
crossattn_cache: dict | None = None,
|
|
) -> torch.Tensor:
|
|
del context_lens
|
|
context_img = context[:, :257]
|
|
context = context[:, 257:]
|
|
n, d = self.tp_num_heads, self.head_dim
|
|
q = self.norm_q(self.q(x)).unflatten(2, (n, d))
|
|
# context (text) and context_img (clip features) are constant within a
|
|
# session, so k/v and k_img/v_img are cached on first call and reused.
|
|
# context_img == img_emb(clip feature) is session-invariant (state.clip_feas
|
|
# is set once at current_start_frame==0 and only cleared on reset, which also
|
|
# rebuilds crossattn_cache). Caching k_img/v_img removes a per-step
|
|
# norm_k_img all-reduce plus the k_img/v_img projection GEMMs in steady state.
|
|
# Only q depends on the per-step x and is always recomputed.
|
|
if crossattn_cache is not None:
|
|
if not crossattn_cache["is_init"]:
|
|
crossattn_cache["is_init"] = True
|
|
k = self.norm_k(self.k(context)).unflatten(2, (n, d))
|
|
v = self.v(context).unflatten(2, (n, d))
|
|
k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
|
|
v_img = self.v_img(context_img).unflatten(2, (n, d))
|
|
crossattn_cache["k"] = k
|
|
crossattn_cache["v"] = v
|
|
crossattn_cache["k_img"] = k_img
|
|
crossattn_cache["v_img"] = v_img
|
|
else:
|
|
k = crossattn_cache["k"]
|
|
v = crossattn_cache["v"]
|
|
k_img = crossattn_cache["k_img"]
|
|
v_img = crossattn_cache["v_img"]
|
|
else:
|
|
k = self.norm_k(self.k(context)).unflatten(2, (n, d))
|
|
v = self.v(context).unflatten(2, (n, d))
|
|
k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d))
|
|
v_img = self.v_img(context_img).unflatten(2, (n, d))
|
|
x = self.attn(q, k, v)
|
|
img_x = self.attn(q, k_img, v_img)
|
|
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,
|
|
}
|
|
|
|
|
|
# ── Self-Attention with causal masking + KV cache ───────────────────
|
|
|
|
|
|
class CausalWanSelfAttention(nn.Module):
|
|
"""Causal self-attention with KV cache + action/state tokens."""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_heads: int,
|
|
frame_seqlen: int,
|
|
local_attn_size: int = -1,
|
|
sink_size: int = 0,
|
|
num_frame_per_block: int = 1,
|
|
qk_norm: bool = True,
|
|
eps: float = 1e-6,
|
|
num_action_per_block: int = 32,
|
|
num_state_per_block: int = 1,
|
|
) -> None:
|
|
super().__init__()
|
|
if dim % num_heads != 0:
|
|
raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
|
|
self.dim = dim
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
if num_heads % tp_size != 0:
|
|
raise ValueError(f"num_heads={num_heads} must be divisible by tp_size={tp_size}.")
|
|
self.tp_num_heads = num_heads // tp_size
|
|
self.tp_inner_dim = self.tp_num_heads * self.head_dim
|
|
self.local_attn_size = local_attn_size
|
|
self.num_frame_per_block = num_frame_per_block
|
|
self.frame_seqlen = frame_seqlen
|
|
self.num_action_per_block = num_action_per_block
|
|
self.num_state_per_block = num_state_per_block
|
|
self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen
|
|
# Fused QKV projection: q/k/v all come from x, so a single column-parallel
|
|
# GEMM replaces three (fewer launches / better GEMM util in the decode-like
|
|
# per-step forward). No GQA here: total_num_kv_heads defaults to num_heads.
|
|
self.qkv = QKVParallelLinear(
|
|
hidden_size=dim,
|
|
head_size=self.head_dim,
|
|
total_num_heads=num_heads,
|
|
bias=True,
|
|
)
|
|
# The forward splits qkv into three equal q/k/v shards, which is only
|
|
# correct without GQA. total_num_kv_heads defaults to total_num_heads
|
|
# here; fail loud if a future checkpoint introduces GQA so the equal
|
|
# split does not silently misalign k/v.
|
|
if self.qkv.total_num_kv_heads != self.qkv.total_num_heads:
|
|
raise ValueError("Self-attn QKV fusion requires no GQA (total_num_kv_heads == total_num_heads).")
|
|
self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False)
|
|
self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
|
self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
|
self.attn = Attention(
|
|
self.tp_num_heads,
|
|
self.head_dim,
|
|
causal=False,
|
|
softmax_scale=self.head_dim**-0.5,
|
|
skip_sequence_parallel=True,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
freqs: torch.Tensor,
|
|
freqs_action: torch.Tensor,
|
|
freqs_state: torch.Tensor,
|
|
action_register_length: int | None,
|
|
kv_cache: torch.Tensor | Any | None = None,
|
|
current_start_frame: int = 0,
|
|
is_tf: bool = True,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
"""Inference-only forward (KV cache path)."""
|
|
n, d = self.tp_num_heads, self.head_dim
|
|
|
|
# Single fused QKV GEMM, then split into the per-rank q/k/v shards.
|
|
qkv, _ = self.qkv(x)
|
|
qk_size = self.tp_num_heads * self.head_dim
|
|
q, k, v = qkv.split([qk_size, qk_size, qk_size], dim=-1)
|
|
# Fuse q/k qk-norm into a single TP all-reduce (both come from x, same shape).
|
|
q, k = fused_qk_rms_norm(self.norm_q, self.norm_k, q, k)
|
|
q = q.unflatten(2, (n, d))
|
|
k = k.unflatten(2, (n, d))
|
|
v = v.unflatten(2, (n, d))
|
|
|
|
updated_kv_cache: torch.Tensor | None = None
|
|
|
|
if kv_cache is None:
|
|
raise RuntimeError("Inference only: kv_cache is required.")
|
|
|
|
action_state_index = max(0, (current_start_frame - 1) // self.num_frame_per_block)
|
|
|
|
roped_query = causal_rope_action_apply(
|
|
q,
|
|
freqs,
|
|
freqs_action,
|
|
freqs_state,
|
|
action_register_length,
|
|
self.num_action_per_block,
|
|
self.num_state_per_block,
|
|
action_state_index,
|
|
).type_as(v)
|
|
roped_key = causal_rope_action_apply(
|
|
k,
|
|
freqs,
|
|
freqs_action,
|
|
freqs_state,
|
|
action_register_length,
|
|
self.num_action_per_block,
|
|
self.num_state_per_block,
|
|
action_state_index,
|
|
).type_as(v)
|
|
|
|
roped_action_query = None
|
|
roped_action_key = None
|
|
action_v = None
|
|
|
|
if action_register_length is not None:
|
|
roped_action_query = roped_query[:, -action_register_length:]
|
|
roped_query = roped_query[:, :-action_register_length]
|
|
roped_action_key = roped_key[:, -action_register_length:]
|
|
roped_key = roped_key[:, :-action_register_length]
|
|
action_v = v[:, -action_register_length:]
|
|
v = v[:, :-action_register_length]
|
|
|
|
if ARDiffusionPagedLayerInputs is not None and isinstance(kv_cache, ARDiffusionPagedLayerInputs):
|
|
# Fused write+attend custom op: one opaque node in the compiled
|
|
# graph (slot writes + FlashAttention block-table kernel inside).
|
|
# Metadata tensors were prepared once per forward in _forward_blocks.
|
|
if action_register_length is not None:
|
|
q_cat = torch.cat([roped_query, roped_action_query], dim=1)
|
|
k_act, v_act = roped_action_key[0], action_v[0]
|
|
else:
|
|
q_cat = roped_query
|
|
k_act = v_act = None
|
|
x = paged_write_attn(
|
|
kv_cache,
|
|
q_cat[0],
|
|
roped_key[0],
|
|
v[0],
|
|
k_act,
|
|
v_act,
|
|
self.head_dim**-0.5,
|
|
).unsqueeze(0)
|
|
else:
|
|
updated_k = kv_cache[0]
|
|
updated_v = kv_cache[1]
|
|
new_k = torch.cat([updated_k, roped_key], dim=1)
|
|
new_v = torch.cat([updated_v, v], dim=1)
|
|
new_k = new_k[:, -self.max_attention_size :]
|
|
new_v = new_v[:, -self.max_attention_size :]
|
|
|
|
if action_register_length is not None:
|
|
q_cat = torch.cat([roped_query, roped_action_query], dim=1)
|
|
k_cat = torch.cat([new_k, roped_action_key], dim=1)
|
|
v_cat = torch.cat([new_v, action_v], dim=1)
|
|
else:
|
|
q_cat = roped_query
|
|
k_cat = new_k
|
|
v_cat = new_v
|
|
|
|
x = self.attn(q_cat, k_cat, v_cat)
|
|
updated_kv_cache = torch.stack([new_k, new_v], dim=0)
|
|
|
|
x = x.flatten(2)
|
|
x = self.o(x)
|
|
return x, updated_kv_cache
|
|
|
|
|
|
# ── Attention Block ─────────────────────────────────────────────────
|
|
|
|
|
|
class CausalWanAttentionBlock(nn.Module):
|
|
"""Transformer block: self-attn + cross-attn + FFN with 6-param modulation."""
|
|
|
|
def __init__(
|
|
self,
|
|
cross_attn_type: str,
|
|
dim: int,
|
|
ffn_dim: int,
|
|
num_heads: int,
|
|
frame_seqlen: int,
|
|
local_attn_size: int = -1,
|
|
sink_size: int = 0,
|
|
num_frame_per_block: int = 1,
|
|
qk_norm: bool = True,
|
|
cross_attn_norm: bool = False,
|
|
eps: float = 1e-6,
|
|
num_action_per_block: int = 32,
|
|
num_state_per_block: int = 1,
|
|
) -> None:
|
|
super().__init__()
|
|
self.norm1 = WanLayerNorm(dim, eps)
|
|
self.self_attn = CausalWanSelfAttention(
|
|
dim=dim,
|
|
num_heads=num_heads,
|
|
frame_seqlen=frame_seqlen,
|
|
local_attn_size=local_attn_size,
|
|
sink_size=sink_size,
|
|
num_frame_per_block=num_frame_per_block,
|
|
qk_norm=qk_norm,
|
|
eps=eps,
|
|
num_action_per_block=num_action_per_block,
|
|
num_state_per_block=num_state_per_block,
|
|
)
|
|
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(
|
|
ColumnParallelLinear(dim, ffn_dim, bias=True, gather_output=False, return_bias=False),
|
|
nn.GELU(approximate="tanh"),
|
|
RowParallelLinear(ffn_dim, dim, bias=True, input_is_parallel=True, return_bias=False),
|
|
)
|
|
self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
e: torch.Tensor,
|
|
freqs: torch.Tensor,
|
|
freqs_action: torch.Tensor,
|
|
freqs_state: torch.Tensor,
|
|
context: torch.Tensor,
|
|
action_register_length: int | None = None,
|
|
kv_cache: torch.Tensor | Any | None = None,
|
|
crossattn_cache: dict | None = None,
|
|
current_start_frame: int = 0,
|
|
is_tf: bool = True,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
|
e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
|
|
|
|
y, updated_kv_cache = self.self_attn(
|
|
x=(self.norm1(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)),
|
|
freqs=freqs,
|
|
freqs_action=freqs_action,
|
|
freqs_state=freqs_state,
|
|
action_register_length=action_register_length,
|
|
kv_cache=kv_cache,
|
|
is_tf=is_tf,
|
|
current_start_frame=current_start_frame,
|
|
)
|
|
x = x + (y * e[2].squeeze(2))
|
|
|
|
x = x + self.cross_attn(self.norm3(x), context, crossattn_cache=crossattn_cache)
|
|
y = self.ffn(self.norm2(x) * (1 + e[4].squeeze(2)) + e[3].squeeze(2))
|
|
x = x + (y * e[5].squeeze(2))
|
|
return x, updated_kv_cache
|
|
|
|
|
|
# ── Output Head ─────────────────────────────────────────────────────
|
|
|
|
|
|
class CausalHead(nn.Module):
|
|
"""Output norm + linear with 2-param modulation.
|
|
Runs once per step (not TP-critical), uses nn.Linear.
|
|
"""
|
|
|
|
def __init__(self, dim: int, out_dim: int, patch_size: tuple, eps: float = 1e-6) -> None:
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.out_dim = out_dim
|
|
self.patch_size = patch_size
|
|
out_channels = math.prod(patch_size) * out_dim
|
|
self.norm = WanLayerNorm(dim, eps)
|
|
self.head = nn.Linear(dim, out_channels)
|
|
self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
|
|
|
|
def forward(self, x: torch.Tensor, e: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Args:
|
|
x: [B, L1, C]
|
|
e: [B, F, 1, C] (time embedding, unsqueezed)
|
|
"""
|
|
e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
|
|
x = self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2))
|
|
return x
|
|
|
|
|
|
# ── Main Model ──────────────────────────────────────────────────────
|
|
|
|
|
|
class CausalWanModel(nn.Module):
|
|
"""Causal video diffusion transformer for DreamZero.
|
|
|
|
Architecture (14B): 40 layers, dim=5120, heads=40, ffn=13824
|
|
"""
|
|
|
|
_layerwise_offload_blocks_attrs = ["blocks"]
|
|
|
|
def __init__(
|
|
self,
|
|
model_type: str = "t2v",
|
|
patch_size: tuple[int, int, int] = (1, 2, 2),
|
|
frame_seqlen: int = 220,
|
|
text_len: int = 512,
|
|
in_dim: int = 16,
|
|
dim: int = 2048,
|
|
ffn_dim: int = 8192,
|
|
freq_dim: int = 256,
|
|
text_dim: int = 4096,
|
|
out_dim: int = 16,
|
|
num_heads: int = 16,
|
|
num_layers: int = 32,
|
|
max_chunk_size: int = -1,
|
|
sink_size: int = 0,
|
|
qk_norm: bool = True,
|
|
cross_attn_norm: bool = True,
|
|
eps: float = 1e-6,
|
|
num_frame_per_block: int = 1,
|
|
action_dim: int = 32,
|
|
num_registers: int = 8,
|
|
max_state_dim: int = 64,
|
|
max_num_embodiments: int = 32,
|
|
hidden_size: int = 1024,
|
|
diffusion_model_pretrained_path: str | None = None,
|
|
num_action_per_block: int = 32,
|
|
num_state_per_block: int = 1,
|
|
) -> None:
|
|
super().__init__()
|
|
if model_type not in ["t2v", "i2v", "ti2v"]:
|
|
raise ValueError(f"Unsupported model_type={model_type!r}; expected one of ['t2v', 'i2v', 'ti2v'].")
|
|
self.model_type = model_type
|
|
self.patch_size = patch_size
|
|
self.frame_seqlen = frame_seqlen
|
|
self.text_len = text_len
|
|
self.dim = dim
|
|
self.freq_dim = freq_dim
|
|
self.out_dim = out_dim
|
|
self.num_heads = num_heads
|
|
self.num_layers = num_layers
|
|
self.local_attn_size = max_chunk_size * num_frame_per_block + 1 if max_chunk_size != -1 else -1
|
|
self.num_frame_per_block = num_frame_per_block
|
|
self.action_dim = action_dim
|
|
self.num_action_per_block = num_action_per_block
|
|
self.num_state_per_block = num_state_per_block
|
|
|
|
max_num_embodiments_local = 1
|
|
self.state_encoder = CategorySpecificMLP(
|
|
num_categories=max_num_embodiments_local,
|
|
input_dim=max_state_dim,
|
|
hidden_dim=hidden_size,
|
|
output_dim=dim,
|
|
)
|
|
self.action_encoder = MultiEmbodimentActionEncoder(
|
|
action_dim=action_dim,
|
|
hidden_size=dim,
|
|
num_embodiments=max_num_embodiments_local,
|
|
)
|
|
self.action_decoder = CategorySpecificMLP(
|
|
num_categories=max_num_embodiments_local,
|
|
input_dim=dim,
|
|
hidden_dim=hidden_size,
|
|
output_dim=action_dim,
|
|
)
|
|
|
|
# Disable the Conv3d GEMM rewrite for patch embedding.
|
|
self.patch_embedding = Conv3dLayer(
|
|
in_dim,
|
|
dim,
|
|
kernel_size=patch_size,
|
|
stride=patch_size,
|
|
)
|
|
self.patch_embedding.enable_linear = False
|
|
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),
|
|
)
|
|
|
|
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,
|
|
frame_seqlen,
|
|
self.local_attn_size,
|
|
sink_size,
|
|
num_frame_per_block,
|
|
qk_norm,
|
|
cross_attn_norm,
|
|
eps,
|
|
num_action_per_block,
|
|
num_state_per_block,
|
|
)
|
|
for _ in range(num_layers)
|
|
]
|
|
)
|
|
|
|
self.head = CausalHead(dim, out_dim, patch_size, eps)
|
|
|
|
if dim % num_heads != 0:
|
|
raise ValueError(f"dim={dim} must be divisible by num_heads={num_heads}.")
|
|
if (dim // num_heads) % 2 != 0:
|
|
raise ValueError(f"dim // num_heads must be even, got {dim // num_heads}.")
|
|
d = dim // num_heads
|
|
self.freqs_action = rope_params(1024 * 10, d)
|
|
self.freqs_state = rope_params(1024, d)
|
|
self.freqs = [
|
|
rope_params(1024, d - 4 * (d // 6)),
|
|
rope_params(1024, 2 * (d // 6)),
|
|
rope_params(1024, 2 * (d // 6)),
|
|
]
|
|
|
|
if model_type == "i2v":
|
|
self.img_emb = MLPProj(1280, dim)
|
|
|
|
self.init_weights()
|
|
|
|
def init_weights(self) -> None:
|
|
"""Initialize parameters."""
|
|
|
|
def _init_linear_like(module: nn.Module) -> None:
|
|
if isinstance(module, nn.Linear):
|
|
nn.init.xavier_uniform_(module.weight)
|
|
if module.bias is not None:
|
|
nn.init.zeros_(module.bias)
|
|
return
|
|
|
|
if isinstance(module, (ColumnParallelLinear, RowParallelLinear)):
|
|
fan_in = module.input_size
|
|
fan_out = module.output_size
|
|
bound = math.sqrt(6.0 / float(fan_in + fan_out))
|
|
nn.init.uniform_(module.weight, -bound, bound)
|
|
if module.bias is not None:
|
|
nn.init.zeros_(module.bias)
|
|
|
|
for module in self.modules():
|
|
_init_linear_like(module)
|
|
|
|
nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
|
|
if self.patch_embedding.bias is not None:
|
|
fan_in = self.patch_embedding.in_channels * math.prod(self.patch_embedding.kernel_size)
|
|
bound = 1 / math.sqrt(fan_in)
|
|
nn.init.uniform_(self.patch_embedding.bias, -bound, bound)
|
|
|
|
for module in self.text_embedding.modules():
|
|
if isinstance(module, nn.Linear):
|
|
nn.init.normal_(module.weight, std=0.02)
|
|
|
|
for module in self.time_embedding.modules():
|
|
if isinstance(module, nn.Linear):
|
|
nn.init.normal_(module.weight, std=0.02)
|
|
|
|
nn.init.zeros_(self.head.head.weight)
|
|
|
|
def _create_freqs(self, grid_size: torch.Tensor, start_frame: int) -> torch.Tensor:
|
|
"""Create 3D RoPE frequency tensor."""
|
|
device = self.patch_embedding.weight.device
|
|
if any(freq.device != device for freq in self.freqs):
|
|
self.freqs = [freq.to(device) for freq in self.freqs]
|
|
if self.freqs_action.device != device:
|
|
self.freqs_action = self.freqs_action.to(device)
|
|
if self.freqs_state.device != device:
|
|
self.freqs_state = self.freqs_state.to(device)
|
|
|
|
f, h, w = grid_size.tolist()
|
|
freqs = torch.cat(
|
|
[
|
|
self.freqs[0][start_frame : start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1),
|
|
self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
|
|
self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
|
|
],
|
|
dim=-1,
|
|
).reshape(f * h * w, 1, -1)
|
|
return freqs
|
|
|
|
def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor) -> torch.Tensor:
|
|
"""Reconstruct video from patch embeddings."""
|
|
B = x.shape[0]
|
|
c = self.out_dim
|
|
grid_size = grid_size.tolist()
|
|
expected_seq_len = math.prod(grid_size)
|
|
if x.shape[1] != expected_seq_len:
|
|
raise ValueError(f"x sequence length must equal product(grid_size)={expected_seq_len}, got {x.shape[1]}.")
|
|
x = x.view(B, *grid_size, *self.patch_size, c)
|
|
x = torch.einsum("bfhwpqrc->bcfphqwr", x)
|
|
x = x.reshape(B, c, *[i * j for i, j in zip(grid_size, self.patch_size)])
|
|
return x
|
|
|
|
def _forward_blocks(
|
|
self,
|
|
x: torch.Tensor,
|
|
seq_len: int,
|
|
freqs: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
context: torch.Tensor,
|
|
clip_feature: torch.Tensor | None,
|
|
embodiment_id: torch.Tensor | None,
|
|
action: torch.Tensor | None,
|
|
timestep_action: torch.Tensor | None,
|
|
state: torch.Tensor | None,
|
|
kv_cache: list[Any],
|
|
current_start_frame: int,
|
|
crossattn_cache: list[dict] | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor | None]]:
|
|
x = x.flatten(start_dim=2).transpose(1, 2)
|
|
B = x.shape[0]
|
|
F_t = timestep.shape[1]
|
|
|
|
if action is not None:
|
|
# Current DreamZero checkpoints have one local action/state adapter.
|
|
# Global embodiment IDs are used by transforms and normalization.
|
|
adapter_category_id = torch.zeros(B, dtype=torch.long, device=x.device)
|
|
action_features = self.action_encoder(action, timestep_action, adapter_category_id)
|
|
state_features = self.state_encoder(state, adapter_category_id)
|
|
action_register = torch.cat([action_features, state_features], dim=1)
|
|
action_length = action_features.shape[1]
|
|
action_register_length = action_register.shape[1]
|
|
x = torch.cat([x, action_register], dim=1)
|
|
else:
|
|
action_length = 0
|
|
action_register_length = None
|
|
|
|
timestep = timestep.unsqueeze(-1).expand(B, F_t, seq_len // F_t).reshape(B, -1)
|
|
if action is not None:
|
|
if timestep_action is None or state is None:
|
|
raise RuntimeError("timestep_action and state are required when action is provided.")
|
|
state_features_t = self.state_encoder(state, adapter_category_id)
|
|
stride = timestep_action.shape[1] // state_features_t.shape[1]
|
|
timestep_state = timestep_action[:, ::stride]
|
|
timestep = torch.cat([timestep, timestep_action, timestep_state], dim=1)
|
|
|
|
e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).type_as(x))
|
|
e = e.unflatten(dim=0, sizes=(B, -1))
|
|
e0 = self.time_projection(e)
|
|
e0 = e0.unflatten(dim=2, sizes=(6, self.dim))
|
|
|
|
context = self.text_embedding(context)
|
|
if clip_feature is not None:
|
|
clip_embedding = self.img_emb(clip_feature)
|
|
context = torch.cat([clip_embedding, context], dim=1)
|
|
|
|
# AR-Diffusion paged path: host-side prep ONCE per branch forward, outside
|
|
# the compiled blocks — allocate current video/action slots, build the
|
|
# padded block-table metadata all 40 layers share, and hand the compiled
|
|
# region plain tensors (ARDiffusionPagedLayerInputs) instead of a Python
|
|
# context object. Lazy-allocation contract preserved: only the branch this
|
|
# CFG-parallel rank executes reaches its _forward_blocks.
|
|
if kv_cache and getattr(kv_cache[0], "is_ar_diffusion_paged_context", False):
|
|
if B != 1:
|
|
raise RuntimeError("AR-Diffusion paged self-attention currently expects batch_size=1")
|
|
fctx = kv_cache[0].forward_ctx
|
|
if seq_len != fctx.seq_len:
|
|
raise RuntimeError(
|
|
f"AR-Diffusion paged context seq_len={fctx.seq_len} but current video KV has {seq_len} tokens"
|
|
)
|
|
fctx.prepare(
|
|
device=x.device,
|
|
action_len=int(action_register_length or 0),
|
|
query_len=int(x.shape[1]),
|
|
)
|
|
kv_cache = [c.to_layer_inputs() for c in kv_cache]
|
|
|
|
updated_kv_caches: list[torch.Tensor | None] = []
|
|
for block_index, block in enumerate(self.blocks):
|
|
x, updated_kv_cache = block(
|
|
x=x,
|
|
e=e0,
|
|
freqs=freqs,
|
|
freqs_action=self.freqs_action,
|
|
freqs_state=self.freqs_state,
|
|
context=context,
|
|
action_register_length=action_register_length,
|
|
kv_cache=kv_cache[block_index] if kv_cache else None,
|
|
crossattn_cache=crossattn_cache[block_index] if crossattn_cache else None,
|
|
current_start_frame=current_start_frame,
|
|
)
|
|
updated_kv_caches.append(updated_kv_cache)
|
|
|
|
if action is not None:
|
|
action_noise_pred = x[:, seq_len : seq_len + action_length]
|
|
action_noise_pred = self.action_decoder(action_noise_pred, adapter_category_id)
|
|
else:
|
|
action_noise_pred = None
|
|
|
|
x_video = x[:, :seq_len]
|
|
e_video = e[:, :seq_len]
|
|
x_video = self.head(x_video, e_video.unsqueeze(2))
|
|
|
|
return x_video, action_noise_pred, updated_kv_caches
|
|
|
|
def _forward_inference(
|
|
self,
|
|
x: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
context: torch.Tensor,
|
|
seq_len: int,
|
|
kv_cache: list[Any],
|
|
current_start_frame: int,
|
|
crossattn_cache: list[dict] | None = None,
|
|
y: torch.Tensor | None = None,
|
|
clip_feature: torch.Tensor | None = None,
|
|
action: torch.Tensor | None = None,
|
|
timestep_action: torch.Tensor | None = None,
|
|
state: torch.Tensor | None = None,
|
|
embodiment_id: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
|
|
if self.model_type == "i2v":
|
|
if clip_feature is None or y is None:
|
|
raise RuntimeError("clip_feature and y are required for i2v inference.")
|
|
if context.shape[1] != self.text_len:
|
|
raise ValueError(f"context length must be {self.text_len}, got {context.shape[1]}.")
|
|
|
|
if y is not None:
|
|
x = torch.cat([x, y.to(dtype=x.dtype)], dim=1)
|
|
|
|
x = self.patch_embedding(x)
|
|
grid_size = torch.tensor(x.shape[2:], dtype=torch.long)
|
|
freqs = self._create_freqs(grid_size, current_start_frame)
|
|
|
|
x_video, action_noise_pred, updated_kv_caches = self._forward_blocks(
|
|
x=x,
|
|
seq_len=seq_len,
|
|
freqs=freqs,
|
|
timestep=timestep,
|
|
context=context,
|
|
clip_feature=clip_feature,
|
|
embodiment_id=embodiment_id,
|
|
action=action,
|
|
timestep_action=timestep_action,
|
|
state=state,
|
|
kv_cache=kv_cache,
|
|
crossattn_cache=crossattn_cache,
|
|
current_start_frame=current_start_frame,
|
|
)
|
|
|
|
x_video = x_video.clone()
|
|
if action_noise_pred is not None:
|
|
action_noise_pred = action_noise_pred.clone()
|
|
|
|
video_noise_pred = self.unpatchify(x_video, grid_size)
|
|
return video_noise_pred, action_noise_pred, updated_kv_caches
|
|
|
|
def forward(self, *args: Any, **kwargs: Any):
|
|
"""Inference only. Requires kv_cache."""
|
|
return self._forward_inference(*args, **kwargs)
|