[Diffusion] DreamZero world model integration with CFG parallel + OpenPI serving (#2162)

Signed-off-by: Meng <meng_chen99@163.com>
Signed-off-by: Yangshen Deng <yangshen.d@outlook.com>
Signed-off-by: Suli Wang <50828270+wsl2000@users.noreply.github.com>
Co-authored-by: Meng <meng_chen99@163.com>
Co-authored-by: Suli Wang <50828270+wsl2000@users.noreply.github.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

vllm_omni/diffusion/models/dreamzero/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/models/dreamzero/action_encoder.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Action encoder/decoder for DreamZero."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def swish(x: torch.Tensor) -> torch.Tensor:
+    """swish activation: x * sigmoid(x)"""
+    return x * torch.sigmoid(x)
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    """Sinusoidal encoding: (B, T) timesteps → (B, T, dim)"""
+
+    def __init__(self, embedding_dim: int) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        timesteps = timesteps.float()
+        half_dim = self.embedding_dim // 2
+        exponent = -torch.arange(half_dim, dtype=torch.float, device=timesteps.device) * (
+            torch.log(torch.tensor(10000.0)) / half_dim
+        )
+        freqs = timesteps.unsqueeze(-1) * exponent.exp()
+        return torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1)
+
+
+class CategorySpecificLinear(nn.Module):
+    """Per-category linear: W[cat_id] @ x + b[cat_id]
+
+    Attributes:
+        W: (num_categories, input_dim, hidden_dim)  — note: 0.02 * randn init
+        b: (num_categories, hidden_dim)              — zero init
+    """
+
+    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim))
+        self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim))
+
+    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        selected_W = self.W[cat_ids]
+        selected_b = self.b[cat_ids]
+        return torch.bmm(x, selected_W) + selected_b.unsqueeze(1)
+
+
+class CategorySpecificMLP(nn.Module):
+    """Two-layer MLP: layer1 (relu) → layer2"""
+
+    def __init__(self, num_categories: int, input_dim: int, hidden_dim: int, output_dim: int) -> None:
+        super().__init__()
+        self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim)
+        self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim)
+
+    def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        hidden = F.relu(self.layer1(x, cat_ids))
+        return self.layer2(hidden, cat_ids)
+
+
+class MultiEmbodimentActionEncoder(nn.Module):
+    """Encode actions with embodiment-specific weights + sinusoidal timestep.
+
+    Flow: actions → W1 → concat(a_emb, pos_enc(timesteps)) → W2 (swish) → W3
+
+    Args:
+        action_dim: action vector dimension (e.g. 32)
+        hidden_size: output/hidden dimension (e.g. 5120 = model dim)
+        num_embodiments: number of robot types (e.g. 32)
+    """
+
+    def __init__(self, action_dim: int, hidden_size: int, num_embodiments: int) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size)
+        self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size)
+        self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size)
+        self.pos_encoding = SinusoidalPositionalEncoding(hidden_size)
+
+    def forward(self, actions: torch.Tensor, timesteps: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            actions:   (B, T, action_dim)
+            timesteps: (B, T) — per-token timestep
+            cat_ids:   (B,)   — embodiment id per sample
+        Returns:
+            (B, T, hidden_size)
+        """
+        a_emb = self.W1(actions, cat_ids)
+        tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype)
+        x = torch.cat([a_emb, tau_emb], dim=-1)
+        x = swish(self.W2(x, cat_ids))
+        x = self.W3(x, cat_ids)
+        return x

vllm_omni/diffusion/models/dreamzero/causal_wan_model.py

@@ -0,0 +1,954 @@
+# 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,
+    get_tp_group,
+)
+from vllm.model_executor.layers.conv import Conv3dLayer
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    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,
+)
+
+# ── 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 forward(self, x: torch.Tensor) -> torch.Tensor:
+        tp_size = get_tensor_model_parallel_world_size()
+        x_float = x.float()
+        local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
+        local_count = x.shape[-1]
+
+        if tp_size > 1:
+            global_sum_sq = local_sum_sq.clone()
+            torch.distributed.all_reduce(global_sum_sq, group=get_tp_group().device_group)
+            global_count = local_count * tp_size
+        else:
+            global_sum_sq = local_sum_sq
+            global_count = local_count
+
+        mean_sq = global_sum_sq / global_count
+        return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight
+
+
+# ── 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))
+        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)
+        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))
+        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
+        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,
+        freqs: torch.Tensor,
+        freqs_action: torch.Tensor,
+        freqs_state: torch.Tensor,
+        action_register_length: int | None,
+        kv_cache: torch.Tensor | 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
+
+        q = self.norm_q(self.q(x)).unflatten(2, (n, d))
+        k = self.norm_k(self.k(x)).unflatten(2, (n, d))
+        v = self.v(x).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]
+
+        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 | 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
+    """
+
+    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[torch.Tensor],
+        current_start_frame: int,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
+        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)
+
+        updated_kv_caches: list[torch.Tensor] = []
+        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,
+                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[torch.Tensor],
+        crossattn_cache: list[torch.Tensor],
+        current_start_frame: int,
+        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,
+            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)

vllm_omni/diffusion/models/dreamzero/image_encoder.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero image encoder.
+
+Only the visual tower used by DreamZero I2V inference is ported here. The
+checkpoint keys under `action_head.image_encoder.*` load via simple prefix
+stripping.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as T
+
+
+class DreamZeroLayerNorm(nn.LayerNorm):
+    """LayerNorm that preserves the input dtype."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return super().forward(x).type_as(x)
+
+
+class DreamZeroVisionSelfAttention(nn.Module):
+    """Self-attention for the vision tower."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        proj_dropout: float = 0.0,
+    ) -> 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
+        self.proj_dropout = proj_dropout
+
+        self.to_qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = x.permute(0, 2, 1, 3).reshape(batch_size, seq_len, self.dim)
+        x = self.proj(x)
+        return F.dropout(x, self.proj_dropout, self.training)
+
+
+class DreamZeroVisionAttentionBlock(nn.Module):
+    """Attention block for the vision tower."""
+
+    def __init__(
+        self,
+        dim: int,
+        mlp_ratio: float,
+        num_heads: int,
+        post_norm: bool = False,
+        activation: str = "gelu",
+        proj_dropout: float = 0.0,
+        norm_eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        if activation != "gelu":
+            raise ValueError(f"DreamZero image encoder uses GELU; got activation={activation!r}.")
+        self.post_norm = post_norm
+        hidden_dim = int(dim * mlp_ratio)
+
+        self.norm1 = DreamZeroLayerNorm(dim, eps=norm_eps)
+        self.attn = DreamZeroVisionSelfAttention(
+            dim,
+            num_heads,
+            proj_dropout=proj_dropout,
+        )
+        self.norm2 = DreamZeroLayerNorm(dim, eps=norm_eps)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(proj_dropout),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.post_norm:
+            x = x + self.norm1(self.attn(x))
+            x = x + self.norm2(self.mlp(x))
+        else:
+            x = x + self.attn(self.norm1(x))
+            x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class DreamZeroVisionTransformer(nn.Module):
+    """Vision transformer used by the image encoder."""
+
+    def __init__(
+        self,
+        image_size: int = 224,
+        patch_size: int = 14,
+        dim: int = 1280,
+        mlp_ratio: float = 4.0,
+        out_dim: int = 1024,
+        num_heads: int = 16,
+        num_layers: int = 32,
+        pool_type: str = "token",
+        pre_norm: bool = True,
+        post_norm: bool = False,
+        activation: str = "gelu",
+        proj_dropout: float = 0.0,
+        embedding_dropout: float = 0.0,
+        norm_eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        if pool_type != "token":
+            raise ValueError(f"DreamZero image encoder only supports pool_type='token', got {pool_type!r}.")
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_patches = (image_size // patch_size) ** 2
+        self.dim = dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.pool_type = pool_type
+
+        gain = 1.0 / math.sqrt(dim)
+        self.patch_embedding = nn.Conv2d(
+            3,
+            dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=not pre_norm,
+        )
+        self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
+        self.pos_embedding = nn.Parameter(
+            gain * torch.randn(1, self.num_patches + 1, dim),
+        )
+        self.dropout = nn.Dropout(embedding_dropout)
+        self.pre_norm = DreamZeroLayerNorm(dim, eps=norm_eps) if pre_norm else None
+        self.transformer = nn.Sequential(
+            *[
+                DreamZeroVisionAttentionBlock(
+                    dim=dim,
+                    mlp_ratio=mlp_ratio,
+                    num_heads=num_heads,
+                    post_norm=post_norm,
+                    activation=activation,
+                    proj_dropout=proj_dropout,
+                    norm_eps=norm_eps,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.post_norm = DreamZeroLayerNorm(dim, eps=norm_eps)
+        self.head = nn.Parameter(gain * torch.randn(dim, out_dim))
+
+    def forward(self, x: torch.Tensor, use_31_block: bool = False) -> torch.Tensor:
+        batch_size = x.shape[0]
+        x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
+        x = torch.cat(
+            [
+                self.cls_embedding.expand(batch_size, -1, -1).to(dtype=x.dtype, device=x.device),
+                x,
+            ],
+            dim=1,
+        )
+        x = self.dropout(x + self.pos_embedding.to(dtype=x.dtype, device=x.device))
+        if self.pre_norm is not None:
+            x = self.pre_norm(x)
+
+        if use_31_block:
+            return self.transformer[:-1](x)
+        return self.transformer(x)
+
+
+class _DreamZeroCLIPContainer(nn.Module):
+    """Container matching checkpoint names under `model.visual.*`."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.log_scale = nn.Parameter(torch.ones(()))
+        self.visual = DreamZeroVisionTransformer(
+            image_size=224,
+            patch_size=14,
+            dim=1280,
+            mlp_ratio=4.0,
+            out_dim=1024,
+            num_heads=16,
+            num_layers=32,
+            pool_type="token",
+            pre_norm=True,
+            post_norm=False,
+            activation="gelu",
+            proj_dropout=0.0,
+            embedding_dropout=0.0,
+            norm_eps=1e-5,
+        )
+
+
+class DreamZeroImageEncoder(nn.Module):
+    """Image encoder wrapper."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.model = _DreamZeroCLIPContainer()
+        # returns a composed transform whose last stage is CLIP normalization.
+        self.transforms = T.Compose(
+            [
+                T.Normalize(
+                    mean=[0.48145466, 0.4578275, 0.40821073],
+                    std=[0.26862954, 0.26130258, 0.27577711],
+                ),
+            ]
+        )
+
+    def encode_image(self, videos: torch.Tensor) -> torch.Tensor:
+        """Encode images for I2V conditioning."""
+        size = (self.model.visual.image_size,) * 2
+        videos = torch.cat(
+            [
+                F.interpolate(
+                    frame_batch,
+                    size=size,
+                    mode="bicubic",
+                    align_corners=False,
+                )
+                for frame_batch in videos
+            ]
+        )
+        videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5))
+
+        param_dtype = next(iter(self.model.visual.parameters())).dtype
+        videos = videos.to(dtype=param_dtype)
+        out = self.model.visual(videos, use_31_block=True)
+        return out.clone()

vllm_omni/diffusion/models/dreamzero/state_dreamzero.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero pipeline persistent state."""
+
+from __future__ import annotations
+
+import logging
+from collections import deque
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Number of frames per chunk for subsequent calls (first call uses 1)
+FRAMES_PER_CHUNK = 4
+
+
+class DreamZeroState:
+    """Pipeline persistent state across forward() calls.
+
+    Lifecycle:
+        - Created once in DreamZeroPipeline.__init__()
+        - Mutated every forward() call (frame append, KV cache grow)
+        - reset() on new session / language change / local_attn_size exceeded
+    """
+
+    def __init__(self) -> None:
+        self.reset()
+
+    # ------------------------------------------------------------------
+    # Frame accumulation (single stitched buffer)
+    # Transform outputs stitched single frame per call.
+    # We accumulate here to build multi-frame video for AR inference.
+    # ------------------------------------------------------------------
+
+    def accumulate_frames(self, stitched: np.ndarray) -> np.ndarray:
+        """Accumulate stitched frames and return multi-frame video.
+
+        Args:
+            stitched: (H, W, C) single frame or (T, H, W, C) multi-frame,
+                      already stitched by transform.
+
+        Returns:
+            (T, H, W, C) ndarray. T=1 for first call, T=FRAMES_PER_CHUNK(4) after.
+        """
+        if stitched.ndim == 3:
+            self.stitched_buffer.append(stitched)
+        elif stitched.ndim == 4:
+            self.stitched_buffer.extend(list(stitched))
+        else:
+            raise ValueError(f"Expected 3D or 4D stitched, got {stitched.ndim}D")
+
+        num_frames = 1 if self.call_count == 0 else FRAMES_PER_CHUNK
+
+        buffer_frames = list(self.stitched_buffer)
+        if len(buffer_frames) >= num_frames:
+            frames = buffer_frames[-num_frames:]
+        else:
+            frames = buffer_frames
+            while len(frames) < num_frames:
+                frames.insert(0, buffer_frames[0])
+
+        self.call_count += 1
+        return np.stack(frames, axis=0)  # (T, H, W, C)
+
+    # ------------------------------------------------------------------
+    # Reset / should_reset
+    # ------------------------------------------------------------------
+
+    def reset(self) -> None:
+        """Clear all state."""
+        self.stitched_buffer: deque[np.ndarray] = deque(maxlen=FRAMES_PER_CHUNK)
+        self.call_count: int = 0
+
+        # KV cache once robot-policy diffusion supports that integration.
+        self.kv_cache: list[torch.Tensor] | None = None
+        self.kv_cache_neg: list[torch.Tensor] | None = None
+        self.crossattn_cache: list[dict[str, bool | torch.Tensor | None]] | None = None
+        self.crossattn_cache_neg: list[dict[str, bool | torch.Tensor | None]] | None = None
+        self.current_start_frame: int = 0
+
+        self.clip_feas: torch.Tensor | None = None
+        self.ys: torch.Tensor | None = None
+        self.language: torch.Tensor | None = None
+
+    def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
+        """Determine if state should be reset before this forward()."""
+        if self.language is None:
+            logger.info("language is None, resetting")
+            return True
+
+        if text_tokens is not None and not torch.equal(self.language, text_tokens):
+            logger.info("language changed, resetting")
+            return True
+
+        # NOTE: after accumulate_frames, num_video_frames is the accumulated T
+        # (1 for first call, 4 for subsequent). Only reset on true single-frame
+        # which happens when the stitched_buffer was cleared externally.
+        if num_video_frames == 1 and self.call_count > 1:
+            logger.info("single frame input after first call, resetting")
+            return True
+
+        if local_attn_size != -1 and self.current_start_frame >= local_attn_size:
+            logger.info(
+                "current_start_frame %d >= local_attn_size %d, resetting", self.current_start_frame, local_attn_size
+            )
+            return True
+
+        return False
+
+    # ------------------------------------------------------------------
+    # KV cache management
+    # ------------------------------------------------------------------
+
+    def create_kv_caches(
+        self,
+        batch_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        num_layers: int,
+        num_heads: int,
+        head_dim: int,
+    ) -> None:
+        """Initialize empty KV caches and cross-attention caches."""
+        self.kv_cache = [
+            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
+        ]
+        self.kv_cache_neg = [
+            torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
+        ]
+
+        self.crossattn_cache = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
+        self.crossattn_cache_neg = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)]
+
+    def update_kv_cache(
+        self,
+        layer_index: int,
+        updated_kv: torch.Tensor,
+        is_negative: bool = False,
+    ) -> None:
+        """Update a single layer's KV cache after prefill."""
+        cache = self.kv_cache_neg if is_negative else self.kv_cache
+        if cache is None:
+            raise RuntimeError("KV caches not initialized, call create_kv_caches first.")
+        cache[layer_index] = updated_kv.clone()
+
+    def get_kv_caches(self, is_negative: bool = False) -> list[torch.Tensor]:
+        """Get KV caches for the specified branch."""
+        cache = self.kv_cache_neg if is_negative else self.kv_cache
+        if cache is None:
+            raise RuntimeError("KV caches not initialized.")
+        return cache
+
+    def get_crossattn_caches(self, is_negative: bool = False) -> list[dict[str, bool | torch.Tensor | None]]:
+        """Get cross-attention caches for the specified branch."""
+        cache = self.crossattn_cache_neg if is_negative else self.crossattn_cache
+        if cache is None:
+            raise RuntimeError("Cross-attn caches not initialized.")
+        return cache

vllm_omni/diffusion/models/dreamzero/transform/__init__.py

@@ -0,0 +1,31 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import importlib
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.models.dreamzero.transform.base import TRANSFORMS
+
+logger = init_logger(__name__)
+
+DEFAULT_EMBODIMENT = "roboarena"
+_BUILTIN_TRANSFORM_MODULES = (
+    "vllm_omni.diffusion.models.dreamzero.transform.droid",
+    "vllm_omni.diffusion.models.dreamzero.transform.roboarena",
+)
+
+
+def ensure_transforms_loaded() -> None:
+    """Import DreamZero transform modules and verify registration."""
+    for module_name in _BUILTIN_TRANSFORM_MODULES:
+        try:
+            importlib.import_module(module_name)
+        except Exception as exc:
+            logger.exception("Failed to import DreamZero transform module %s", module_name)
+            raise RuntimeError(f"Failed to import DreamZero transform module '{module_name}'.") from exc
+
+    if DEFAULT_EMBODIMENT not in TRANSFORMS:
+        raise RuntimeError(f"Built-in DreamZero transform '{DEFAULT_EMBODIMENT}' is not registered after import.")

vllm_omni/diffusion/models/dreamzero/transform/base.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Base transform interface for DreamZero robot policy serving.
+
+Transforms handle dataset-specific concerns ONLY:
+  - Observation key mapping
+  - Multi-view stitching (embodiment-specific layout)
+  - Language template wrapping (embodiment-specific)
+  - Raw state extraction (dataset-specific keys)
+  - Output action slicing (to actual action_dim)
+
+Model-specific concerns belong in the pipeline:
+  - Tokenization (pipeline owns tokenizer)
+  - State padding (pipeline knows MAX_STATE_DIM)
+  - Negative prompt (pipeline owns the string)
+  - Noise generation, encoding, decoding
+
+Flow:
+  raw obs (dataset format)
+    → DreamZeroPipeline selects transform by embodiment
+      → unified dict (stitched video, templated prompt str, raw state)
+        → tokenize, pad, encode, denoise
+          → transform_action_output()
+            → ndarray (N, action_dim)
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+
+class RobotPolicyTransform:
+    """Base class for dataset-specific observation transforms.
+
+    Subclasses MUST define:
+        IMAGE_KEY_MAP: dict — dataset obs keys → unified keys
+        EMBODIMENT_NAME: str — embodiment identity (pipeline maps to numeric ID)
+        ACTION_DIM: int — actual action dimensions (for output slicing)
+
+    Subclasses MUST override:
+        _stitch_views()      — multi-view → single stitched image
+        _language_template() — prompt → embodiment-aware template
+        _extract_raw_state() — obs → raw state ndarray
+    """
+
+    IMAGE_KEY_MAP: dict[str, str]
+    EMBODIMENT_NAME: str
+    ACTION_DIM: int
+
+    def transform_input(self, obs: dict) -> dict:
+        """Dataset-specific transform: key map → stitch → template → state."""
+        # 1. Map image keys → unified keys
+        images: dict[str, np.ndarray] = {}
+        for src_key, dst_key in self.IMAGE_KEY_MAP.items():
+            if src_key in obs:
+                images[dst_key] = np.asarray(obs[src_key])
+
+        # 2. Multi-view stitching
+        stitched = self._stitch_views(images)
+
+        # 3. Language template (string only, pipeline tokenizes)
+        prompt = obs.get("prompt", "")
+        templated_prompt = self._language_template(prompt)
+
+        # 4. Raw state extraction (pipeline pads)
+        raw_state = self._extract_raw_state(obs)
+
+        # 5. Build unified output
+        unified: dict[str, Any] = {
+            "images": stitched,  # ndarray (T, H_out, W_out, 3)
+            "prompt": templated_prompt,  # str (templated, not tokenized)
+            "state": raw_state,  # ndarray (state_dim,) — pipeline pads
+            "embodiment_name": self.EMBODIMENT_NAME,
+        }
+        if "session_id" in obs:
+            unified["session_id"] = obs["session_id"]
+        return unified
+
+    def transform_action_output(self, actions: Any) -> np.ndarray:
+        """Adapt model action output to this transform's action dimensions."""
+        actions = np.asarray(actions, dtype=np.float32)
+        # Handle any remaining batch dims: squeeze to 2D (horizon, dim)
+        while actions.ndim > 2:
+            actions = actions[0]
+        # Slice padded dim to actual ACTION_DIM
+        if actions.ndim == 2 and actions.shape[-1] > self.ACTION_DIM:
+            actions = actions[:, : self.ACTION_DIM]
+        return actions
+
+    # ------------------------------------------------------------------
+    # Subclass MUST override
+    # ------------------------------------------------------------------
+
+    def _stitch_views(self, images: dict[str, np.ndarray]) -> np.ndarray:
+        """Stitch camera views into single image.
+        Input: unified key → ndarray (H,W,3) or (T,H,W,3).
+        Output: ndarray (T, H_out, W_out, 3).
+        """
+        raise NotImplementedError
+
+    def _language_template(self, prompt: str) -> str:
+        """Wrap prompt in embodiment-specific template string."""
+        raise NotImplementedError
+
+    def _extract_raw_state(self, obs: dict) -> np.ndarray:
+        """Extract raw state vector from obs.
+        Returns: ndarray (state_dim,) float64. Pipeline handles padding.
+        """
+        raise NotImplementedError
+
+
+# Transform registry — keyed by embodiment/dataset name
+TRANSFORMS: dict[str, RobotPolicyTransform] = {}
+
+
+def register_transform(name: str, transform: RobotPolicyTransform) -> None:
+    TRANSFORMS[name] = transform
+
+
+def get_transform(name: str) -> RobotPolicyTransform:
+    if name not in TRANSFORMS:
+        raise KeyError(f"Unknown transform '{name}'. Available: {list(TRANSFORMS.keys())}")
+    return TRANSFORMS[name]

vllm_omni/diffusion/models/dreamzero/transform/droid.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DROID dataset transform.
+
+DROID uses 1-indexed exterior cameras, 3 views total (OXE_DROID embodiment).
+Stitching layout (same as RoboArena — both are OXE_DROID):
+    ┌─────────────────────────┐
+    │   wrist (2x width)      │  ← pixel-repeat along width
+    ├────────────┬────────────┤
+    │  left ext  │ right ext  │
+    └────────────┴────────────┘
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torchvision.transforms.v2 as T
+
+from vllm_omni.diffusion.models.dreamzero.transform.base import (
+    RobotPolicyTransform,
+    register_transform,
+)
+
+
+class DroidTransform(RobotPolicyTransform):
+    """Transform for DROID dataset (OXE_DROID embodiment).
+
+    DROID observation keys (1-indexed exterior cameras):
+        observation/exterior_image_1_left  → left exterior
+        observation/exterior_image_2_left  → right exterior
+        observation/wrist_image_left       → wrist
+    """
+
+    IMAGE_KEY_MAP = {
+        "observation/exterior_image_1_left": "images/exterior_0",
+        "observation/exterior_image_2_left": "images/exterior_1",
+        "observation/wrist_image_left": "images/wrist",
+    }
+    EMBODIMENT_NAME = "oxe_droid"
+    ACTION_DIM = 8  # 7 joint + 1 gripper
+    _VIDEO_CROP_SCALE = 0.95
+    _VIDEO_RESIZE_HW = (176, 320)
+
+    @classmethod
+    def _preprocess_view(cls, arr: np.ndarray) -> np.ndarray:
+        """Apply per-view crop and resize before stitching."""
+        frames = torch.from_numpy(arr).to(torch.float32).permute(0, 3, 1, 2) / 255.0
+        crop_h = int(arr.shape[1] * cls._VIDEO_CROP_SCALE)
+        crop_w = int(arr.shape[2] * cls._VIDEO_CROP_SCALE)
+        frames = T.CenterCrop((crop_h, crop_w))(frames)
+        frames = T.Resize(
+            cls._VIDEO_RESIZE_HW,
+            interpolation=T.InterpolationMode.BILINEAR,
+            antialias=True,
+        )(frames)
+        return (frames.permute(0, 2, 3, 1) * 255.0).to(torch.uint8).cpu().numpy()
+
+    def _stitch_views(self, images: dict[str, np.ndarray]) -> np.ndarray:
+        """OXE_DROID 2x2 stitching: wrist top (2x wide), exteriors bottom."""
+        left_ext = images.get("images/exterior_0")
+        right_ext = images.get("images/exterior_1")
+        wrist = images.get("images/wrist")
+
+        # Ensure 4D: (T, H, W, C)
+        def ensure_4d(arr: np.ndarray | None) -> np.ndarray | None:
+            if arr is None:
+                return None
+            return arr if arr.ndim == 4 else arr[np.newaxis]
+
+        left_ext = ensure_4d(left_ext)
+        right_ext = ensure_4d(right_ext)
+        wrist = ensure_4d(wrist)
+
+        # Determine shape from first available view.
+        ref = next((v for v in [wrist, left_ext, right_ext] if v is not None), None)
+        if ref is None:
+            return np.zeros((1, 352, 640, 3), dtype=np.uint8)
+
+        # Apply per-view crop + resize before stitching.
+        def maybe_preprocess(arr: np.ndarray | None) -> np.ndarray | None:
+            if arr is None:
+                return None
+            return self._preprocess_view(arr)
+
+        left_ext = maybe_preprocess(left_ext)
+        right_ext = maybe_preprocess(right_ext)
+        wrist = maybe_preprocess(wrist)
+        ref = next((v for v in [wrist, left_ext, right_ext] if v is not None), None)
+        if ref is None:
+            raise RuntimeError("Expected at least one DROID camera view after preprocessing.")
+        t, h, w, c = ref.shape
+
+        out = np.zeros((t, 2 * h, 2 * w, c), dtype=ref.dtype)  # (T, 2H, 2W, C)
+
+        # Top row: wrist repeated 2x along width.
+        if wrist is not None:
+            wrist_wide = np.repeat(wrist, 2, axis=2)  # (T, H, 2W, C)
+            out[:, :h, :] = wrist_wide
+
+        # Bottom row: left exterior | right exterior.
+        if left_ext is not None:
+            out[:, h:, :w] = left_ext
+        if right_ext is not None:
+            out[:, h:, w:] = right_ext
+
+        return out
+
+    def _language_template(self, prompt: str) -> str:
+        """Expand the language prompt for the OXE_DROID multi-view format."""
+        prompt = (prompt or "Perform the default behavior.").strip()
+        prompt_lower = prompt.lower()
+        return (
+            "A multi-view video shows that a robot "
+            + prompt_lower
+            + " The video is split into three views: The top view shows the "
+            + "camera view from the robot's wrist, the bottom-left view shows "
+            + "the camera view from the left exterior camera, and the "
+            + "bottom-right view shows the camera view from the right exterior "
+            + "camera. During training, one of the two bottom exterior views "
+            + "may be a black screen (dropped view). The robot "
+            + prompt_lower
+        )
+
+    def _extract_raw_state(self, obs: dict) -> np.ndarray:
+        """OXE_DROID state: 7 joint + 1 gripper = 8 dims."""
+        parts = []
+        if "observation/joint_position" in obs:
+            parts.append(np.asarray(obs["observation/joint_position"], dtype=np.float64).flatten())
+        if "observation/gripper_position" in obs:
+            parts.append(np.asarray(obs["observation/gripper_position"], dtype=np.float64).flatten())
+        if parts:
+            return np.concatenate(parts)
+        return np.zeros(8, dtype=np.float64)
+
+
+register_transform("droid", DroidTransform())

vllm_omni/diffusion/models/dreamzero/transform/roboarena.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""RoboArena dataset transform.
+
+RoboArena uses 0-indexed exterior cameras, 3 views total (OXE_DROID embodiment).
+Same stitching layout as DROID — both map to OXE_DROID in DreamZero.
+
+"""
+
+from __future__ import annotations
+
+from vllm_omni.diffusion.models.dreamzero.transform.base import (
+    register_transform,
+)
+from vllm_omni.diffusion.models.dreamzero.transform.droid import (
+    DroidTransform,
+)
+
+
+class RoboArenaTransform(DroidTransform):
+    """Transform for RoboArena dataset.
+
+    Same embodiment as DROID (OXE_DROID), same stitching and template.
+    Only difference: 0-indexed exterior camera keys.
+
+    RoboArena observation keys (0-indexed):
+        observation/exterior_image_0_left  → left exterior
+        observation/exterior_image_1_left  → right exterior
+        observation/wrist_image_left       → wrist
+
+    """
+
+    IMAGE_KEY_MAP = {
+        "observation/exterior_image_0_left": "images/exterior_0",
+        "observation/exterior_image_1_left": "images/exterior_1",
+        "observation/wrist_image_left": "images/wrist",
+    }
+
+
+register_transform("roboarena", RoboArenaTransform())

vllm_omni/diffusion/models/dreamzero/utils.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DreamZero model constants shared by the pipeline."""
+
+DEFAULT_NUM_INFERENCE_STEPS = 16
+DEFAULT_CFG_SCALE = 5.0
+DEFAULT_SIGMA_SHIFT = 5.0
+DEFAULT_SEED = 1140
+
+DEFAULT_NEGATIVE_PROMPT = (
+    "Vibrant colors, overexposed, static, blurry details, text, subtitles, "
+    "style, artwork, painting, image, still, grayscale, dull, worst quality, "
+    "low quality, JPEG artifacts, ugly, mutilated, extra fingers, bad hands, "
+    "bad face, deformed, disfigured, mutated limbs, fused fingers, stagnant "
+    "image, cluttered background, three legs, many people in the background, "
+    "walking backwards."
+)
+
+DEFAULT_EMBODIMENT_NAME_TO_ID = {
+    "oxe_droid": 17,
+    "agibot": 26,
+    "gr1_unified": 24,
+    "xdof": 22,
+    "yam": 32,
+    "mecka_hands": 27,
+    "lapa": 27,
+    "dream": 31,
+}

vllm_omni/diffusion/models/dreamzero/video_export_worker.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import torch
+
+
+class DreamZeroVideoExportWorkerExtension:
+    """DreamZero worker RPCs used by offline example video export."""
+
+    def decode_video_latents_to_uint8(self, video_latents: torch.Tensor) -> torch.Tensor:
+        if self.model_runner is None or self.model_runner.pipeline is None:
+            raise RuntimeError("DreamZero pipeline is not initialized on this worker.")
+
+        with torch.inference_mode():
+            decoded = self.model_runner.pipeline.decode_video_latents(video_latents)
+            decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()
+            decoded = decoded.clamp(-1, 1) * 0.5 + 0.5
+            decoded = (decoded * 255.0).round().to(torch.uint8).cpu()
+        return decoded