1345 lines
57 KiB
Python
1345 lines
57 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Checkpoint-compatible causal DiT for the LingBot World v2 model package."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from collections.abc import Iterable
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from typing import Any, Self, cast
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
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.model_loader.weight_utils import default_weight_loader
|
|
from vllm.model_executor.utils import set_weight_attrs
|
|
|
|
from vllm_omni.diffusion.attention.layer import Attention
|
|
from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D
|
|
from vllm_omni.diffusion.distributed.parallel_state import get_sp_group
|
|
from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
|
|
from vllm_omni.diffusion.layers.norm import LayerNorm
|
|
from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
|
|
ARDiffusionPagedLayerContext,
|
|
ARDiffusionPagedLayerInputs,
|
|
ar_diffusion_paged_attention,
|
|
paged_write_attn,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class LingBotAttentionCache:
|
|
"""K/V storage plus its logical cursor metadata.
|
|
|
|
``end`` is occupied storage, ``absolute_end`` is the next global token
|
|
position, and ``sink_end`` separates permanently retained prefix tokens
|
|
from the sliding local window. ``last_start`` rejects overlapping or
|
|
out-of-order causal chunks.
|
|
"""
|
|
|
|
key: torch.Tensor
|
|
value: torch.Tensor
|
|
end: int = 0
|
|
absolute_end: int = 0
|
|
last_start: int | None = None
|
|
sink_end: int = 0
|
|
|
|
|
|
@dataclass
|
|
class LingBotTransformerCache:
|
|
"""One request's per-layer video K/V and reusable text K/V."""
|
|
|
|
self_attention: list[LingBotAttentionCache | ARDiffusionPagedLayerContext | ARDiffusionPagedLayerInputs]
|
|
cross_attention: list[LingBotAttentionCache | None]
|
|
|
|
|
|
def allocate_lingbot_cache(
|
|
*,
|
|
batch_size: int,
|
|
num_layers: int,
|
|
max_tokens: int,
|
|
num_local_heads: int,
|
|
head_dim: int,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
) -> LingBotTransformerCache:
|
|
# Cross-attention starts empty because its token count is known after text encoding.
|
|
shape = (batch_size, max_tokens, num_local_heads, head_dim)
|
|
self_attention = [
|
|
LingBotAttentionCache(
|
|
key=torch.zeros(shape, device=device, dtype=dtype),
|
|
value=torch.zeros(shape, device=device, dtype=dtype),
|
|
)
|
|
for _ in range(num_layers)
|
|
]
|
|
return LingBotTransformerCache(
|
|
self_attention=self_attention,
|
|
cross_attention=[None for _ in range(num_layers)],
|
|
)
|
|
|
|
|
|
class _LingBotRMSNorm(nn.Module):
|
|
"""RMSNorm over a tensor-parallel sharded projection."""
|
|
|
|
def __init__(self, hidden_size: int, eps: float) -> None:
|
|
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
|
|
shard_start = get_tensor_model_parallel_rank() * shard_size
|
|
shard = loaded_weight.narrow(0, shard_start, 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 forward(self, value: torch.Tensor) -> torch.Tensor:
|
|
tp_size = get_tensor_model_parallel_world_size()
|
|
value_float = value.float()
|
|
sum_of_squares = value_float.pow(2).sum(dim=-1, keepdim=True)
|
|
element_count = value.shape[-1]
|
|
if tp_size > 1:
|
|
sum_of_squares = tensor_model_parallel_all_reduce(sum_of_squares)
|
|
element_count *= tp_size
|
|
rms = torch.sqrt(sum_of_squares / element_count + self.eps)
|
|
return (value_float / rms * self.weight.float()).to(value.dtype)
|
|
|
|
|
|
def _projection_prefix(prefix: str, name: str) -> str:
|
|
return f"{prefix}.{name}" if prefix else name
|
|
|
|
|
|
def _ulysses_state() -> tuple[int, int, torch.distributed.ProcessGroup | None]:
|
|
coordinator = get_sp_group()
|
|
return (
|
|
int(coordinator.ulysses_world_size),
|
|
int(coordinator.ulysses_rank),
|
|
coordinator.ulysses_group,
|
|
)
|
|
|
|
|
|
class _LingBotSPPrepare(nn.Module):
|
|
"""Expand frame conditioning only when Ulysses shards the token sequence."""
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
camera_hidden_states: torch.Tensor,
|
|
timestep_projection: torch.Tensor,
|
|
rotary_emb: tuple[torch.Tensor, torch.Tensor],
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
num_frames = timestep_projection.shape[1]
|
|
if hidden_states.shape[1] % num_frames:
|
|
raise ValueError("LingBot token count must be divisible by the latent frame count.")
|
|
tokens_per_frame = hidden_states.shape[1] // num_frames
|
|
if get_sp_group().ulysses_world_size > 1:
|
|
timestep_projection = (
|
|
timestep_projection.unsqueeze(2).expand(-1, -1, tokens_per_frame, -1, -1).flatten(1, 2)
|
|
)
|
|
cosine, sine = rotary_emb
|
|
return hidden_states, camera_hidden_states, timestep_projection, cosine, sine
|
|
|
|
|
|
class LingBotSelfAttention(nn.Module):
|
|
"""Block-causal self-attention over retained history and one full chunk."""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_heads: int,
|
|
*,
|
|
eps: float = 1e-6,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
if dim % num_heads != 0:
|
|
raise ValueError(f"dim={dim} must be divisible by num_heads={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.dim = dim
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
self.qkv = QKVParallelLinear(
|
|
hidden_size=dim,
|
|
head_size=self.head_dim,
|
|
total_num_heads=num_heads,
|
|
bias=True,
|
|
prefix=_projection_prefix(prefix, "qkv"),
|
|
)
|
|
self.num_local_heads = self.qkv.num_heads
|
|
self.ulysses_world_size, self.ulysses_rank, self.ulysses_group = _ulysses_state()
|
|
if self.num_local_heads % self.ulysses_world_size:
|
|
raise ValueError(
|
|
"LingBot local attention heads must be divisible by the Ulysses degree: "
|
|
f"heads={self.num_local_heads}, ulysses={self.ulysses_world_size}."
|
|
)
|
|
self.num_sp_heads = self.num_local_heads // self.ulysses_world_size
|
|
self.tp_inner_dim = self.num_local_heads * self.head_dim
|
|
self.o = RowParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=True,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "o"),
|
|
)
|
|
self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps)
|
|
self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps)
|
|
self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True)
|
|
self.attn = Attention(
|
|
num_heads=self.num_sp_heads,
|
|
head_size=self.head_dim,
|
|
num_kv_heads=self.num_sp_heads,
|
|
softmax_scale=self.head_dim**-0.5,
|
|
causal=False,
|
|
role="self",
|
|
qkv_layout="BSND",
|
|
prefix=prefix,
|
|
skip_sequence_parallel=True,
|
|
)
|
|
|
|
def _update_cache(
|
|
self,
|
|
cache: LingBotAttentionCache,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
current_start: int,
|
|
*,
|
|
sink_tokens: int,
|
|
update_cache: bool,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
chunk_tokens = key.shape[1]
|
|
capacity = cache.key.shape[1]
|
|
next_sink_end = cache.sink_end
|
|
|
|
if cache.last_start is None:
|
|
if chunk_tokens > capacity:
|
|
raise ValueError(
|
|
f"Current chunk has {chunk_tokens} tokens but the cache can hold only {capacity}; "
|
|
"the full current chunk must remain visible."
|
|
)
|
|
next_key = key
|
|
next_value = value
|
|
next_sink_end = max(0, min(chunk_tokens, sink_tokens - current_start))
|
|
elif current_start == cache.last_start:
|
|
if current_start + chunk_tokens != cache.absolute_end:
|
|
raise ValueError("A repeated current_start must overwrite the same-size current chunk.")
|
|
if chunk_tokens > cache.end:
|
|
raise ValueError("The current chunk is no longer fully retained in the cache.")
|
|
prefix_end = cache.end - chunk_tokens
|
|
next_key = torch.cat((cache.key[:, :prefix_end], key), dim=1)
|
|
next_value = torch.cat((cache.value[:, :prefix_end], value), dim=1)
|
|
else:
|
|
if current_start < cache.last_start:
|
|
raise ValueError(f"current_start={current_start} precedes the latest chunk start {cache.last_start}.")
|
|
if current_start < cache.absolute_end:
|
|
raise ValueError(
|
|
f"current_start={current_start} overlaps cached tokens ending at {cache.absolute_end}."
|
|
)
|
|
if current_start > cache.absolute_end:
|
|
raise ValueError(
|
|
f"New chunks must be contiguous: current_start={current_start}, expected {cache.absolute_end}."
|
|
)
|
|
|
|
incoming_sink_tokens = max(0, min(chunk_tokens, sink_tokens - current_start))
|
|
old_sink_key = cache.key[:, : cache.sink_end]
|
|
old_sink_value = cache.value[:, : cache.sink_end]
|
|
new_sink_key = key[:, :incoming_sink_tokens]
|
|
new_sink_value = value[:, :incoming_sink_tokens]
|
|
next_sink_end = cache.sink_end + incoming_sink_tokens
|
|
|
|
old_local_key = cache.key[:, cache.sink_end : cache.end]
|
|
old_local_value = cache.value[:, cache.sink_end : cache.end]
|
|
new_local_key = key[:, incoming_sink_tokens:]
|
|
new_local_value = value[:, incoming_sink_tokens:]
|
|
local_capacity = capacity - next_sink_end
|
|
if new_local_key.shape[1] > local_capacity:
|
|
raise ValueError("The configured cache cannot retain all sink tokens and the full current chunk.")
|
|
retained_local_tokens = min(
|
|
old_local_key.shape[1],
|
|
local_capacity - new_local_key.shape[1],
|
|
)
|
|
if retained_local_tokens:
|
|
old_local_key = old_local_key[:, -retained_local_tokens:]
|
|
old_local_value = old_local_value[:, -retained_local_tokens:]
|
|
else:
|
|
old_local_key = old_local_key[:, :0]
|
|
old_local_value = old_local_value[:, :0]
|
|
|
|
next_key = torch.cat((old_sink_key, new_sink_key, old_local_key, new_local_key), dim=1)
|
|
next_value = torch.cat((old_sink_value, new_sink_value, old_local_value, new_local_value), dim=1)
|
|
|
|
next_end = next_key.shape[1]
|
|
if update_cache:
|
|
cache.key[:, :next_end].copy_(next_key)
|
|
cache.value[:, :next_end].copy_(next_value)
|
|
cache.end = next_end
|
|
cache.absolute_end = current_start + chunk_tokens
|
|
cache.last_start = current_start
|
|
cache.sink_end = next_sink_end
|
|
return next_key, next_value
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
*,
|
|
cache: LingBotAttentionCache | ARDiffusionPagedLayerInputs,
|
|
current_start: int,
|
|
rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
sink_tokens: int,
|
|
update_cache: bool = True,
|
|
) -> torch.Tensor:
|
|
# Project logical [B, S, D] tokens into TP-local
|
|
# [B, S, num_local_heads, head_dim].
|
|
qkv, _ = self.qkv(hidden_states)
|
|
query, key, value = qkv.split((self.tp_inner_dim, self.tp_inner_dim, self.tp_inner_dim), dim=-1)
|
|
query = self.norm_q(query)
|
|
key = self.norm_k(key)
|
|
|
|
query = query.unflatten(2, (self.num_local_heads, self.head_dim))
|
|
key = key.unflatten(2, (self.num_local_heads, self.head_dim))
|
|
value = value.unflatten(2, (self.num_local_heads, self.head_dim))
|
|
if rotary_emb is not None:
|
|
cos, sin = rotary_emb
|
|
query = self.rotary_embedding(query, cos, sin)
|
|
key = self.rotary_embedding(key, cos, sin)
|
|
|
|
if self.ulysses_world_size > 1:
|
|
query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False)
|
|
key = SeqAllToAll4D.apply(self.ulysses_group, key, 2, 1, False)
|
|
value = SeqAllToAll4D.apply(self.ulysses_group, value, 2, 1, False)
|
|
|
|
if isinstance(cache, ARDiffusionPagedLayerInputs):
|
|
if query.shape[0] != 1:
|
|
raise RuntimeError("LingBot AR-Diffusion paged attention requires batch_size=1.")
|
|
output = paged_write_attn(
|
|
cache,
|
|
query[0],
|
|
key[0],
|
|
value[0],
|
|
None,
|
|
None,
|
|
self.head_dim**-0.5,
|
|
).unsqueeze(0)
|
|
else:
|
|
# The direct path sees history + current K/V even when the caller
|
|
# asks not to mutate persistent cache state.
|
|
visible_key, visible_value = self._update_cache(
|
|
cache,
|
|
key,
|
|
value,
|
|
current_start,
|
|
sink_tokens=sink_tokens,
|
|
update_cache=update_cache,
|
|
)
|
|
if query.is_cuda and query.shape[0] == 1:
|
|
# Use the same block-table FlashAttention entry point as the
|
|
# realtime path so direct replay is a numerical oracle for
|
|
# paged execution, rather than a comparison between two
|
|
# different attention kernels.
|
|
block_size = key.shape[1]
|
|
key_cache = visible_key[0].unflatten(0, (-1, block_size))
|
|
value_cache = visible_value[0].unflatten(0, (-1, block_size))
|
|
block_count = key_cache.shape[0]
|
|
block_table = torch.arange(
|
|
block_count,
|
|
dtype=torch.int32,
|
|
device=query.device,
|
|
).unsqueeze(0)
|
|
query_start_loc = torch.tensor(
|
|
[0, query.shape[1]],
|
|
dtype=torch.int32,
|
|
device=query.device,
|
|
)
|
|
seq_lens = torch.tensor(
|
|
[visible_key.shape[1]],
|
|
dtype=torch.int32,
|
|
device=query.device,
|
|
)
|
|
output = ar_diffusion_paged_attention(
|
|
query,
|
|
key_cache,
|
|
value_cache,
|
|
block_table=block_table,
|
|
query_start_loc=query_start_loc,
|
|
seq_lens=seq_lens,
|
|
max_query_len=query.shape[1],
|
|
max_seq_len=visible_key.shape[1],
|
|
softmax_scale=self.head_dim**-0.5,
|
|
)
|
|
else:
|
|
output = self.attn(query, visible_key, visible_value)
|
|
if self.ulysses_world_size > 1:
|
|
output = SeqAllToAll4D.apply(self.ulysses_group, output, 1, 2, False)
|
|
return self.o(output.flatten(2, 3))
|
|
|
|
|
|
class LingBotCrossAttention(nn.Module):
|
|
"""Cross-attention with caller-owned request-local encoder K/V reuse."""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_heads: int,
|
|
*,
|
|
eps: float = 1e-6,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
if dim % num_heads != 0:
|
|
raise ValueError(f"dim={dim} must be divisible by num_heads={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.dim = dim
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
self.num_local_heads = num_heads // tp_size
|
|
self.ulysses_world_size, self.ulysses_rank, self.ulysses_group = _ulysses_state()
|
|
if self.num_local_heads % self.ulysses_world_size:
|
|
raise ValueError(
|
|
"LingBot local cross-attention heads must be divisible by the Ulysses degree: "
|
|
f"heads={self.num_local_heads}, ulysses={self.ulysses_world_size}."
|
|
)
|
|
self.num_sp_heads = self.num_local_heads // self.ulysses_world_size
|
|
self.tp_inner_dim = self.num_local_heads * self.head_dim
|
|
|
|
self.q = ColumnParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=True,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "q"),
|
|
)
|
|
self.k = ColumnParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=True,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "k"),
|
|
)
|
|
self.v = ColumnParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=True,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "v"),
|
|
)
|
|
self.o = RowParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=True,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "o"),
|
|
)
|
|
self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps)
|
|
self.norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps)
|
|
self.attn = Attention(
|
|
num_heads=self.num_sp_heads,
|
|
head_size=self.head_dim,
|
|
num_kv_heads=self.num_sp_heads,
|
|
softmax_scale=self.head_dim**-0.5,
|
|
causal=False,
|
|
role="cross",
|
|
qkv_layout="BSND",
|
|
prefix=prefix,
|
|
skip_sequence_parallel=True,
|
|
disable_kv_quant=True,
|
|
)
|
|
|
|
def shard_kv_heads(self, value: torch.Tensor) -> torch.Tensor:
|
|
"""Select this Ulysses rank's projected K/V heads."""
|
|
|
|
if self.ulysses_world_size == 1:
|
|
return value
|
|
start = self.ulysses_rank * self.num_sp_heads
|
|
return value[:, :, start : start + self.num_sp_heads].clone(memory_format=torch.contiguous_format)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | None,
|
|
*,
|
|
cache: LingBotAttentionCache | None,
|
|
) -> tuple[torch.Tensor, LingBotAttentionCache]:
|
|
query = self.norm_q(self.q(hidden_states))
|
|
query = query.unflatten(2, (self.num_local_heads, self.head_dim))
|
|
if self.ulysses_world_size > 1:
|
|
query = SeqAllToAll4D.apply(self.ulysses_group, query, 2, 1, False)
|
|
|
|
# Text K/V is constant within a request and is projected once per layer.
|
|
if cache is None:
|
|
if encoder_hidden_states is None:
|
|
raise ValueError("encoder_hidden_states are required when the cross-attention cache is empty.")
|
|
key = self.norm_k(self.k(encoder_hidden_states))
|
|
value = self.v(encoder_hidden_states)
|
|
key = self.shard_kv_heads(key.unflatten(2, (self.num_local_heads, self.head_dim)))
|
|
value = self.shard_kv_heads(value.unflatten(2, (self.num_local_heads, self.head_dim)))
|
|
cache = LingBotAttentionCache(
|
|
key=key,
|
|
value=value,
|
|
end=key.shape[1],
|
|
absolute_end=key.shape[1],
|
|
last_start=0,
|
|
)
|
|
else:
|
|
key = cache.key[:, : cache.end]
|
|
value = cache.value[:, : cache.end]
|
|
|
|
output = self.attn(query, key, value)
|
|
if self.ulysses_world_size > 1:
|
|
output = SeqAllToAll4D.apply(self.ulysses_group, output, 1, 2, False)
|
|
return self.o(output.flatten(2, 3)), cache
|
|
|
|
|
|
class LingBotAttentionBlock(nn.Module):
|
|
"""Checkpoint-compatible LingBot block with causal video attention."""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_heads: int,
|
|
*,
|
|
ffn_dim: int | None = None,
|
|
cross_attn_norm: bool = True,
|
|
eps: float = 1e-6,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
ffn_dim = ffn_dim or dim * 4
|
|
self.dim = dim
|
|
self.norm1 = LayerNorm(dim, eps=eps, elementwise_affine=False)
|
|
self.self_attn = LingBotSelfAttention(
|
|
dim,
|
|
num_heads,
|
|
eps=eps,
|
|
prefix=_projection_prefix(prefix, "self_attn"),
|
|
)
|
|
self.cross_attn = LingBotCrossAttention(
|
|
dim,
|
|
num_heads,
|
|
eps=eps,
|
|
prefix=_projection_prefix(prefix, "cross_attn"),
|
|
)
|
|
self.norm2 = LayerNorm(dim, eps=eps, elementwise_affine=False)
|
|
self.norm3 = LayerNorm(dim, eps=eps) if cross_attn_norm else nn.Identity()
|
|
self.ffn = nn.Sequential(
|
|
ColumnParallelLinear(
|
|
dim,
|
|
ffn_dim,
|
|
bias=True,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "ffn.0"),
|
|
),
|
|
nn.GELU(approximate="tanh"),
|
|
RowParallelLinear(
|
|
ffn_dim,
|
|
dim,
|
|
bias=True,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "ffn.2"),
|
|
),
|
|
)
|
|
self.modulation = nn.Parameter(torch.randn(1, 6, dim) / math.sqrt(dim))
|
|
self.cam_injector_layer1 = ColumnParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=True,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "cam_injector_layer1"),
|
|
)
|
|
self.cam_injector_layer2 = RowParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=True,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "cam_injector_layer2"),
|
|
)
|
|
self.cam_scale_layer = nn.Linear(dim, dim)
|
|
self.cam_shift_layer = nn.Linear(dim, dim)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | None,
|
|
timestep_projection: torch.Tensor,
|
|
camera_hidden_states: torch.Tensor,
|
|
*,
|
|
self_cache: LingBotAttentionCache | ARDiffusionPagedLayerInputs,
|
|
cross_cache: LingBotAttentionCache | None,
|
|
current_start: int,
|
|
sink_tokens: int,
|
|
update_cache: bool,
|
|
rotary_emb: tuple[torch.Tensor, torch.Tensor],
|
|
) -> tuple[torch.Tensor, LingBotAttentionCache]:
|
|
batch_size, token_count, dim = hidden_states.shape
|
|
if (
|
|
timestep_projection.ndim != 4
|
|
or timestep_projection.shape[0] != batch_size
|
|
or timestep_projection.shape[2:] != (6, dim)
|
|
):
|
|
raise ValueError("LingBot timestep projection must have shape (batch, frames or tokens, 6, dim).")
|
|
num_frames = timestep_projection.shape[1]
|
|
if num_frames == 0 or token_count % num_frames:
|
|
raise ValueError("LingBot timestep projection must align with hidden tokens.")
|
|
if self.self_attn.ulysses_world_size > 1 and num_frames != token_count:
|
|
raise ValueError("LingBot Ulysses requires one timestep projection per local token.")
|
|
# Without SP, broadcast per frame. With SP, each entry represents one token.
|
|
tokens_per_frame = token_count // num_frames
|
|
# Timestep, camera, and text remain separate conditioning paths.
|
|
modulation = self.modulation.unsqueeze(1) + timestep_projection.float()
|
|
shift_msa, scale_msa, gate_msa, shift_ffn, scale_ffn, gate_ffn = modulation.chunk(6, dim=2)
|
|
|
|
hidden_grid = hidden_states.unflatten(1, (num_frames, tokens_per_frame))
|
|
normalized = self.norm1(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame))
|
|
normalized = (normalized * (1 + scale_msa) + shift_msa).flatten(1, 2).to(hidden_states.dtype)
|
|
attention_output = self.self_attn(
|
|
normalized,
|
|
cache=self_cache,
|
|
current_start=current_start,
|
|
rotary_emb=rotary_emb,
|
|
sink_tokens=sink_tokens,
|
|
update_cache=update_cache,
|
|
)
|
|
hidden_grid = hidden_grid + attention_output.unflatten(1, (num_frames, tokens_per_frame)) * gate_msa
|
|
hidden_states = hidden_grid.flatten(1, 2).to(hidden_states.dtype)
|
|
|
|
camera_features = self.cam_injector_layer2(F.silu(self.cam_injector_layer1(camera_hidden_states)))
|
|
camera_features = camera_features + camera_hidden_states
|
|
camera_scale = self.cam_scale_layer(camera_features)
|
|
camera_shift = self.cam_shift_layer(camera_features)
|
|
hidden_states = ((1 + camera_scale) * hidden_states + camera_shift).to(hidden_states.dtype)
|
|
|
|
attention_output, cross_cache = self.cross_attn(
|
|
self.norm3(hidden_states),
|
|
encoder_hidden_states,
|
|
cache=cross_cache,
|
|
)
|
|
hidden_states = hidden_states + attention_output
|
|
|
|
hidden_grid = hidden_states.unflatten(1, (num_frames, tokens_per_frame))
|
|
normalized = self.norm2(hidden_states.float()).unflatten(1, (num_frames, tokens_per_frame))
|
|
normalized = (normalized * (1 + scale_ffn) + shift_ffn).flatten(1, 2).to(hidden_states.dtype)
|
|
ffn_output = self.ffn(normalized).unflatten(1, (num_frames, tokens_per_frame))
|
|
hidden_states = (hidden_grid + ffn_output * gate_ffn).flatten(1, 2).to(hidden_states.dtype)
|
|
return hidden_states, cast(LingBotAttentionCache, cross_cache)
|
|
|
|
|
|
class _LingBotCameraPatchEmbedding(nn.Module):
|
|
"""Linear camera patchifier matching the checkpoint's direct parameter names."""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
dim: int,
|
|
patch_size: tuple[int, int, int],
|
|
) -> None:
|
|
super().__init__()
|
|
self.patch_size = patch_size
|
|
self.in_features = in_channels * math.prod(patch_size)
|
|
self.weight = nn.Parameter(torch.empty(dim, self.in_features))
|
|
self.bias = nn.Parameter(torch.empty(dim))
|
|
nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
|
bound = 1 / math.sqrt(self.in_features)
|
|
nn.init.uniform_(self.bias, -bound, bound)
|
|
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
batch_size, channels, frames, height, width = hidden_states.shape
|
|
patch_frames, patch_height, patch_width = self.patch_size
|
|
# Preserve Conv3d patch order while retaining checkpoint weight/bias names.
|
|
hidden_states = hidden_states.reshape(
|
|
batch_size,
|
|
channels,
|
|
frames // patch_frames,
|
|
patch_frames,
|
|
height // patch_height,
|
|
patch_height,
|
|
width // patch_width,
|
|
patch_width,
|
|
)
|
|
hidden_states = hidden_states.permute(0, 2, 4, 6, 1, 3, 5, 7)
|
|
hidden_states = hidden_states.reshape(batch_size, -1, self.in_features)
|
|
return F.linear(hidden_states, self.weight, self.bias)
|
|
|
|
|
|
class _LingBotHead(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
out_channels: int,
|
|
patch_size: tuple[int, int, int],
|
|
eps: float,
|
|
) -> None:
|
|
super().__init__()
|
|
self.norm = LayerNorm(dim, eps=eps, elementwise_affine=False)
|
|
self.head = nn.Linear(dim, out_channels * math.prod(patch_size))
|
|
self.modulation = nn.Parameter(torch.randn(1, 2, dim) / math.sqrt(dim))
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
timestep_embedding: torch.Tensor,
|
|
*,
|
|
tokens_per_frame: int | None = None,
|
|
token_offset: int = 0,
|
|
) -> torch.Tensor:
|
|
num_frames = timestep_embedding.shape[1]
|
|
modulation = self.modulation.unsqueeze(1) + timestep_embedding.unsqueeze(2).float()
|
|
shift, scale = modulation.chunk(2, dim=2)
|
|
normalized = self.norm(hidden_states.float())
|
|
if tokens_per_frame is None:
|
|
normalized = normalized.unflatten(1, (num_frames, -1))
|
|
normalized = (normalized * (1 + scale) + shift).flatten(1, 2)
|
|
else:
|
|
# A sequence shard can begin/end inside a frame.
|
|
frames = (
|
|
torch.arange(hidden_states.shape[1], device=hidden_states.device) + token_offset
|
|
) // tokens_per_frame
|
|
scale = scale.squeeze(2).index_select(1, frames)
|
|
shift = shift.squeeze(2).index_select(1, frames)
|
|
normalized = normalized * (1 + scale) + shift
|
|
normalized = normalized.to(hidden_states.dtype)
|
|
return self.head(normalized)
|
|
|
|
|
|
def _sinusoidal_embedding(dim: int, timestep: torch.Tensor) -> torch.Tensor:
|
|
if dim % 2:
|
|
raise ValueError(f"freq_dim must be even, got {dim}.")
|
|
half_dim = dim // 2
|
|
timestep = timestep.to(torch.float64)
|
|
frequencies = torch.pow(
|
|
10000,
|
|
-torch.arange(half_dim, device=timestep.device, dtype=torch.float64) / half_dim,
|
|
)
|
|
phase = torch.outer(timestep, frequencies)
|
|
return torch.cat((phase.cos(), phase.sin()), dim=1)
|
|
|
|
|
|
def _rope_axis(max_seq_len: int, dim: int, *, start: int = 0) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if dim == 0:
|
|
empty = torch.empty(max_seq_len, 0, dtype=torch.float32)
|
|
return empty, empty.clone()
|
|
if dim % 2:
|
|
raise ValueError(f"RoPE axis dimension must be even, got {dim}.")
|
|
frequencies = 1.0 / torch.pow(
|
|
10000,
|
|
torch.arange(0, dim, 2, dtype=torch.float64) / dim,
|
|
)
|
|
phase = torch.outer(torch.arange(start, start + max_seq_len, dtype=torch.float64), frequencies)
|
|
return phase.cos().float(), phase.sin().float()
|
|
|
|
|
|
class CausalLingBotWorldTransformer3DModel(nn.Module):
|
|
"""Checkpoint-compatible causal LingBot World video transformer."""
|
|
|
|
_repeated_blocks = ["LingBotAttentionBlock"]
|
|
packed_modules_mapping = {"qkv": ["q", "k", "v"]}
|
|
_layerwise_offload_blocks_attrs = ["blocks"]
|
|
_sp_plan = {
|
|
"sp_prepare": {
|
|
0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True),
|
|
1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True),
|
|
2: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
|
|
3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
|
|
4: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
|
|
},
|
|
"sp_output_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
patch_size: tuple[int, int, int] = (1, 2, 2),
|
|
num_attention_heads: int = 40,
|
|
attention_head_dim: int = 128,
|
|
in_channels: int = 36,
|
|
out_channels: int = 16,
|
|
text_dim: int = 4096,
|
|
freq_dim: int = 256,
|
|
ffn_dim: int = 13824,
|
|
num_layers: int = 40,
|
|
cross_attn_norm: bool = True,
|
|
eps: float = 1e-6,
|
|
image_dim: int | None = None,
|
|
added_kv_proj_dim: int | None = None,
|
|
rope_max_seq_len: int = 1024,
|
|
pos_embed_seq_len: int | None = None,
|
|
qk_norm: str = "rms_norm_across_heads",
|
|
sink_size: int = 9,
|
|
num_frames_per_block: int = 3,
|
|
sliding_window_num_frames: int = 18,
|
|
local_attn_size: int = -1,
|
|
quant_config: object | None = None,
|
|
prefix: str = "",
|
|
) -> None:
|
|
super().__init__()
|
|
if len(patch_size) != 3 or any(size <= 0 for size in patch_size):
|
|
raise ValueError("patch_size must contain three positive integers.")
|
|
if num_attention_heads <= 0 or attention_head_dim <= 0:
|
|
raise ValueError("Attention head counts and dimensions must be positive.")
|
|
if attention_head_dim % 2:
|
|
raise ValueError("attention_head_dim must be even for RoPE.")
|
|
if num_layers <= 0 or ffn_dim <= 0 or freq_dim <= 0:
|
|
raise ValueError("num_layers, ffn_dim, and freq_dim must be positive.")
|
|
if sink_size < 0 or num_frames_per_block <= 0 or sliding_window_num_frames <= 0:
|
|
raise ValueError("Causal frame-window sizes must be positive, except sink_size which may be zero.")
|
|
if local_attn_size != -1 and local_attn_size <= 0:
|
|
raise ValueError("local_attn_size must be -1 or a positive frame count.")
|
|
cache_window_frames = local_attn_size if local_attn_size != -1 else sliding_window_num_frames
|
|
if cache_window_frames < sink_size + num_frames_per_block:
|
|
raise ValueError("The causal cache window must fit all sink frames and one complete current block.")
|
|
if qk_norm != "rms_norm_across_heads":
|
|
raise ValueError(
|
|
f"qk_norm must be 'rms_norm_across_heads' for the LingBot World v2 checkpoint, got {qk_norm!r}."
|
|
)
|
|
for field_name, field_value in (
|
|
("image_dim", image_dim),
|
|
("added_kv_proj_dim", added_kv_proj_dim),
|
|
("pos_embed_seq_len", pos_embed_seq_len),
|
|
):
|
|
if field_value is not None:
|
|
raise ValueError(f"{field_name} must be None because LingBot World v2 has no image embedding path.")
|
|
if quant_config is not None:
|
|
raise RuntimeError(
|
|
"quant_config is not supported by the LingBot World transformer; construct the unquantized model."
|
|
)
|
|
|
|
dim = num_attention_heads * attention_head_dim
|
|
self.dim = dim
|
|
self.config = SimpleNamespace(
|
|
patch_size=patch_size,
|
|
num_attention_heads=num_attention_heads,
|
|
attention_head_dim=attention_head_dim,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
text_dim=text_dim,
|
|
freq_dim=freq_dim,
|
|
ffn_dim=ffn_dim,
|
|
num_layers=num_layers,
|
|
cross_attn_norm=cross_attn_norm,
|
|
eps=eps,
|
|
image_dim=image_dim,
|
|
added_kv_proj_dim=added_kv_proj_dim,
|
|
rope_max_seq_len=rope_max_seq_len,
|
|
pos_embed_seq_len=pos_embed_seq_len,
|
|
qk_norm=qk_norm,
|
|
sink_size=sink_size,
|
|
num_frames_per_block=num_frames_per_block,
|
|
sliding_window_num_frames=sliding_window_num_frames,
|
|
local_attn_size=local_attn_size,
|
|
)
|
|
|
|
self.patch_embedding = Conv3dLayer(
|
|
in_channels=in_channels,
|
|
out_channels=dim,
|
|
kernel_size=patch_size,
|
|
stride=patch_size,
|
|
)
|
|
self.patch_embedding_wancamctrl = _LingBotCameraPatchEmbedding(6 * 8 * 8, dim, patch_size)
|
|
self.sp_prepare = _LingBotSPPrepare()
|
|
self.sp_output_gather = nn.Identity()
|
|
self.c2ws_hidden_states_layer1 = ColumnParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=True,
|
|
gather_output=False,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "c2ws_hidden_states_layer1"),
|
|
)
|
|
self.c2ws_hidden_states_layer2 = RowParallelLinear(
|
|
dim,
|
|
dim,
|
|
bias=True,
|
|
input_is_parallel=True,
|
|
return_bias=False,
|
|
prefix=_projection_prefix(prefix, "c2ws_hidden_states_layer2"),
|
|
)
|
|
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),
|
|
)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
LingBotAttentionBlock(
|
|
dim,
|
|
num_attention_heads,
|
|
ffn_dim=ffn_dim,
|
|
cross_attn_norm=cross_attn_norm,
|
|
eps=eps,
|
|
prefix=_projection_prefix(prefix, f"blocks.{index}"),
|
|
)
|
|
for index in range(num_layers)
|
|
]
|
|
)
|
|
self.head = _LingBotHead(dim, out_channels, patch_size, eps)
|
|
|
|
temporal_dim = attention_head_dim - 4 * (attention_head_dim // 6)
|
|
height_dim = width_dim = 2 * (attention_head_dim // 6)
|
|
for axis, axis_dim in (("temporal", temporal_dim), ("height", height_dim), ("width", width_dim)):
|
|
cosine, sine = _rope_axis(rope_max_seq_len, axis_dim)
|
|
self.register_buffer(f"_rope_{axis}_cosine", cosine, persistent=False)
|
|
self.register_buffer(f"_rope_{axis}_sine", sine, persistent=False)
|
|
|
|
@property
|
|
def dtype(self) -> torch.dtype:
|
|
"""Return the dtype used by the transformer parameters."""
|
|
|
|
return next(self.parameters()).dtype
|
|
|
|
@classmethod
|
|
def from_config(
|
|
cls,
|
|
config: dict[str, Any],
|
|
*,
|
|
quant_config: Any | None = None,
|
|
prefix: str = "",
|
|
) -> Self:
|
|
checkpoint_contract = {
|
|
"patch_size": (1, 2, 2),
|
|
"num_attention_heads": 40,
|
|
"attention_head_dim": 128,
|
|
"in_channels": 36,
|
|
"out_channels": 16,
|
|
"text_dim": 4096,
|
|
"freq_dim": 256,
|
|
"ffn_dim": 13824,
|
|
"num_layers": 40,
|
|
"cross_attn_norm": True,
|
|
"eps": 1e-6,
|
|
"image_dim": None,
|
|
"added_kv_proj_dim": None,
|
|
"rope_max_seq_len": 1024,
|
|
"pos_embed_seq_len": None,
|
|
"qk_norm": "rms_norm_across_heads",
|
|
"sink_size": 9,
|
|
"num_frames_per_block": 3,
|
|
"sliding_window_num_frames": 18,
|
|
"local_attn_size": -1,
|
|
}
|
|
class_name = config.get("_class_name")
|
|
if class_name is not None and class_name != "CausalLingBotWorldTransformer3DModel":
|
|
raise ValueError(
|
|
"_class_name must be 'CausalLingBotWorldTransformer3DModel' for the LingBot World v2 checkpoint."
|
|
)
|
|
kwargs = {name: value for name, value in config.items() if name in checkpoint_contract}
|
|
if "patch_size" in kwargs:
|
|
kwargs["patch_size"] = tuple(kwargs["patch_size"])
|
|
for name, expected in checkpoint_contract.items():
|
|
actual = kwargs.get(name, expected)
|
|
if actual != expected:
|
|
raise ValueError(f"LingBot checkpoint config {name} must be {expected!r}, got {actual!r}.")
|
|
return cls(**kwargs, quant_config=quant_config, prefix=prefix)
|
|
|
|
def allocate_cache(
|
|
self,
|
|
*,
|
|
batch_size: int,
|
|
latent_height: int,
|
|
latent_width: int,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
) -> LingBotTransformerCache:
|
|
"""Allocate one caller-owned cache using this transformer's geometry."""
|
|
|
|
if isinstance(batch_size, bool) or not isinstance(batch_size, int) or batch_size <= 0:
|
|
raise ValueError("batch_size must be a positive integer for LingBot cache allocation.")
|
|
if (
|
|
isinstance(latent_height, bool)
|
|
or not isinstance(latent_height, int)
|
|
or latent_height <= 0
|
|
or isinstance(latent_width, bool)
|
|
or not isinstance(latent_width, int)
|
|
or latent_width <= 0
|
|
):
|
|
raise ValueError("latent_height and latent_width must be positive integers.")
|
|
|
|
# Cache capacity is measured in post-patch tokens, not raw/latent
|
|
# pixels: window_frames * (latent_height/patch_h) * (latent_width/patch_w).
|
|
patch_frames, patch_height, patch_width = self.config.patch_size
|
|
if patch_frames != 1 or latent_height % patch_height or latent_width % patch_width:
|
|
raise ValueError("latent height/width must align with the configured LingBot patch size.")
|
|
post_patch_height = latent_height // patch_height
|
|
post_patch_width = latent_width // patch_width
|
|
window_frames = (
|
|
self.config.local_attn_size if self.config.local_attn_size != -1 else self.config.sliding_window_num_frames
|
|
)
|
|
max_tokens = int(window_frames * post_patch_height * post_patch_width)
|
|
num_local_heads = self.blocks[0].self_attn.num_sp_heads
|
|
return allocate_lingbot_cache(
|
|
batch_size=batch_size,
|
|
num_layers=self.config.num_layers,
|
|
max_tokens=max_tokens,
|
|
num_local_heads=num_local_heads,
|
|
head_dim=self.config.attention_head_dim,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
|
|
def _rotary_embedding(
|
|
self,
|
|
*,
|
|
frames: int,
|
|
height: int,
|
|
width: int,
|
|
start_frame: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if height > self.config.rope_max_seq_len or width > self.config.rope_max_seq_len:
|
|
raise ValueError("Spatial RoPE positions exceed rope_max_seq_len.")
|
|
if start_frame + frames <= self._rope_temporal_cosine.shape[0]:
|
|
temporal_cosine = self._rope_temporal_cosine[start_frame : start_frame + frames]
|
|
temporal_sine = self._rope_temporal_sine[start_frame : start_frame + frames]
|
|
else:
|
|
# Keep absolute positions and checkpoint frequencies, computing only
|
|
# this block instead of growing a table for the whole session.
|
|
temporal_dim = self._rope_temporal_cosine.shape[1] * 2
|
|
temporal_cosine, temporal_sine = _rope_axis(frames, temporal_dim, start=start_frame)
|
|
temporal_cosine = temporal_cosine.to(self._rope_temporal_cosine)
|
|
temporal_sine = temporal_sine.to(self._rope_temporal_sine)
|
|
|
|
def expand_axis(table: torch.Tensor, axis: str) -> torch.Tensor:
|
|
if axis == "temporal":
|
|
return table.view(frames, 1, 1, -1).expand(frames, height, width, -1)
|
|
if axis == "height":
|
|
return table[:height].view(1, height, 1, -1).expand(frames, height, width, -1)
|
|
return table[:width].view(1, 1, width, -1).expand(frames, height, width, -1)
|
|
|
|
cosine = torch.cat(
|
|
(
|
|
expand_axis(temporal_cosine, "temporal"),
|
|
expand_axis(self._rope_height_cosine, "height"),
|
|
expand_axis(self._rope_width_cosine, "width"),
|
|
),
|
|
dim=-1,
|
|
)
|
|
sine = torch.cat(
|
|
(
|
|
expand_axis(temporal_sine, "temporal"),
|
|
expand_axis(self._rope_height_sine, "height"),
|
|
expand_axis(self._rope_width_sine, "width"),
|
|
),
|
|
dim=-1,
|
|
)
|
|
return (
|
|
cosine.reshape(frames * height * width, -1).to(device=device, dtype=dtype),
|
|
sine.reshape(frames * height * width, -1).to(device=device, dtype=dtype),
|
|
)
|
|
|
|
def _validate_inputs(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
camera_hidden_states: torch.Tensor,
|
|
cache: LingBotTransformerCache,
|
|
start_frame: int,
|
|
) -> tuple[int, int, int, int, int]:
|
|
if hidden_states.ndim != 5:
|
|
raise ValueError("hidden_states must be rank 5 [batch, channels, frames, height, width].")
|
|
if hidden_states.shape[1] != self.config.in_channels:
|
|
raise ValueError(f"hidden_states must have {self.config.in_channels} channels.")
|
|
if not isinstance(start_frame, int) or isinstance(start_frame, bool) or start_frame < 0:
|
|
raise ValueError("start_frame must be a non-negative integer latent-frame offset.")
|
|
patch_frames, patch_height, patch_width = self.config.patch_size
|
|
batch_size, _, frames, height, width = hidden_states.shape
|
|
if frames % patch_frames or height % patch_height or width % patch_width:
|
|
raise ValueError("hidden_states dimensions must be divisible by patch_size.")
|
|
post_patch_frames = frames // patch_frames
|
|
if post_patch_frames != self.config.num_frames_per_block:
|
|
raise ValueError(
|
|
f"Each LingBot forward must contain exactly {self.config.num_frames_per_block} post-patch frames, "
|
|
f"got {post_patch_frames}."
|
|
)
|
|
if start_frame % patch_frames:
|
|
raise ValueError("start_frame must align to the temporal patch size.")
|
|
if encoder_hidden_states.ndim != 3:
|
|
raise ValueError("encoder_hidden_states must be rank 3 [batch, tokens, text_dim].")
|
|
if encoder_hidden_states.shape[0] != batch_size or encoder_hidden_states.shape[2] != self.config.text_dim:
|
|
raise ValueError("encoder_hidden_states batch/width do not match the LingBot config.")
|
|
if encoder_hidden_states.shape[1] == 0:
|
|
raise ValueError("encoder_hidden_states must contain at least one token.")
|
|
expected_camera_shape = (batch_size, 6 * 8 * 8, frames, height, width)
|
|
if camera_hidden_states.shape != expected_camera_shape:
|
|
raise ValueError(
|
|
"camera_hidden_states must use folded [batch, 6*8*8, frames, height, width] layout "
|
|
f"matching the video; expected {expected_camera_shape}, got {tuple(camera_hidden_states.shape)}."
|
|
)
|
|
if hidden_states.device != encoder_hidden_states.device or hidden_states.device != camera_hidden_states.device:
|
|
raise ValueError("Video, text, and camera tensors must be on the same device.")
|
|
if hidden_states.dtype != encoder_hidden_states.dtype or hidden_states.dtype != camera_hidden_states.dtype:
|
|
raise ValueError("Video, text, and camera tensors must use the same dtype.")
|
|
if timestep.device != hidden_states.device:
|
|
raise ValueError("timestep must be on the same device as hidden_states.")
|
|
if len(cache.self_attention) != self.config.num_layers or len(cache.cross_attention) != self.config.num_layers:
|
|
raise ValueError("LingBotTransformerCache layer counts must match num_layers.")
|
|
return batch_size, frames, height, width, patch_frames
|
|
|
|
def _timestep_embeddings(
|
|
self,
|
|
timestep: torch.Tensor,
|
|
*,
|
|
batch_size: int,
|
|
frames: int,
|
|
dtype: torch.dtype,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if timestep.ndim == 0:
|
|
if batch_size != 1:
|
|
raise ValueError("A scalar timestep is only valid for batch_size=1.")
|
|
timestep = timestep.reshape(1)
|
|
if timestep.ndim == 1:
|
|
if timestep.shape[0] != batch_size:
|
|
raise ValueError("A rank-1 timestep must contain one value per batch item.")
|
|
timestep = timestep.unsqueeze(1).expand(batch_size, frames)
|
|
elif timestep.ndim == 2:
|
|
if timestep.shape != (batch_size, frames):
|
|
raise ValueError(f"A rank-2 timestep must have shape {(batch_size, frames)}.")
|
|
else:
|
|
raise ValueError("timestep must be scalar, rank 1 [batch], or rank 2 [batch, frames].")
|
|
|
|
frequency_embedding = _sinusoidal_embedding(self.config.freq_dim, timestep.reshape(-1)).to(dtype=dtype)
|
|
timestep_embedding = self.time_embedding(frequency_embedding).unflatten(0, (batch_size, frames))
|
|
timestep_projection = self.time_projection(timestep_embedding).unflatten(2, (6, self.dim))
|
|
return timestep_embedding, timestep_projection
|
|
|
|
def _unpatchify(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
*,
|
|
batch_size: int,
|
|
frames: int,
|
|
height: int,
|
|
width: int,
|
|
) -> torch.Tensor:
|
|
patch_frames, patch_height, patch_width = self.config.patch_size
|
|
hidden_states = hidden_states.reshape(
|
|
batch_size,
|
|
frames,
|
|
height,
|
|
width,
|
|
patch_frames,
|
|
patch_height,
|
|
patch_width,
|
|
self.config.out_channels,
|
|
)
|
|
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
|
|
return hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor,
|
|
camera_hidden_states: torch.Tensor,
|
|
*,
|
|
cache: LingBotTransformerCache,
|
|
start_frame: int,
|
|
update_cache: bool,
|
|
) -> torch.Tensor:
|
|
if not isinstance(update_cache, bool):
|
|
raise ValueError("update_cache must be a boolean.")
|
|
batch_size, frames, height, width, patch_frames = self._validate_inputs(
|
|
hidden_states,
|
|
timestep,
|
|
encoder_hidden_states,
|
|
camera_hidden_states,
|
|
cache,
|
|
start_frame,
|
|
)
|
|
patch_height, patch_width = self.config.patch_size[1:]
|
|
patched_frames = frames // patch_frames
|
|
patched_height = height // patch_height
|
|
patched_width = width // patch_width
|
|
tokens_per_frame = patched_height * patched_width
|
|
patched_start_frame = start_frame // patch_frames
|
|
current_start = patched_start_frame * tokens_per_frame
|
|
sink_tokens = self.config.sink_size * tokens_per_frame
|
|
# Phase 1: create absolute 3D positions for this causal block. The
|
|
# spatial coordinates restart for every frame; temporal coordinates
|
|
# begin at ``start_frame`` so cached blocks retain their true ordering.
|
|
rotary_emb = self._rotary_embedding(
|
|
frames=patched_frames,
|
|
height=patched_height,
|
|
width=patched_width,
|
|
start_frame=patched_start_frame,
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
)
|
|
# Phase 2: independently patchify video and camera grids to identical
|
|
# token layouts, then project timestep and text conditions.
|
|
hidden_states = self.patch_embedding(hidden_states).flatten(2).transpose(1, 2)
|
|
camera_hidden_states = self.patch_embedding_wancamctrl(camera_hidden_states)
|
|
camera_hidden_states = camera_hidden_states + self.c2ws_hidden_states_layer2(
|
|
F.silu(self.c2ws_hidden_states_layer1(camera_hidden_states))
|
|
)
|
|
if hidden_states.shape != camera_hidden_states.shape:
|
|
raise ValueError("Patched video and camera token shapes must match.")
|
|
|
|
timestep_embedding, timestep_projection = self._timestep_embeddings(
|
|
timestep,
|
|
batch_size=batch_size,
|
|
frames=patched_frames,
|
|
dtype=hidden_states.dtype,
|
|
)
|
|
projected_text = (
|
|
self.text_embedding(encoder_hidden_states)
|
|
if any(value is None for value in cache.cross_attention)
|
|
else None
|
|
)
|
|
if cache.self_attention and isinstance(cache.self_attention[0], ARDiffusionPagedLayerContext):
|
|
if batch_size != 1:
|
|
raise RuntimeError("LingBot AR-Diffusion paged attention requires batch_size=1.")
|
|
forward_context = cache.self_attention[0].forward_ctx
|
|
expected_seq_len = patched_frames * tokens_per_frame
|
|
if forward_context.seq_len != expected_seq_len:
|
|
raise RuntimeError(
|
|
"LingBot AR-Diffusion paged context token count does not "
|
|
f"match this block: {forward_context.seq_len} != "
|
|
f"{expected_seq_len}."
|
|
)
|
|
forward_context.prepare(
|
|
device=hidden_states.device,
|
|
action_len=0,
|
|
query_len=hidden_states.shape[1],
|
|
)
|
|
cache.self_attention = [layer_context.to_layer_inputs() for layer_context in cache.self_attention]
|
|
sp_size = self.blocks[0].self_attn.ulysses_world_size
|
|
global_tokens = patched_frames * tokens_per_frame
|
|
hidden_dim = hidden_states.shape[-1]
|
|
rope_dim = rotary_emb[0].shape[-1]
|
|
if global_tokens % sp_size:
|
|
raise ValueError("LingBot Ulysses requires the token count to be divisible by its degree.")
|
|
hidden_states, camera_hidden_states, timestep_projection, cosine, sine = self.sp_prepare(
|
|
hidden_states,
|
|
camera_hidden_states,
|
|
timestep_projection,
|
|
rotary_emb,
|
|
)
|
|
if sp_size > 1:
|
|
local_tokens = global_tokens // sp_size
|
|
expected_hidden = (batch_size, local_tokens, hidden_dim)
|
|
if (
|
|
hidden_states.shape != expected_hidden
|
|
or camera_hidden_states.shape != expected_hidden
|
|
or timestep_projection.shape != (batch_size, local_tokens, 6, hidden_dim)
|
|
or cosine.shape != (local_tokens, rope_dim)
|
|
or sine.shape != (local_tokens, rope_dim)
|
|
):
|
|
raise RuntimeError(
|
|
"LingBot SP input hooks did not shard all conditioning tensors consistently; "
|
|
f"expected {local_tokens} tokens per rank. Check SP hook registration and input dimensions."
|
|
)
|
|
rotary_emb = (cosine, sine)
|
|
# Phase 3: each layer receives its own cache entry. Text K/V is passed
|
|
# only when absent; the returned cache is stored for subsequent DMD
|
|
# steps and causal blocks in this request.
|
|
for index, block in enumerate(self.blocks):
|
|
hidden_states, cross_cache = block(
|
|
hidden_states,
|
|
projected_text if cache.cross_attention[index] is None else None,
|
|
timestep_projection,
|
|
camera_hidden_states,
|
|
self_cache=cache.self_attention[index],
|
|
cross_cache=cache.cross_attention[index],
|
|
current_start=current_start,
|
|
sink_tokens=sink_tokens,
|
|
update_cache=update_cache,
|
|
rotary_emb=rotary_emb,
|
|
)
|
|
cache.cross_attention[index] = cross_cache
|
|
|
|
# Phase 4: project local tokens before gathering the much narrower flow values.
|
|
hidden_states = self.head(
|
|
hidden_states,
|
|
timestep_embedding,
|
|
tokens_per_frame=tokens_per_frame if sp_size > 1 else None,
|
|
token_offset=get_sp_group().ulysses_rank * (global_tokens // sp_size) if sp_size > 1 else 0,
|
|
)
|
|
hidden_states = self.sp_output_gather(hidden_states)
|
|
projected_dim = self.config.out_channels * math.prod(self.config.patch_size)
|
|
if sp_size > 1 and hidden_states.shape != (batch_size, global_tokens, projected_dim):
|
|
raise RuntimeError("LingBot SP output hook did not gather the full projected token sequence.")
|
|
return self._unpatchify(
|
|
hidden_states,
|
|
batch_size=batch_size,
|
|
frames=patched_frames,
|
|
height=patched_height,
|
|
width=patched_width,
|
|
)
|
|
|
|
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
"""Load exact checkpoint names, delegating TP sharding to parameters."""
|
|
|
|
params = dict(self.named_parameters())
|
|
loaded: set[str] = set()
|
|
qkv_mapping = (("q", "q"), ("k", "k"), ("v", "v"))
|
|
for checkpoint_name, loaded_weight in weights:
|
|
name = checkpoint_name
|
|
shard_id = None
|
|
for projection_name, projection_shard in qkv_mapping:
|
|
marker = f".self_attn.{projection_name}."
|
|
if marker in name:
|
|
name = name.replace(marker, ".self_attn.qkv.")
|
|
shard_id = projection_shard
|
|
break
|
|
if name not in params:
|
|
raise KeyError(f"Unexpected LingBot model weight name: {checkpoint_name}")
|
|
param = params[name]
|
|
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
|
if shard_id is None:
|
|
weight_loader(param, loaded_weight)
|
|
else:
|
|
if weight_loader is default_weight_loader:
|
|
raise RuntimeError(f"Fused LingBot QKV parameter {name} has no stacked weight loader.")
|
|
weight_loader(param, loaded_weight, shard_id)
|
|
# DiffusersPipelineLoader compares this return value with the
|
|
# model parameter namespace. Packed Q/K/V checkpoint tensors
|
|
# therefore need to report both their source names and the fused
|
|
# model parameter they populated, matching Wan's loader contract.
|
|
loaded.add(checkpoint_name)
|
|
loaded.add(name)
|
|
return loaded
|