Skip to content

vllm_omni.diffusion.models.lingbot_world.transformer

Checkpoint-compatible causal DiT for the LingBot World v2 model package.

CameraModulationCache

Per-block camera scale and shift, held by the caller across the forwards of one AR block.

The camera injector is four projections and a SiLU per block and reads only the camera tokens, which do not change between the denoise steps and the clean commit of one block. The caller that owns that invariant creates one of these per block and passes it to every forward of the block; forward fills a block's entry on first use and reads it after. Its lifetime is the caller's: the AR runner keeps one per block, generate_block one per loop, and nothing here checks the camera tensor for change.

entries instance-attribute 

entries: dict[int, tuple[Tensor, Tensor]] = {}

CausalLingBotWorldTransformer3DModel

Bases: Module

Checkpoint-compatible causal LingBot World video transformer.

blocks instance-attribute 

blocks = nn.ModuleList(
    [
        LingBotAttentionBlock(
            dim,
            num_attention_heads,
            ffn_dim=ffn_dim,
            cross_attn_norm=cross_attn_norm,
            eps=eps,
            quant_config=quant_config,
            prefix=_projection_prefix(
                prefix, f"blocks.{index}"
            ),
        )
        for index in range(num_layers)
    ]
)

c2ws_hidden_states_layer1 instance-attribute 

c2ws_hidden_states_layer1 = ColumnParallelLinear(
    dim,
    dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=_projection_prefix(
        prefix, "c2ws_hidden_states_layer1"
    ),
)

c2ws_hidden_states_layer2 instance-attribute 

c2ws_hidden_states_layer2 = RowParallelLinear(
    dim,
    dim,
    bias=True,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=_projection_prefix(
        prefix, "c2ws_hidden_states_layer2"
    ),
)

config instance-attribute 

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,
)

dim instance-attribute 

dim = dim

dtype property 

dtype: dtype

Return the dtype used by the transformer parameters.

head instance-attribute 

head = _LingBotHead(dim, out_channels, patch_size, eps)

packed_modules_mapping class-attribute instance-attribute 

packed_modules_mapping = {'qkv': ['q', 'k', 'v']}

patch_embedding instance-attribute 

patch_embedding = Conv3dLayer(
    in_channels=in_channels,
    out_channels=dim,
    kernel_size=patch_size,
    stride=patch_size,
)

patch_embedding_wancamctrl instance-attribute 

patch_embedding_wancamctrl = _LingBotCameraPatchEmbedding(
    6 * 8 * 8, dim, patch_size
)

sp_output_gather instance-attribute 

sp_output_gather = nn.Identity()

sp_prepare instance-attribute 

sp_prepare = _LingBotSPPrepare()

text_embedding instance-attribute 

text_embedding = nn.Sequential(
    nn.Linear(text_dim, dim),
    nn.GELU(approximate="tanh"),
    nn.Linear(dim, dim),
)

time_embedding instance-attribute 

time_embedding = nn.Sequential(
    nn.Linear(freq_dim, dim),
    nn.SiLU(),
    nn.Linear(dim, dim),
)

time_projection instance-attribute 

time_projection = nn.Sequential(
    nn.SiLU(), nn.Linear(dim, dim * 6)
)

allocate_cache 

allocate_cache(
    *,
    batch_size: int,
    latent_height: int,
    latent_width: int,
    device: device,
    dtype: dtype,
) -> LingBotTransformerCache

Allocate one caller-owned cache using this transformer's geometry.

forward 

forward(
    hidden_states: Tensor,
    timestep: Tensor,
    encoder_hidden_states: Tensor,
    camera_hidden_states: Tensor,
    *,
    cache: LingBotTransformerCache,
    start_frame: int,
    update_cache: bool,
    camera_modulation_cache: CameraModulationCache
    | None = None,
) -> Tensor

from_config classmethod 

from_config(
    config: dict[str, Any],
    *,
    quant_config: QuantizationConfig | None = None,
    prefix: str = "",
) -> Self

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

Load exact checkpoint names, delegating TP sharding to parameters.

LingBotAttentionBlock

Bases: Module

Checkpoint-compatible LingBot block with causal video attention.

cam_injector_layer1 instance-attribute 

cam_injector_layer1 = ColumnParallelLinear(
    dim,
    dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=_projection_prefix(
        prefix, "cam_injector_layer1"
    ),
)

cam_injector_layer2 instance-attribute 

cam_injector_layer2 = RowParallelLinear(
    dim,
    dim,
    bias=True,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=_projection_prefix(
        prefix, "cam_injector_layer2"
    ),
)

cam_scale_layer instance-attribute 

cam_scale_layer = nn.Linear(dim, dim)

cam_shift_layer instance-attribute 

cam_shift_layer = nn.Linear(dim, dim)

cross_attn instance-attribute 

cross_attn = LingBotCrossAttention(
    dim,
    num_heads,
    eps=eps,
    quant_config=quant_config,
    prefix=_projection_prefix(prefix, "cross_attn"),
)

dim instance-attribute 

dim = dim

ffn instance-attribute 

ffn = nn.Sequential(
    ColumnParallelLinear(
        dim,
        ffn_dim,
        bias=True,
        gather_output=False,
        return_bias=False,
        quant_config=quant_config,
        prefix=_projection_prefix(prefix, "ffn.0"),
    ),
    nn.GELU(approximate="tanh"),
    RowParallelLinear(
        ffn_dim,
        dim,
        bias=True,
        input_is_parallel=True,
        return_bias=False,
        quant_config=quant_config,
        prefix=_projection_prefix(prefix, "ffn.2"),
    ),
)

modulation instance-attribute 

modulation = nn.Parameter(
    torch.randn(1, 6, dim) / math.sqrt(dim)
)

norm1 instance-attribute 

norm1 = LayerNorm(dim, eps=eps, elementwise_affine=False)

norm2 instance-attribute 

norm2 = LayerNorm(dim, eps=eps, elementwise_affine=False)

norm3 instance-attribute 

norm3 = (
    LayerNorm(dim, eps=eps)
    if cross_attn_norm
    else nn.Identity()
)

self_attn instance-attribute 

self_attn = LingBotSelfAttention(
    dim,
    num_heads,
    eps=eps,
    quant_config=quant_config,
    prefix=_projection_prefix(prefix, "self_attn"),
)

camera_modulation 

camera_modulation(
    camera_hidden_states: Tensor,
) -> tuple[Tensor, Tensor]

Build this block's camera scale and shift from the patchified camera tokens.

Depends only on camera_hidden_states, so a caller holding the camera trajectory fixed -- every denoise forward of one AR chunk does -- can build this once and pass it back through forward.

forward 

forward(
    hidden_states: Tensor,
    encoder_hidden_states: Tensor | None,
    timestep_projection: Tensor,
    camera_hidden_states: Tensor,
    *,
    self_cache: LingBotAttentionCache
    | ARDiffusionPagedLayerInputs,
    cross_cache: LingBotAttentionCache | None,
    current_start: int,
    sink_tokens: int,
    update_cache: bool,
    rotary_emb: tuple[Tensor, Tensor],
    camera_modulation: tuple[Tensor, Tensor] | None = None,
) -> tuple[Tensor, LingBotAttentionCache]

LingBotAttentionCache dataclass

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.

absolute_end class-attribute instance-attribute 

absolute_end: int = 0

end class-attribute instance-attribute 

end: int = 0

key instance-attribute 

key: Tensor

last_start class-attribute instance-attribute 

last_start: int | None = None

sink_end class-attribute instance-attribute 

sink_end: int = 0

value instance-attribute 

value: Tensor

LingBotCrossAttention

Bases: Module

Cross-attention with caller-owned request-local encoder K/V reuse.

attn instance-attribute 

attn = Attention(
    num_heads=self.num_local_heads,
    head_size=self.head_dim,
    num_kv_heads=self.num_local_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,
)

dim instance-attribute 

dim = dim

head_dim instance-attribute 

head_dim = dim // num_heads

k instance-attribute 

k = ColumnParallelLinear(
    dim,
    dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=_projection_prefix(prefix, "k"),
)

norm_k instance-attribute 

norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps)

norm_q instance-attribute 

norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps)

num_heads instance-attribute 

num_heads = num_heads

num_local_heads instance-attribute 

num_local_heads = num_heads // tp_size

o instance-attribute 

o = RowParallelLinear(
    dim,
    dim,
    bias=True,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=_projection_prefix(prefix, "o"),
)

q instance-attribute 

q = ColumnParallelLinear(
    dim,
    dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=_projection_prefix(prefix, "q"),
)

tp_inner_dim instance-attribute 

tp_inner_dim = self.num_local_heads * self.head_dim

v instance-attribute 

v = ColumnParallelLinear(
    dim,
    dim,
    bias=True,
    gather_output=False,
    return_bias=False,
    quant_config=quant_config,
    prefix=_projection_prefix(prefix, "v"),
)

forward 

forward(
    hidden_states: Tensor,
    encoder_hidden_states: Tensor | None,
    *,
    cache: LingBotAttentionCache | None,
) -> tuple[Tensor, LingBotAttentionCache]

LingBotSelfAttention

Bases: Module

Block-causal self-attention over retained history and one full chunk.

attn instance-attribute 

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,
)

dim instance-attribute 

dim = dim

head_dim instance-attribute 

head_dim = dim // num_heads

norm_k instance-attribute 

norm_k = _LingBotRMSNorm(self.tp_inner_dim, eps)

norm_q instance-attribute 

norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps)

num_heads instance-attribute 

num_heads = num_heads

num_local_heads instance-attribute 

num_local_heads = self.qkv.num_heads

num_sp_heads instance-attribute 

num_sp_heads = (
    self.num_local_heads // self.ulysses_world_size
)

o instance-attribute 

o = RowParallelLinear(
    dim,
    dim,
    bias=True,
    input_is_parallel=True,
    return_bias=False,
    quant_config=quant_config,
    prefix=_projection_prefix(prefix, "o"),
)

qkv instance-attribute 

qkv = QKVParallelLinear(
    hidden_size=dim,
    head_size=self.head_dim,
    total_num_heads=num_heads,
    bias=True,
    quant_config=quant_config,
    prefix=_projection_prefix(prefix, "qkv"),
)

rotary_embedding instance-attribute 

rotary_embedding = RotaryEmbeddingWan(
    is_neox_style=False, half_head_dim=True
)

tp_inner_dim instance-attribute 

tp_inner_dim = self.num_local_heads * self.head_dim

forward 

forward(
    hidden_states: Tensor,
    *,
    cache: LingBotAttentionCache
    | ARDiffusionPagedLayerInputs,
    current_start: int,
    rotary_emb: tuple[Tensor, Tensor] | None = None,
    sink_tokens: int,
    update_cache: bool = True,
) -> Tensor

LingBotTransformerCache dataclass

One request's per-layer video K/V and reusable text K/V.

cross_attention instance-attribute 

cross_attention: list[LingBotAttentionCache | None]

self_attention instance-attribute 

self_attention: list[
    LingBotAttentionCache
    | ARDiffusionPagedLayerContext
    | ARDiffusionPagedLayerInputs
]

allocate_lingbot_cache 

allocate_lingbot_cache(
    *,
    batch_size: int,
    num_layers: int,
    max_tokens: int,
    num_local_heads: int,
    head_dim: int,
    device: device,
    dtype: dtype,
) -> LingBotTransformerCache