Merge branch 'main' into fix/pr4709-review-followup

vllm_omni/diffusion/models/audiox/pipeline_audiox.py

@@ -60,6 +60,15 @@ def get_audiox_post_process_func(od_config: OmniDiffusionConfig):
     """Convert the pipeline's float audio tensor to a CPU numpy array for serving."""
 
     def post_process_func(audio: torch.Tensor) -> Any:
+        if isinstance(audio, dict) and isinstance(audio.get("payload"), dict):
+            payload = dict(audio["payload"])
+            audio_payload = payload.get("audio")
+            if isinstance(audio_payload, torch.Tensor):
+                payload["audio"] = audio_payload.detach().cpu().float().numpy()
+            return {
+                "payload": payload,
+                "metadata": audio.get("metadata") or {},
+            }
         if isinstance(audio, torch.Tensor):
             return audio.detach().cpu().float().numpy()
         return audio
@@ -915,9 +924,10 @@ class AudioXPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfilerMix
             audio = audio[..., :target_samples]
 
         stage_durations = self.stage_durations if getattr(self, "enable_diffusion_pipeline_profiler", False) else None
-        custom_output = {"audiox_task": task_norm}
         return DiffusionOutput(
-            output=audio,
-            custom_output=custom_output,
+            output={
+                "payload": {"audio": audio},
+                "metadata": {"audio": {"audiox_task": task_norm}},
+            },
             stage_durations=stage_durations,
         )

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -739,8 +739,10 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
                         text_output = text_output.split("<|im_start|>")[-1]
 
             return DiffusionOutput(
-                output=text_output,
-                custom_output={"text_output": text_output},
+                output={
+                    "payload": {"text": text_output},
+                    "metadata": {"text": {"text_output": text_output}},
+                },
                 stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
             )
 
@@ -847,22 +849,28 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         if trajectory_log_probs:
             trajectory_log_probs_stacked = torch.stack(trajectory_log_probs)
 
-        custom = {}
+        payload = {"image": img}
+        metadata = {}
         if think_text is not None:
-            custom["think_text"] = think_text
-        # Mirror the PIL image into ``custom_output`` so callers reading via
-        # the orchestrator IPC boundary (which strips the bare ``output``
-        # field) can still recover the result.  ``video_frames`` already
-        # uses this pattern.
-        custom["image"] = img
+            metadata["text"] = {"think_text": think_text}
+        trajectory_payload = {}
+        if trajectory_latents_stacked is not None:
+            trajectory_payload["latents"] = trajectory_latents_stacked
+        if trajectory_timesteps_stacked is not None:
+            trajectory_payload["timesteps"] = trajectory_timesteps_stacked
+        if trajectory_log_probs_stacked is not None:
+            trajectory_payload["log_probs"] = trajectory_log_probs_stacked
+        if trajectory_decoded is not None:
+            trajectory_payload["decoded"] = trajectory_decoded
+        if trajectory_payload:
+            payload["trajectory"] = trajectory_payload
+            metadata["trajectory"] = {"type": "denoising"}
 
         return DiffusionOutput(
-            output=img,
-            trajectory_latents=trajectory_latents_stacked,
-            trajectory_timesteps=trajectory_timesteps_stacked,
-            trajectory_log_probs=trajectory_log_probs_stacked,
-            trajectory_decoded=trajectory_decoded,
-            custom_output=custom,
+            output={
+                "payload": payload,
+                "metadata": metadata,
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 

vllm_omni/diffusion/models/cosmos3/__init__.py

@@ -3,16 +3,16 @@
 
 from .pipeline_cosmos3 import (
     Cosmos3OmniDiffusersPipeline,
-    get_cosmos3_action_post_process_func,
     get_cosmos3_post_process_func,
     get_cosmos3_pre_process_func,
 )
 from .transformer_cosmos3 import Cosmos3VFMTransformer
+from .transformer_cosmos3_edge import Cosmos3EdgeVFMTransformer
 
 __all__ = [
     "Cosmos3OmniDiffusersPipeline",
-    "get_cosmos3_action_post_process_func",
     "get_cosmos3_post_process_func",
     "get_cosmos3_pre_process_func",
     "Cosmos3VFMTransformer",
+    "Cosmos3EdgeVFMTransformer",
 ]

vllm_omni/diffusion/models/cosmos3/action.py

@@ -4,8 +4,8 @@
 
 These helpers cover action modes that use action tokens as an auxiliary output
 stream: ``policy``, ``forward_dynamics``, and ``inverse_dynamics``. The pipeline
-returns predicted actions through ``custom_output`` rather than as a normal
-``modalities`` entry.
+returns predicted actions through the diffusion output payload/metadata
+envelope rather than as a video frame stream.
 """
 
 from __future__ import annotations

vllm_omni/diffusion/models/cosmos3/guardrails.py

@@ -21,6 +21,7 @@ from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.progress_bar import _is_rank_zero
 from vllm_omni.errors import GuardrailViolationError
+from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -28,6 +29,28 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+# ---------------------------------------------------------------------------
+# NPU/XPU compatibility: cosmos_guardrail.cosmos_utils.load_model calls
+# ``torch.load(path, weights_only=True)`` without ``map_location``.
+# On NPU/XPU, ``torch.cuda.is_available()`` is False, so deserializing a
+# CUDA checkpoint with ``weights_only=True`` raises an UnpicklingError.
+# Work around by patching ``torch.load`` to default ``map_location="cpu"``
+# when CUDA is unavailable and no explicit map_location was given.
+# ---------------------------------------------------------------------------
+_original_torch_load = torch.load
+
+
+def _patched_torch_load(*args, **kwargs):
+    if (
+        "map_location" not in kwargs
+        and not torch.cuda.is_available()
+        and (current_omni_platform.is_npu() or current_omni_platform.is_xpu())
+    ):
+        kwargs["map_location"] = "cpu"
+    return _original_torch_load(*args, **kwargs)
+
+
+torch.load = _patched_torch_load  # type: ignore[assignment]
 
 try:
     from cosmos_guardrail import CosmosSafetyChecker
@@ -75,7 +98,7 @@ def _build_text_guardrail(checker: Any) -> TextGuardrailFn:
 
 def _build_video_guardrail(checker: Any, offload_to_cpu: bool) -> VideoGuardrailFn:
     video_models = _nn_models(checker.video_guardrail)
-    compute_device = "cuda"
+    compute_device = current_omni_platform.device_type
 
     def video_guardrail(frames: np.ndarray) -> np.ndarray:
         if offload_to_cpu:
@@ -112,9 +135,9 @@ def _init_default_guardrails(offload_to_cpu: bool = False) -> None:
     checker = CosmosSafetyChecker()
 
     # Place text models on their resting device permanently. Video models
-    # idle on CPU when offload is on and move to GPU per-call (handled in
+    # idle on CPU when offload is on and move to compute device per-call (handled in
     # the video guardrail closure).
-    idle_device = "cpu" if offload_to_cpu else "cuda"
+    idle_device = "cpu" if offload_to_cpu else current_omni_platform.device_type
     for m in _nn_models(checker.text_guardrail):
         m.to(idle_device)
     for m in _nn_models(checker.video_guardrail):

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -12,7 +12,9 @@ Ported from the TRT-LLM integration (tekit branch user/shreyasm/cosmos3).
 from __future__ import annotations
 
 import math
-from typing import Any
+from collections.abc import Iterator
+from contextlib import AbstractContextManager, contextmanager, nullcontext
+from typing import TYPE_CHECKING, Any
 
 import torch
 import torch.distributed as dist
@@ -38,6 +40,9 @@ from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_
 from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
 from vllm_omni.platforms import current_omni_platform
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+
 logger = init_logger(__name__)
 
 
@@ -391,7 +396,6 @@ class TimestepEmbedder(nn.Module):
         hidden_size: int,
         frequency_embedding_size: int = 256,
         max_period: int = 10000,
-        target_dtype: torch.dtype = torch.bfloat16,
     ) -> None:
         super().__init__()
         # Following diffusers naming pattern here for checkpoint compatibility.
@@ -402,7 +406,7 @@ class TimestepEmbedder(nn.Module):
         self.hidden_size = hidden_size
 
         half = frequency_embedding_size // 2
-        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half)
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half)
         self.register_buffer("freqs", freqs, persistent=False)
 
     def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -582,6 +586,7 @@ class Cosmos3CrossAttention(nn.Module):
         head_dim: int,
         rms_norm_eps: float,
         quant_config: QuantizationConfig | None = None,
+        qk_norm: bool = True,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -631,8 +636,10 @@ class Cosmos3CrossAttention(nn.Module):
             prefix=f"{prefix}.to_out",
         )
 
-        self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
-        self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.qk_norm = qk_norm
+        if self.qk_norm:
+            self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps)
+            self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps)
 
         self.attn = FrameworkAttention(
             num_heads=self.num_heads,
@@ -716,8 +723,9 @@ class Cosmos3CrossAttention(nn.Module):
         v = self.to_v(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim)
 
         # Per-head QK norm
-        q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
-        k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
+        if self.qk_norm:
+            q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon)
+            k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon)
 
         # Qwen3-style RoPE
         q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
@@ -801,6 +809,8 @@ class Cosmos3GenDecoderLayer(nn.Module):
         head_dim: int,
         rms_norm_eps: float,
         quant_config: QuantizationConfig | None = None,
+        mlp_cls: type[nn.Module] = Cosmos3GatedMLP,
+        qk_norm: bool = True,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -812,11 +822,12 @@ class Cosmos3GenDecoderLayer(nn.Module):
             head_dim=head_dim,
             rms_norm_eps=rms_norm_eps,
             quant_config=quant_config,
+            qk_norm=qk_norm,
             prefix=f"{prefix}.cross_attention",
         )
         self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
-        self.mlp = Cosmos3GatedMLP(
+        self.mlp = mlp_cls(
             hidden_size=hidden_size,
             intermediate_size=intermediate_size,
             quant_config=quant_config,
@@ -869,6 +880,8 @@ class Cosmos3LanguageModel(nn.Module):
     computed once and reused across all sampling steps.
     """
 
+    _layerwise_offload_blocks_attrs = ["layers"]
+
     def __init__(
         self,
         *,
@@ -959,7 +972,13 @@ class Cosmos3VFMTransformer(nn.Module):
     runs at each denoising step over the target video/image latent stream and
     optional transfer-control, action, and sound latent streams.
 
-    Layerwise offloading uses ``gen_layers`` as the block container.
+    Layerwise offloading uses ``gen_layers`` as the GEN block container.  The
+    nested UND language model declares its own ``layers`` container so the two
+    pathways are offloaded as independent rings.
+
+    Model-level CPU offload is Cosmos3-local: the UND ``reasoner`` and GEN
+    ``generator`` components are swapped at the phase boundaries marked by
+    ``with self._offload_context(...)``.
 
     Sequence parallelism uses ``_sp_plan`` to shard/gather the GEN pathway at
     module boundaries. ``Cosmos3CrossAttention`` checks
@@ -983,6 +1002,9 @@ class Cosmos3VFMTransformer(nn.Module):
 
     packed_modules_mapping = {}
 
+    _language_model_cls = Cosmos3LanguageModel
+    _gen_mlp_cls = Cosmos3GatedMLP
+
     @staticmethod
     def _is_transformer_block(name: str, module) -> bool:
         return ("gen_layers" in name or "language_model.layers" in name) and name.split(".")[-1].isdigit()
@@ -1019,6 +1041,25 @@ class Cosmos3VFMTransformer(nn.Module):
             if actual != expected:
                 raise ValueError(f"Unsupported Cosmos3 transformer config: {key}={actual!r}; expected {expected!r}.")
 
+    @classmethod
+    def _resolve_rms_norm_eps(cls, model_config: Any) -> float:
+        return float(_tf_config_get(model_config, "rms_norm_eps", 1e-6))
+
+    @classmethod
+    def _resolve_rope_theta(cls, model_config: Any) -> float:
+        return float(_tf_config_get(model_config, "rope_theta", 5_000_000))
+
+    @classmethod
+    def _resolve_mrope_section(cls, model_config: Any) -> list[int]:
+        rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {}
+        return list(rope_scaling.get("mrope_section", [24, 20, 20]))
+
+    def _language_model_kwargs(self) -> dict[str, Any]:
+        return {}
+
+    def validate_loaded_weights(self, loaded: set[str]) -> None:
+        del loaded
+
     def __init__(
         self,
         od_config: OmniDiffusionConfig,
@@ -1031,7 +1072,6 @@ class Cosmos3VFMTransformer(nn.Module):
         super().__init__()
         model_config = od_config.tf_model_config
         self._validate_supported_config(model_config)
-        rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {}
 
         self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096))
         self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36))
@@ -1040,9 +1080,10 @@ class Cosmos3VFMTransformer(nn.Module):
         self.head_dim = int(_tf_config_get(model_config, "head_dim", 128))
         self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288))
         self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936))
-        self.rms_norm_eps = float(_tf_config_get(model_config, "rms_norm_eps", 1e-6))
-        self.rope_theta = float(_tf_config_get(model_config, "rope_theta", 5_000_000))
-        self.mrope_section = list(rope_scaling.get("mrope_section", [24, 20, 20]))
+        self.rms_norm_eps = self._resolve_rms_norm_eps(model_config)
+        self.rope_theta = self._resolve_rope_theta(model_config)
+        self.mrope_section = self._resolve_mrope_section(model_config)
+        self.qk_norm_for_diffusion = bool(_tf_config_get(model_config, "qk_norm_for_diffusion", True))
         self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2))
         self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48))
         self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001))
@@ -1082,10 +1123,12 @@ class Cosmos3VFMTransformer(nn.Module):
         )
         self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size
 
+        self.use_k_norm_und_for_gen = _tf_config_get(model_config, "use_k_norm_und_for_gen", None)
+
         dtype = od_config.dtype
         quant_config = getattr(od_config, "quantization_config", None) if od_config else None
 
-        self.language_model = Cosmos3LanguageModel(
+        self.language_model = self._language_model_cls(
             hidden_size=self.hidden_size,
             intermediate_size=self.intermediate_size,
             num_hidden_layers=self.num_hidden_layers,
@@ -1098,12 +1141,13 @@ class Cosmos3VFMTransformer(nn.Module):
             mrope_section=self.mrope_section,
             quant_config=quant_config,
             prefix="language_model",
+            **self._language_model_kwargs(),
         )
 
         # Video projection layers are small; not worth quantizing.
         self.proj_in = nn.Linear(self.patch_latent_dim, self.hidden_size)
         self.proj_out = nn.Linear(self.hidden_size, self.patch_latent_dim)
-        self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=dtype)
+        self.time_embedder = TimestepEmbedder(self.hidden_size)
         if self.action_gen:
             self.action_proj_in = DomainAwareLinear(
                 self.action_dim,
@@ -1134,6 +1178,8 @@ class Cosmos3VFMTransformer(nn.Module):
                     head_dim=self.head_dim,
                     rms_norm_eps=self.rms_norm_eps,
                     quant_config=quant_config,
+                    mlp_cls=self._gen_mlp_cls,
+                    qk_norm=self.qk_norm_for_diffusion,
                     prefix=f"gen_layers.{i}",
                 )
                 for i in range(self.num_hidden_layers)
@@ -1148,10 +1194,125 @@ class Cosmos3VFMTransformer(nn.Module):
         self.cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None
         self.cached_freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None
 
+        self._model_cpu_offload_enabled = False
+        self._model_cpu_offload_device: torch.device | None = None
+        self._model_cpu_offload_mover: SequentialOffloadHook | None = None
+        self._active_model_cpu_offload_component: str | None = None
+
     @property
     def device(self) -> torch.device:
+        offload_device = getattr(self, "_model_cpu_offload_device", None)
+        if getattr(self, "_model_cpu_offload_enabled", False) and offload_device is not None:
+            return offload_device
         return next(self.parameters()).device
 
+    def _model_cpu_offload_components(self) -> dict[str, list[nn.Module]]:
+        """Cosmos3's mutually-exclusive reasoner/generator component sets."""
+        return {
+            "reasoner": [self.language_model.layers],
+            "generator": [self.gen_layers],
+        }
+
+    def _model_cpu_offload_component_tensor_ids(self) -> set[int]:
+        component_tensors: set[int] = set()
+        for modules in self._model_cpu_offload_components().values():
+            for module in modules:
+                component_tensors.update(id(param) for param in module.parameters())
+                component_tensors.update(id(buffer) for buffer in module.buffers())
+        return component_tensors
+
+    def _move_model_cpu_offload_residents(self) -> None:
+        """Keep non-reasoner/non-generator weights resident on the target device."""
+        component_tensors = self._model_cpu_offload_component_tensor_ids()
+        device = self._model_cpu_offload_device
+        if device is None:
+            return
+        with torch.no_grad():
+            for param in self.parameters():
+                if id(param) not in component_tensors and param.data.device != device:
+                    param.data = param.data.to(device, non_blocking=False)
+            for buffer in self.buffers():
+                if id(buffer) not in component_tensors and buffer.device != device:
+                    buffer.data = buffer.data.to(device, non_blocking=False)
+
+    def _offload_model_cpu_component(self, name: str) -> None:
+        mover = self._model_cpu_offload_mover
+        if mover is None:
+            raise RuntimeError("Cosmos3 model CPU offload is not enabled")
+        for module in self._model_cpu_offload_components()[name]:
+            mover._to_cpu(module)
+
+    def _load_model_cpu_component(self, name: str) -> None:
+        mover = self._model_cpu_offload_mover
+        if mover is None:
+            raise RuntimeError("Cosmos3 model CPU offload is not enabled")
+        for module in self._model_cpu_offload_components()[name]:
+            mover._to_gpu(module)
+
+    def enable_model_cpu_offload(
+        self,
+        *,
+        device: torch.device,
+        pin_memory: bool = True,
+        use_hsdp: bool = False,
+    ) -> None:
+        """Enable Cosmos3 reasoner/generator CPU swapping inside ``forward``."""
+        if getattr(self, "_model_cpu_offload_enabled", False):
+            return
+
+        from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook
+
+        self._model_cpu_offload_device = torch.device(device)
+        self._model_cpu_offload_mover = SequentialOffloadHook(
+            offload_targets=[],
+            device=self._model_cpu_offload_device,
+            pin_memory=pin_memory,
+            use_hsdp=use_hsdp,
+        )
+        self._move_model_cpu_offload_residents()
+        for name in self._model_cpu_offload_components():
+            self._offload_model_cpu_component(name)
+        self._model_cpu_offload_enabled = True
+        self._active_model_cpu_offload_component = None
+        logger.info("Cosmos3 component-level CPU offload enabled on %s", self._model_cpu_offload_device)
+
+    def disable_model_cpu_offload(self) -> None:
+        if not getattr(self, "_model_cpu_offload_enabled", False):
+            return
+        for name in self._model_cpu_offload_components():
+            self._load_model_cpu_component(name)
+        if self._model_cpu_offload_device is not None and self._model_cpu_offload_device.type != "cpu":
+            current_omni_platform.synchronize()
+        self._model_cpu_offload_enabled = False
+        self._model_cpu_offload_device = None
+        self._model_cpu_offload_mover = None
+        self._active_model_cpu_offload_component = None
+
+    def _activate_model_cpu_offload_component(self, name: str) -> None:
+        if not getattr(self, "_model_cpu_offload_enabled", False):
+            return
+        components = self._model_cpu_offload_components()
+        if name not in components:
+            raise ValueError(f"Unknown Cosmos3 offload component: {name!r} (known: {list(components)})")
+        if self._active_model_cpu_offload_component == name:
+            return
+        self._active_model_cpu_offload_component = None
+        for other in components:
+            if other != name:
+                self._offload_model_cpu_component(other)
+        self._load_model_cpu_component(name)
+        self._active_model_cpu_offload_component = name
+
+    @contextmanager
+    def _model_cpu_offload_context(self, name: str) -> Iterator[None]:
+        self._activate_model_cpu_offload_component(name)
+        yield
+
+    def _offload_context(self, name: str) -> AbstractContextManager[None]:
+        if not getattr(self, "_model_cpu_offload_enabled", False):
+            return nullcontext()
+        return self._model_cpu_offload_context(name)
+
     # -- Patchify / Unpatchify -----------------------------------------------
 
     def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]:
@@ -1483,26 +1644,14 @@ class Cosmos3VFMTransformer(nn.Module):
         # Query Ulysses state at runtime
         ulysses_size, _, _ = _get_ulysses_state()
 
-        # Patchify latents and project to hidden space
-        hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w))
-        s_video = hidden_video.shape[1]
-        s_control = 0
-        hidden_controls: list[torch.Tensor] = []
-        for idx, control in enumerate(control_latent_list):
-            if control.shape != hidden_states.shape:
-                raise ValueError(
-                    "Cosmos3 transfer control latent shape must match target latent shape: "
-                    f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}."
-                )
-            hidden_control = self.proj_in(
-                self.patchify(control.to(device=hidden_states.device, dtype=hidden_states.dtype), t, h, w)
-            )
-            hidden_controls.append(hidden_control)
-            s_control += hidden_control.shape[1]
+        # Pack action/sound tokens (no learned weights) up front so the UND
+        # cache sizing knows their token lengths.  The modality projections are
+        # deferred into the generator offload context below, so model-level
+        # offload does not stage GEN weights before swapping to the reasoner.
+        action_tokens = None
+        sound_tokens = None
         s_action = 0
-        hidden_action = None
         s_sound = 0
-        hidden_sound = None
         if action_latents is not None:
             if action_latents.shape[0] != hidden_states.shape[0]:
                 raise ValueError(
@@ -1511,60 +1660,16 @@ class Cosmos3VFMTransformer(nn.Module):
                 )
             if action_domain_ids is None:
                 action_domain_ids = torch.zeros(action_latents.shape[0], dtype=torch.long, device=action_latents.device)
-            hidden_action = self.action_proj_in(self.pack_action(action_latents), action_domain_ids)
-            hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype)
-            s_action = hidden_action.shape[1]
+            action_tokens = self.pack_action(action_latents)
+            s_action = action_tokens.shape[1]
         if sound_latents is not None:
             if sound_latents.shape[0] != hidden_states.shape[0]:
                 raise ValueError(
                     "Cosmos3 sound and video batch sizes must match: "
                     f"video={hidden_states.shape[0]}, sound={sound_latents.shape[0]}."
                 )
-            hidden_sound = self.audio_proj_in(self.pack_sound(sound_latents))
-            hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype)
-            s_sound = hidden_sound.shape[1]
-
-        # Timestep embedding (fp32 for precision).
-        # For I2V: only add to noisy tokens, not conditioned ones.
-        # Conditioned frames are clean context and should not receive
-        # the diffusion timestep signal.
-        with torch.autocast(current_omni_platform.device_type, enabled=False):
-            time_embed = self.time_embedder((timestep * self.timestep_scale).float())
-        time_embed = time_embed.to(hidden_states.dtype)
-
-        if noisy_frame_mask is not None:
-            # Build per-token mask from per-frame mask.
-            # noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1]
-            token_noisy_mask = (
-                noisy_frame_mask[:, 0, :, 0, 0]  # [B, t]
-                .unsqueeze(-1)  # [B, t, 1]
-                .expand(-1, -1, hp * wp)  # [B, t, hp*wp]
-                .reshape(hidden_video.shape[0], -1, 1)  # [B, t*hp*wp, 1]
-            )
-            hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask
-        else:
-            hidden_video = hidden_video + time_embed.unsqueeze(1)
-
-        if hidden_action is not None:
-            if action_noisy_mask is None:
-                hidden_action = hidden_action + time_embed.unsqueeze(1)
-            else:
-                if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1):
-                    raise ValueError(
-                        "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], "
-                        f"got {tuple(action_noisy_mask.shape)}."
-                    )
-                action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device)
-                hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask
-
-        if hidden_sound is not None:
-            hidden_sound = hidden_sound + time_embed.unsqueeze(1)
-        hidden_parts = [*hidden_controls, hidden_video]
-        if hidden_action is not None:
-            hidden_parts.append(hidden_action)
-        if hidden_sound is not None:
-            hidden_parts.append(hidden_sound)
-        hidden_gen = torch.cat(hidden_parts, dim=1)
+            sound_tokens = self.pack_sound(sound_latents)
+            s_sound = sound_tokens.shape[1]
 
         # Run UND pathway once and cache K/V (replicated across all ranks)
         if self.cached_kv is None:
@@ -1583,90 +1688,162 @@ class Cosmos3VFMTransformer(nn.Module):
                 num_vision_items=len(control_latent_list) + 1,
                 share_vision_temporal_positions=transfer_share_vision_temporal_positions,
             )
-            cached_kv_full = self.language_model(text_ids, freqs_und)
+            with self._offload_context("reasoner"):
+                cached_kv_full = self.language_model(text_ids, freqs_und)
             self.cached_freqs_gen = freqs_gen
 
             # Trim to real text length (remove padding).  K/V stay replicated;
             # the framework Attention layer head-slices them via joint_key/value.
             self.cached_kv = [(k[:, :max_real_len], v[:, :max_real_len]) for k, v in cached_kv_full]
 
-        # Run GEN layers.  UND K/V (replicated) is passed to each layer;
-        # the Cosmos3CrossAttention forwards them as joint_key/value so the
-        # framework Attention handles the Ulysses head-slicing internally.
-        if self.cached_kv is None or self.cached_freqs_gen is None:
-            raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
-        self._validate_gen_sequence_parallel(
-            s_gen=hidden_gen.shape[1],
-            s_video=s_video,
-            s_control=s_control,
-            s_action=s_action,
-            s_sound=s_sound,
-            has_action=has_action,
-            has_sound=has_sound,
-            has_control=has_control,
-            ulysses_size=ulysses_size,
-        )
-        freqs_cos, freqs_sin = self.cached_freqs_gen
-        hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin)
-        freqs_gen = (freqs_cos, freqs_sin)
-
-        if len(self.gen_layers) == len(self.cached_kv):
-            for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True):
-                hidden_gen = layer(
-                    hidden_gen,
-                    k_und=k_und,
-                    v_und=v_und,
-                    freqs_cos=freqs_cos,
-                    freqs_sin=freqs_sin,
+        with self._offload_context("generator"):
+            # Patchify latents and project to hidden space after UND cache
+            # construction, so model-level offload does not stage GEN weights
+            # before immediately swapping to the reasoner.
+            hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w))
+            s_video = hidden_video.shape[1]
+            s_control = 0
+            hidden_controls: list[torch.Tensor] = []
+            for idx, control in enumerate(control_latent_list):
+                if control.shape != hidden_states.shape:
+                    raise ValueError(
+                        "Cosmos3 transfer control latent shape must match target latent shape: "
+                        f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}."
+                    )
+                hidden_control = self.proj_in(
+                    self.patchify(control.to(device=hidden_states.device, dtype=hidden_states.dtype), t, h, w)
                 )
-                # Cache-dit's block wrapper may return a tuple; unwrap it.
-                if isinstance(hidden_gen, tuple):
-                    hidden_gen = hidden_gen[0]
-        else:
-            # Cache-dit patches gen_layers to a grouped wrapper.
-            for layer in self.gen_layers:
-                hidden_gen = layer(
-                    hidden_gen,
-                    cached_kv=self.cached_kv,
-                    freqs_gen=freqs_gen,
+                hidden_controls.append(hidden_control)
+                s_control += hidden_control.shape[1]
+            hidden_action = None
+            hidden_sound = None
+            if action_tokens is not None:
+                assert action_domain_ids is not None
+                hidden_action = self.action_proj_in(action_tokens, action_domain_ids)
+                hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype)
+            if sound_tokens is not None:
+                hidden_sound = self.audio_proj_in(sound_tokens)
+                hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype)
+
+            # Timestep embedding (fp32 for precision).
+            # For I2V: only add to noisy tokens, not conditioned ones.
+            # Conditioned frames are clean context and should not receive
+            # the diffusion timestep signal.
+            with torch.autocast(current_omni_platform.device_type, enabled=False):
+                time_embed = self.time_embedder((timestep * self.timestep_scale).float())
+            time_embed = time_embed.to(hidden_states.dtype)
+
+            if noisy_frame_mask is not None:
+                # Build per-token mask from per-frame mask.
+                # noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1]
+                token_noisy_mask = (
+                    noisy_frame_mask[:, 0, :, 0, 0]  # [B, t]
+                    .unsqueeze(-1)  # [B, t, 1]
+                    .expand(-1, -1, hp * wp)  # [B, t, hp*wp]
+                    .reshape(hidden_video.shape[0], -1, 1)  # [B, t*hp*wp, 1]
                 )
-                if isinstance(hidden_gen, tuple):
-                    hidden_gen = hidden_gen[0]
+                hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask
+            else:
+                hidden_video = hidden_video + time_embed.unsqueeze(1)
 
-        hidden_gen = self.gen_sp_gather(hidden_gen)
+            if hidden_action is not None:
+                if action_noisy_mask is None:
+                    hidden_action = hidden_action + time_embed.unsqueeze(1)
+                else:
+                    if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1):
+                        raise ValueError(
+                            "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], "
+                            f"got {tuple(action_noisy_mask.shape)}."
+                        )
+                    action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device)
+                    hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask
 
-        # Final norm and project back to latent space
-        hidden_gen = self.norm_moe_gen(hidden_gen)
-        if not has_action and not has_sound and not has_control:
-            return self.unpatchify(self.proj_out(hidden_gen), t, h, w)
+            if hidden_sound is not None:
+                hidden_sound = hidden_sound + time_embed.unsqueeze(1)
+            hidden_parts = [*hidden_controls, hidden_video]
+            if hidden_action is not None:
+                hidden_parts.append(hidden_action)
+            if hidden_sound is not None:
+                hidden_parts.append(hidden_sound)
+            hidden_gen = torch.cat(hidden_parts, dim=1)
 
-        split_sizes = []
-        if has_control:
-            split_sizes.append(s_control)
-        split_sizes.append(s_video)
-        if has_action:
-            split_sizes.append(s_action)
-        if has_sound:
-            split_sizes.append(s_sound)
-        split_hidden = hidden_gen.split(split_sizes, dim=1)
-        split_idx = 0
-        if has_control:
+            # Run GEN layers.  UND K/V (replicated) is passed to each layer;
+            # the Cosmos3CrossAttention forwards them as joint_key/value so the
+            # framework Attention handles the Ulysses head-slicing internally.
+            if self.cached_kv is None or self.cached_freqs_gen is None:
+                raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.")
+            self._validate_gen_sequence_parallel(
+                s_gen=hidden_gen.shape[1],
+                s_video=s_video,
+                s_control=s_control,
+                s_action=s_action,
+                s_sound=s_sound,
+                has_action=has_action,
+                has_sound=has_sound,
+                has_control=has_control,
+                ulysses_size=ulysses_size,
+            )
+            freqs_cos, freqs_sin = self.cached_freqs_gen
+            hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin)
+            freqs_gen = (freqs_cos, freqs_sin)
+
+            if len(self.gen_layers) == len(self.cached_kv):
+                for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True):
+                    hidden_gen = layer(
+                        hidden_gen,
+                        k_und=k_und,
+                        v_und=v_und,
+                        freqs_cos=freqs_cos,
+                        freqs_sin=freqs_sin,
+                    )
+                    # Cache-dit's block wrapper may return a tuple; unwrap it.
+                    if isinstance(hidden_gen, tuple):
+                        hidden_gen = hidden_gen[0]
+            else:
+                # Cache-dit patches gen_layers to a grouped wrapper.
+                for layer in self.gen_layers:
+                    hidden_gen = layer(
+                        hidden_gen,
+                        cached_kv=self.cached_kv,
+                        freqs_gen=freqs_gen,
+                    )
+                    if isinstance(hidden_gen, tuple):
+                        hidden_gen = hidden_gen[0]
+
+            hidden_gen = self.gen_sp_gather(hidden_gen)
+
+            # Final norm and project back to latent space
+            hidden_gen = self.norm_moe_gen(hidden_gen)
+            if not has_action and not has_sound and not has_control:
+                return self.unpatchify(self.proj_out(hidden_gen), t, h, w)
+
+            split_sizes = []
+            if has_control:
+                split_sizes.append(s_control)
+            split_sizes.append(s_video)
+            if has_action:
+                split_sizes.append(s_action)
+            if has_sound:
+                split_sizes.append(s_sound)
+            split_hidden = hidden_gen.split(split_sizes, dim=1)
+            split_idx = 0
+            if has_control:
+                split_idx += 1
+            hidden_video = split_hidden[split_idx]
             split_idx += 1
-        hidden_video = split_hidden[split_idx]
-        split_idx += 1
-        video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w)
-        if has_control:
-            return video_pred
-        outputs: list[torch.Tensor] = [video_pred]
-        if has_action:
-            hidden_action = split_hidden[split_idx]
-            split_idx += 1
-            assert action_domain_ids is not None
-            outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids)))
-        if has_sound:
-            hidden_sound = split_hidden[split_idx]
-            outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound)))
-        return tuple(outputs)
+            video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w)
+            if has_control:
+                return video_pred
+            outputs: list[torch.Tensor] = [video_pred]
+            if has_action:
+                hidden_action = split_hidden[split_idx]
+                split_idx += 1
+                assert action_domain_ids is not None
+                outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids)))
+            if has_sound:
+                hidden_sound = split_hidden[split_idx]
+                outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound)))
+            return tuple(outputs)
 
     def post_load_weights(self) -> None:
         """Post-load processing: ensure correct dtypes."""

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py

@@ -0,0 +1,396 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cosmos3 Edge transformer variant with a Nemotron dense UND backbone."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
+
+from .transformer_cosmos3 import (
+    Cosmos3VFMTransformer,
+    Qwen3VLTextRotaryEmbedding,
+    RMSNorm,
+    _apply_rotary_pos_emb,
+    _nested_get,
+    _tf_config_get,
+)
+
+COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense"
+
+
+def _tf_config_nested_get(config: Any, key: str, default: Any = None) -> Any:
+    if config is None:
+        return default
+    if isinstance(config, dict):
+        value = _nested_get(config, key)
+    elif hasattr(config, "to_dict"):
+        value = _nested_get(config.to_dict(), key)
+    elif hasattr(config, "params"):
+        value = _nested_get(config.params, key)
+    else:
+        value = getattr(config, key, None)
+    return default if value is None else value
+
+
+class Cosmos3Relu2MLP(nn.Module):
+    """Nemotron dense MLP: down_proj(relu(up_proj(x)) ** 2)."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.up_proj = ColumnParallelLinear(
+            hidden_size,
+            intermediate_size,
+            bias=False,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            intermediate_size,
+            hidden_size,
+            bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        hidden = F.relu(self.up_proj(x))
+        return self.down_proj(hidden * hidden)
+
+
+class Cosmos3EdgeCausalAttention(nn.Module):
+    """Edge UND causal attention.
+
+    Reasoner self-attention uses raw UND Q/K with RoPE because
+    ``qk_norm_for_text=false``.  The GEN-facing UND key cache is a separate
+    normalized-then-RoPE'd view of the raw K tensor when GEN QK norm is enabled;
+    no-GEN-QK-norm checkpoints reuse the raw RoPE'd UND K.
+    """
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        rms_norm_eps: float,
+        use_k_norm_und_for_gen: bool,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_attention_heads
+        self.num_kv_heads = num_key_value_heads
+        self.head_dim = head_dim
+
+        tp_size = get_tensor_model_parallel_world_size()
+        self.num_heads_local = self.num_heads // tp_size
+        self.num_kv_heads_local = self.num_kv_heads // tp_size
+
+        self.to_q = ColumnParallelLinear(
+            hidden_size,
+            self.num_heads * self.head_dim,
+            bias=False,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_q",
+        )
+        self.to_k = ColumnParallelLinear(
+            hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=False,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_k",
+        )
+        self.to_v = ColumnParallelLinear(
+            hidden_size,
+            self.num_kv_heads * self.head_dim,
+            bias=False,
+            gather_output=False,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_v",
+        )
+        self.to_out = RowParallelLinear(
+            self.num_heads * self.head_dim,
+            hidden_size,
+            bias=False,
+            input_is_parallel=True,
+            return_bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.to_out",
+        )
+
+        self.k_norm_und_for_gen = RMSNorm(self.head_dim, eps=rms_norm_eps) if use_k_norm_und_for_gen else None
+        self.attn = FrameworkAttention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            causal=True,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            num_kv_heads=self.num_kv_heads,
+            skip_sequence_parallel=True,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        freqs_cos: torch.Tensor,
+        freqs_sin: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        B, S, _ = hidden_states.shape
+
+        q = self.to_q(hidden_states).view(B, S, self.num_heads_local, self.head_dim)
+        k = self.to_k(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
+        v = self.to_v(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim)
+
+        q_reasoner, k_reasoner = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
+
+        if self.k_norm_und_for_gen is not None:
+            k_gen = F.rms_norm(
+                k,
+                (self.head_dim,),
+                self.k_norm_und_for_gen.weight,
+                eps=self.k_norm_und_for_gen.variance_epsilon,
+            )
+            _, k_gen = _apply_rotary_pos_emb(q, k_gen, freqs_cos, freqs_sin)
+        else:
+            k_gen = k_reasoner
+
+        out = self.attn(q_reasoner, k_reasoner, v).reshape(B, S, -1)
+        return self.to_out(out), k_gen, v
+
+
+class Cosmos3EdgeUndDecoderLayer(nn.Module):
+    """Edge UND decoder layer: causal self-attention + ReLU2 MLP."""
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        intermediate_size: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        rms_norm_eps: float,
+        use_k_norm_und_for_gen: bool,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.self_attn = Cosmos3EdgeCausalAttention(
+            hidden_size=hidden_size,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            head_dim=head_dim,
+            rms_norm_eps=rms_norm_eps,
+            use_k_norm_und_for_gen=use_k_norm_und_for_gen,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps)
+        self.mlp = Cosmos3Relu2MLP(
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        freqs: tuple[torch.Tensor, torch.Tensor],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+
+        cos, sin = freqs
+        attn_out, k, v = self.self_attn(hidden_states, cos, sin)
+        hidden_states = residual + attn_out
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = residual + self.mlp(hidden_states)
+
+        return hidden_states, k, v
+
+
+class Cosmos3EdgeLanguageModel(nn.Module):
+    """Nemotron dense UND tower that returns GEN-facing K plus raw V per layer."""
+
+    _layerwise_offload_blocks_attrs = ["layers"]
+
+    def __init__(
+        self,
+        *,
+        hidden_size: int,
+        intermediate_size: int,
+        num_hidden_layers: int,
+        num_attention_heads: int,
+        num_key_value_heads: int,
+        head_dim: int,
+        vocab_size: int,
+        rms_norm_eps: float,
+        rope_theta: float,
+        mrope_section: list[int],
+        use_k_norm_und_for_gen: bool,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.embed_tokens = nn.Embedding(vocab_size, hidden_size)
+        self.rotary_emb = Qwen3VLTextRotaryEmbedding(
+            head_dim=head_dim,
+            rope_theta=rope_theta,
+            mrope_section=mrope_section,
+        )
+        self.layers = nn.ModuleList(
+            [
+                Cosmos3EdgeUndDecoderLayer(
+                    hidden_size=hidden_size,
+                    intermediate_size=intermediate_size,
+                    num_attention_heads=num_attention_heads,
+                    num_key_value_heads=num_key_value_heads,
+                    head_dim=head_dim,
+                    rms_norm_eps=rms_norm_eps,
+                    use_k_norm_und_for_gen=use_k_norm_und_for_gen,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{i}",
+                )
+                for i in range(num_hidden_layers)
+            ]
+        )
+        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
+
+    def forward(
+        self,
+        text_ids: torch.Tensor,
+        freqs: tuple[torch.Tensor, torch.Tensor],
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        hidden = self.embed_tokens(text_ids)
+
+        cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = []
+        for layer in self.layers:
+            hidden, k, v = layer(hidden, freqs)
+            cached_kv.append((k, v))
+
+        return cached_kv
+
+
+class Cosmos3EdgeVFMTransformer(Cosmos3VFMTransformer):
+    """Cosmos3 Edge variant with Nemotron dense UND and shared GEN diffusion."""
+
+    _language_model_cls = Cosmos3EdgeLanguageModel
+    _gen_mlp_cls = Cosmos3Relu2MLP
+
+    @staticmethod
+    def _validate_supported_config(model_config: Any) -> None:
+        expected_values = {
+            "qk_norm_for_text": False,
+            "position_embedding_type": "unified_3d_mrope",
+            "unified_3d_mrope_reset_spatial_ids": True,
+            "joint_attn_implementation": "two_way",
+        }
+        for key, expected in expected_values.items():
+            actual = _tf_config_get(model_config, key, expected)
+            if actual != expected:
+                raise ValueError(
+                    f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}."
+                )
+
+        backbone_type = _tf_config_get(model_config, "backbone_type", None)
+        if backbone_type is None:
+            raise ValueError(
+                f"Cosmos3 Edge transformer config must declare backbone_type={COSMOS3_EDGE_BACKBONE_TYPE!r}."
+            )
+        if backbone_type != COSMOS3_EDGE_BACKBONE_TYPE:
+            raise ValueError(
+                "Unsupported Cosmos3 Edge transformer config: "
+                f"backbone_type={backbone_type!r}; expected {COSMOS3_EDGE_BACKBONE_TYPE!r}."
+            )
+
+        required_values = {
+            "latent_channel": 48,
+            "latent_patch_size": 2,
+            "temporal_compression_factor": 4,
+        }
+        for key, expected in required_values.items():
+            actual = _tf_config_get(model_config, key, None)
+            if actual is None:
+                raise ValueError(f"Cosmos3 Edge transformer config must declare {key}={expected}.")
+            if int(actual) != expected:
+                raise ValueError(
+                    f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}."
+                )
+
+    @classmethod
+    def _resolve_rms_norm_eps(cls, model_config: Any) -> float:
+        return float(
+            _tf_config_get(
+                model_config,
+                "rms_norm_eps",
+                _tf_config_get(model_config, "layer_norm_epsilon", 1e-5),
+            )
+        )
+
+    @classmethod
+    def _resolve_rope_theta(cls, model_config: Any) -> float:
+        return float(_tf_config_get(model_config, "rope_theta", 100_000_000))
+
+    @classmethod
+    def _resolve_mrope_section(cls, model_config: Any) -> list[int]:
+        mrope_section = _tf_config_get(model_config, "mrope_section", None)
+        if mrope_section is None:
+            mrope_section = _tf_config_nested_get(model_config, "mrope_section", [24, 20, 20])
+        return list(mrope_section)
+
+    def _language_model_kwargs(self) -> dict[str, Any]:
+        return {"use_k_norm_und_for_gen": self.use_k_norm_und_for_gen}
+
+    def validate_loaded_weights(self, loaded: set[str]) -> None:
+        missing: list[str] = []
+        for layer_idx in range(self.num_hidden_layers):
+            required_markers = (
+                f"language_model.layers.{layer_idx}.mlp.up_proj.",
+                f"language_model.layers.{layer_idx}.mlp.down_proj.",
+                f"gen_layers.{layer_idx}.mlp.up_proj.",
+                f"gen_layers.{layer_idx}.mlp.down_proj.",
+            )
+            if self.use_k_norm_und_for_gen:
+                required_markers = (
+                    *required_markers,
+                    f"language_model.layers.{layer_idx}.self_attn.k_norm_und_for_gen.",
+                )
+            missing.extend(
+                marker.rstrip(".") for marker in required_markers if not any(marker in name for name in loaded)
+            )
+        if missing:
+            preview = ", ".join(missing[:10])
+            suffix = " ..." if len(missing) > 10 else ""
+            raise ValueError(
+                "Cosmos3 Edge transformer checkpoint is missing required weights for "
+                f"{preview}{suffix}. Use a converted Edge checkpoint with ReLU2 UND/GEN MLP weights."
+            )

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -24,6 +24,11 @@ from torch import nn
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import BasePipelineUtils, get_pipeline_utils
+from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import (
+    apply_diffusers_quantization_config,
+    convert_diffusers_quantization_config,
+    ensure_supported_diffusers_quantization,
+)
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniPromptType, OmniTextPrompt
@@ -31,6 +36,20 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
 
+_DIFFUSERS_CONFIG_LOAD_KWARGS = {
+    "cache_dir",
+    "dduf_entries",
+    "force_download",
+    "local_dir",
+    "local_dir_use_symlinks",
+    "local_files_only",
+    "proxies",
+    "revision",
+    "subfolder",
+    "token",
+    "user_agent",
+}
+
 
 class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
     """Black-box adapter that delegates full pipeline execution to a diffusers pipeline.
@@ -80,12 +99,19 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             "torch_dtype": dtype,
             **self.od_config.diffusers_load_kwargs,
         }
-        logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}")
+        convert_diffusers_quantization_config(load_kwargs)
 
         pipeline_class = self.od_config.diffusers_pipeline_cls
         pipeline_class_name = pipeline_class.__name__ if pipeline_class is not None else None
         self._pipeline_utils = get_pipeline_utils(pipeline_class_name)
         self._pipeline_utils.update_load_kwargs(self.od_config, load_kwargs)
+        component_names = (
+            self._load_diffusers_component_names(model_id, load_kwargs)
+            if self.od_config.quantization_config is not None and "quantization_config" not in load_kwargs
+            else {}
+        )
+        apply_diffusers_quantization_config(self.od_config, load_kwargs, component_names)
+        logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}")
 
         self._pipeline = DiffusionPipeline.from_pretrained(model_id, **load_kwargs)
         self._pipeline_utils.apply_post_load_updates(self._pipeline, self.od_config)
@@ -190,10 +216,38 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
                 "Eager execution is not supported with the diffusers backend. "
                 "Use a native pipeline for continuous batching mode."
             )
-        if self.od_config.quantization_config is not None:
-            raise NotImplementedError(
-                "Quantization is not supported with the diffusers backend. Use a native pipeline for quantization."
+        if (
+            self.od_config.quantization_config is not None
+            and "quantization_config" not in self.od_config.diffusers_load_kwargs
+        ):
+            ensure_supported_diffusers_quantization(self.od_config.quantization_config)
+            if self.od_config.diffusers_load_kwargs.get("dduf_file"):
+                raise NotImplementedError(
+                    "Diffusers backend quantization conversion does not support "
+                    "diffusers_load_kwargs.dduf_file yet. The preflight component "
+                    "discovery would need to mirror Diffusers' DDUF config loading. "
+                    "Use diffusers_load_kwargs.quantization_config for a native "
+                    "Diffusers quantization config, or omit dduf_file for vLLM-Omni "
+                    "quantization conversion."
+                )
+
+    def _load_diffusers_component_names(
+        self,
+        model_id: str,
+        load_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        config_load_kwargs = {k: load_kwargs[k] for k in _DIFFUSERS_CONFIG_LOAD_KWARGS if k in load_kwargs}
+        pipeline_config = DiffusionPipeline.load_config(model_id, **config_load_kwargs)
+        return {
+            name: value
+            for name, value in pipeline_config.items()
+            if (
+                isinstance(value, list)
+                and len(value) > 0
+                and value[0] is not None
+                and (name not in load_kwargs or load_kwargs[name] is not None)
             )
+        }
 
     # ------------------------------------------------------------------
     # Wrap settings, inputs, and outputs

vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py

@@ -0,0 +1,266 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantization helpers for the Diffusers backend."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from diffusers import PipelineQuantizationConfig, TorchAoConfig
+
+logger = logging.getLogger(__name__)
+
+_DIFFUSERS_TRANSFORMER_COMPONENTS = ("transformer", "transformer_2")
+_QuantizationConfigSpec = tuple[Callable[[Any], None], str]
+
+
+def ensure_supported_diffusers_quantization(quant_config: Any) -> None:
+    """Validate that a vLLM-Omni quantization config has a Diffusers mapping."""
+
+    method = _get_quant_method_name(quant_config)
+    validator, _ = _get_quantization_config_spec(method)
+    validator(quant_config)
+
+
+def convert_diffusers_quantization_config(load_kwargs: dict[str, Any]) -> None:
+    """Build Diffusers-native quantization config objects from dict kwargs."""
+
+    quant_config = load_kwargs.get("quantization_config")
+    if isinstance(quant_config, dict):
+        load_kwargs["quantization_config"] = _build_diffusers_native_quantization_config(quant_config)
+
+
+def apply_diffusers_quantization_config(
+    od_config: Any,
+    load_kwargs: dict[str, Any],
+    component_names: dict[str, Any],
+) -> bool:
+    """Inject a courtesy-converted quantization_config into load kwargs.
+
+    ``diffusers_load_kwargs`` is the canonical Diffusers backend configuration
+    path, so an explicit ``quantization_config`` already present in ``load_kwargs``
+    is never replaced. Returns whether a config was injected by this helper.
+    """
+
+    quant_config = getattr(od_config, "quantization_config", None)
+    if quant_config is None:
+        return False
+
+    if "quantization_config" in load_kwargs:
+        logger.warning(
+            "Both vLLM-Omni quantization_config and diffusers_load_kwargs.quantization_config "
+            "were provided for the diffusers backend. Using the Diffusers-native "
+            "quantization_config from diffusers_load_kwargs."
+        )
+        return False
+
+    transformer_components = _get_diffusers_transformer_components(component_names)
+    load_kwargs["quantization_config"] = _build_diffusers_quantization_config(
+        quant_config,
+        transformer_components,
+    )
+    return True
+
+
+def _normalize_method_name(method: Any) -> str:
+    method = getattr(method, "value", method)
+    return str(method).lower().replace("-", "_")
+
+
+def _get_quant_method_name(quant_config: Any) -> str:
+    get_name = getattr(quant_config, "get_name", None)
+    if get_name is not None:
+        return _normalize_method_name(get_name() if callable(get_name) else get_name)
+
+    method = getattr(quant_config, "quant_method", None)
+    if method is None:
+        method = getattr(quant_config, "method", None)
+    if method is None:
+        raise NotImplementedError(
+            "Diffusers backend quantization conversion requires a quantization "
+            "config with get_name(), quant_method, or method."
+        )
+    return _normalize_method_name(method)
+
+
+def _ensure_no_ignored_layers(quant_config: Any) -> None:
+    ignored_layers = getattr(quant_config, "ignored_layers", None)
+    if not ignored_layers:
+        ignored_layers = getattr(quant_config, "modules_to_not_convert", None)
+    if ignored_layers:
+        raise NotImplementedError(
+            "Diffusers backend quantization conversion does not map vLLM "
+            "ignored_layers/modules_to_not_convert names to Diffusers module "
+            "names. Use diffusers_load_kwargs for a native Diffusers config."
+        )
+
+
+def _get_torchao_quant_type_cls(class_name: str) -> type[Any]:
+    try:
+        import torchao.quantization as torchao_quantization
+    except ImportError as exc:
+        raise ImportError(
+            "Diffusers backend quantization conversion for fp8/int8 requires "
+            "torchao. Install torchao or pass a Diffusers-native "
+            "quantization_config through diffusers_load_kwargs."
+        ) from exc
+
+    try:
+        return getattr(torchao_quantization, class_name)
+    except AttributeError as exc:
+        raise ImportError(f"torchao.quantization.{class_name} is required for this quantization mapping.") from exc
+
+
+def _build_diffusers_native_quantization_config(config: dict[str, Any]) -> Any:
+    config = config.copy()
+    quant_mapping = config.get("quant_mapping")
+    if isinstance(quant_mapping, dict):
+        config["quant_mapping"] = {
+            component_name: _build_diffusers_component_quant_config(component_config)
+            for component_name, component_config in quant_mapping.items()
+        }
+    return PipelineQuantizationConfig(**config)
+
+
+def _build_diffusers_component_quant_config(config: Any) -> Any:
+    if not isinstance(config, dict):
+        return config
+
+    quant_method = config.get("quant_method")
+    if quant_method is None:
+        raise ValueError("Diffusers quant_mapping entries must provide quant_method.")
+
+    quant_method = _normalize_method_name(quant_method)
+    config_cls = _get_diffusers_quantization_config_cls(quant_method)
+    from_dict = getattr(config_cls, "from_dict", None)
+    if from_dict is not None:
+        return from_dict(config)
+
+    init_kwargs = config.copy()
+    init_kwargs.pop("quant_method", None)
+    return config_cls(**init_kwargs)
+
+
+def _get_diffusers_quantization_config_cls(quant_method: str) -> type[Any]:
+    from diffusers.quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING as DIFFUSERS_QUANT_CONFIG_MAPPING
+
+    config_cls = DIFFUSERS_QUANT_CONFIG_MAPPING.get(quant_method)
+    if config_cls is not None:
+        return config_cls
+
+    try:
+        from transformers.quantizers.auto import (
+            AUTO_QUANTIZATION_CONFIG_MAPPING as TRANSFORMERS_QUANT_CONFIG_MAPPING,
+        )
+    except ImportError:
+        TRANSFORMERS_QUANT_CONFIG_MAPPING = {}
+
+    config_cls = TRANSFORMERS_QUANT_CONFIG_MAPPING.get(quant_method)
+    if config_cls is None:
+        raise ValueError(f"Diffusers quant_mapping entry uses unknown quant_method={quant_method!r}.")
+    return config_cls
+
+
+def _build_torchao_pipeline_quant_config(
+    torchao_quant_type_name: str,
+    transformer_components: list[str],
+) -> Any:
+    """Build a Diffusers config after the caller has validated the vLLM config."""
+
+    quant_type_cls = _get_torchao_quant_type_cls(torchao_quant_type_name)
+    return PipelineQuantizationConfig(
+        quant_mapping={
+            component_name: TorchAoConfig(
+                quant_type=quant_type_cls(),
+            )
+            for component_name in transformer_components
+        }
+    )
+
+
+def _get_diffusers_transformer_components(component_names: dict[str, Any]) -> list[str]:
+    transformer_components = [
+        component_name for component_name in component_names if component_name in _DIFFUSERS_TRANSFORMER_COMPONENTS
+    ]
+    if transformer_components:
+        return transformer_components
+
+    available_components = ", ".join(sorted(component_names.keys())) or "<none>"
+    raise NotImplementedError(
+        "Diffusers backend quantization conversion currently only supports "
+        "pipelines with transformer or transformer_2 components. "
+        f"Found pipeline components: {available_components}. Use "
+        "diffusers_load_kwargs for a native Diffusers quantization config."
+    )
+
+
+def _validate_fp8_quant_config(quant_config: Any) -> None:
+    if getattr(quant_config, "is_checkpoint_fp8_serialized", False):
+        raise NotImplementedError(
+            "Diffusers backend fp8 conversion only supports online/dynamic "
+            "TorchAO quantization; serialized vLLM fp8 checkpoints are not mapped."
+        )
+
+    activation_scheme = getattr(quant_config, "activation_scheme", "dynamic")
+    if activation_scheme != "dynamic":
+        raise NotImplementedError(
+            f"Diffusers backend fp8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}."
+        )
+
+    weight_block_size = getattr(quant_config, "weight_block_size", None)
+    if weight_block_size is not None:
+        raise NotImplementedError(
+            "Diffusers backend fp8 conversion does not map vLLM weight_block_size "
+            "to TorchAO. Use diffusers_load_kwargs for a native Diffusers config."
+        )
+
+    _ensure_no_ignored_layers(quant_config)
+
+
+def _validate_int8_quant_config(quant_config: Any) -> None:
+    if getattr(quant_config, "is_checkpoint_int8_serialized", False):
+        raise NotImplementedError(
+            "Diffusers backend int8 conversion only supports online/dynamic "
+            "TorchAO quantization; serialized vLLM int8 checkpoints are not mapped."
+        )
+
+    activation_scheme = getattr(quant_config, "activation_scheme", "dynamic")
+    if activation_scheme != "dynamic":
+        raise NotImplementedError(
+            f"Diffusers backend int8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}."
+        )
+
+    _ensure_no_ignored_layers(quant_config)
+
+
+_QUANTIZATION_CONFIG_SPECS: dict[str, _QuantizationConfigSpec] = {
+    "fp8": (_validate_fp8_quant_config, "Float8DynamicActivationFloat8WeightConfig"),
+    "int8": (_validate_int8_quant_config, "Int8DynamicActivationInt8WeightConfig"),
+}
+
+
+def _get_quantization_config_spec(method: str) -> _QuantizationConfigSpec:
+    spec = _QUANTIZATION_CONFIG_SPECS.get(method)
+    if spec is None:
+        raise NotImplementedError(
+            f"Diffusers backend quantization conversion does not support {method!r}. "
+            "Use diffusers_load_kwargs for a native Diffusers quantization config, "
+            "or use a native vLLM-Omni pipeline for this quantization method."
+        )
+    return spec
+
+
+def _build_diffusers_quantization_config(
+    quant_config: Any,
+    transformer_components: list[str],
+) -> Any:
+    """Build a Diffusers PipelineQuantizationConfig from a supported config."""
+
+    method = _get_quant_method_name(quant_config)
+    _, torchao_quant_type_name = _get_quantization_config_spec(method)
+    return _build_torchao_pipeline_quant_config(
+        torchao_quant_type_name,
+        transformer_components,
+    )

vllm_omni/diffusion/models/dreamzero/causal_wan_model.py

@@ -37,6 +37,20 @@ from vllm_omni.diffusion.models.dreamzero.action_encoder import (
     MultiEmbodimentActionEncoder,
 )
 
+# AR-Diffusion paged self-attention (in-tree experimental engine). Import at
+# module level so the isinstance check + custom-op call trace cleanly inside
+# the fullgraph-compiled DiT block (an import inside the traced region would
+# graph-break). The model still works without the engine: the payload type is
+# only ever constructed by the AR-Diffusion runner.
+try:
+    from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
+        ARDiffusionPagedLayerInputs,
+        paged_write_attn,
+    )
+except ImportError:  # pragma: no cover - experimental package always ships in-tree
+    ARDiffusionPagedLayerInputs = None
+    paged_write_attn = None
+
 # ── RoPE utilities ──────────────────────────────────────────────────
 
 
@@ -510,7 +524,7 @@ class CausalWanSelfAttention(nn.Module):
         freqs_action: torch.Tensor,
         freqs_state: torch.Tensor,
         action_register_length: int | None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -567,24 +581,44 @@ class CausalWanSelfAttention(nn.Module):
             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)
+        if ARDiffusionPagedLayerInputs is not None and isinstance(kv_cache, ARDiffusionPagedLayerInputs):
+            # Fused write+attend custom op: one opaque node in the compiled
+            # graph (slot writes + FlashAttention block-table kernel inside).
+            # Metadata tensors were prepared once per forward in _forward_blocks.
+            if action_register_length is not None:
+                q_cat = torch.cat([roped_query, roped_action_query], dim=1)
+                k_act, v_act = roped_action_key[0], action_v[0]
+            else:
+                q_cat = roped_query
+                k_act = v_act = None
+            x = paged_write_attn(
+                kv_cache,
+                q_cat[0],
+                roped_key[0],
+                v[0],
+                k_act,
+                v_act,
+                self.head_dim**-0.5,
+            ).unsqueeze(0)
         else:
-            q_cat = roped_query
-            k_cat = new_k
-            v_cat = new_v
+            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 :]
 
-        x = self.attn(q_cat, k_cat, v_cat)
-        updated_kv_cache = torch.stack([new_k, new_v], dim=0)
+            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)
@@ -646,7 +680,7 @@ class CausalWanAttentionBlock(nn.Module):
         freqs_state: torch.Tensor,
         context: torch.Tensor,
         action_register_length: int | None = None,
-        kv_cache: torch.Tensor | None = None,
+        kv_cache: torch.Tensor | Any | None = None,
         crossattn_cache: dict | None = None,
         current_start_frame: int = 0,
         is_tf: bool = True,
@@ -925,10 +959,10 @@ class CausalWanModel(nn.Module):
         action: torch.Tensor | None,
         timestep_action: torch.Tensor | None,
         state: torch.Tensor | None,
-        kv_cache: list[torch.Tensor],
-        crossattn_cache: list[dict] | None,
+        kv_cache: list[Any],
         current_start_frame: int,
-    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]:
+        crossattn_cache: list[dict] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor | None]]:
         x = x.flatten(start_dim=2).transpose(1, 2)
         B = x.shape[0]
         F_t = timestep.shape[1]
@@ -966,7 +1000,28 @@ class CausalWanModel(nn.Module):
             clip_embedding = self.img_emb(clip_feature)
             context = torch.cat([clip_embedding, context], dim=1)
 
-        updated_kv_caches: list[torch.Tensor] = []
+        # AR-Diffusion paged path: host-side prep ONCE per branch forward, outside
+        # the compiled blocks — allocate current video/action slots, build the
+        # padded block-table metadata all 40 layers share, and hand the compiled
+        # region plain tensors (ARDiffusionPagedLayerInputs) instead of a Python
+        # context object. Lazy-allocation contract preserved: only the branch this
+        # CFG-parallel rank executes reaches its _forward_blocks.
+        if kv_cache and getattr(kv_cache[0], "is_ar_diffusion_paged_context", False):
+            if B != 1:
+                raise RuntimeError("AR-Diffusion paged self-attention currently expects batch_size=1")
+            fctx = kv_cache[0].forward_ctx
+            if seq_len != fctx.seq_len:
+                raise RuntimeError(
+                    f"AR-Diffusion paged context seq_len={fctx.seq_len} but current video KV has {seq_len} tokens"
+                )
+            fctx.prepare(
+                device=x.device,
+                action_len=int(action_register_length or 0),
+                query_len=int(x.shape[1]),
+            )
+            kv_cache = [c.to_layer_inputs() for c in kv_cache]
+
+        updated_kv_caches: list[torch.Tensor | None] = []
         for block_index, block in enumerate(self.blocks):
             x, updated_kv_cache = block(
                 x=x,
@@ -1000,9 +1055,9 @@ class CausalWanModel(nn.Module):
         timestep: torch.Tensor,
         context: torch.Tensor,
         seq_len: int,
-        kv_cache: list[torch.Tensor],
-        crossattn_cache: list[torch.Tensor],
+        kv_cache: list[Any],
         current_start_frame: int,
+        crossattn_cache: list[dict] | None = None,
         y: torch.Tensor | None = None,
         clip_feature: torch.Tensor | None = None,
         action: torch.Tensor | None = None,

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -3,7 +3,7 @@
 
 """DreamZero pipeline for vllm-omni.
 
-Entry point for DiffusionEngine.step() → pipeline.forward(req)
+Entry point for DiffusionEngine.step() -> pipeline.forward(req)
 """
 
 from __future__ import annotations
@@ -94,17 +94,102 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     Multi-output: predict_noise() returns (video_pred, action_pred).
     CFG: video gets standard CFG, action takes positive branch only.
-    State: DreamZeroState manages KV cache + frame buffer across forward() calls.
+
+    KV is managed by the AR-Diffusion engine: ``self._ar_diffusion_kv_state`` is set by the
+    runner before ``forward()`` and the pipeline routes all KV access (get / update /
+    commit / reset) through the pool-backed state. Purely duck-typed (no engine import).
     """
 
+    _ar_diffusion_kv_state = None  # set by the runner before each forward
+
+    def _kv_get(self, state, is_negative, seq_len=None, update_kv_cache=False):
+        return self._ar_diffusion_kv_state.get_kv_caches(
+            is_negative,
+            seq_len=seq_len,
+            commit_current=update_kv_cache,
+        )
+
+    def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim):
+        # The engine owns all KV allocation: self-attn is allocated lazily from
+        # paged contexts, and cross-attn is populated eagerly in _kv_populate_cross.
+        # Nothing is created model-side.
+        return
+
+    def _kv_commit(self, is_negative: bool):
+        self._ar_diffusion_kv_state.commit_paged_context(is_negative)
+
+    def _kv_get_cross(self, state, is_negative):
+        """Cross-attn cache from the engine pool (text k/v + I2V image k_img/v_img)."""
+        return self._ar_diffusion_kv_state.get_cross_kv_caches(is_negative)
+
+    def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None:
+        """Eagerly project cross-attn K/V for all layers into the AR-Diffusion pool.
+
+        Caches the session-invariant cross-attn projections once, per half: the text
+        ``k``/``v`` from ``text_embedding(context)`` survive window-boundary resets
+        (prompt unchanged within a session — only session resets clear them), while
+        the I2V image-token ``k_img``/``v_img`` from ``img_emb(clip_feature)`` (the
+        257 image tokens the forward splits off, cached model-side by #4154) are
+        re-projected on every window restart from the fresh CLIP features. Must run
+        after the image is encoded so ``clip_feature`` is available.
+        """
+        s = self._ar_diffusion_kv_state
+        need_text = not s._cross_text_populated.get(is_negative, False)
+        need_img = (
+            clip_feature is not None
+            and getattr(self.transformer, "model_type", "t2v") == "i2v"
+            and not s._cross_img_populated.get(is_negative, False)
+        )
+        if not need_text and not need_img:
+            return
+        projected = self.transformer.text_embedding(context) if need_text else None
+        img_ctx = self.transformer.img_emb(clip_feature) if need_img else None
+        for i, block in enumerate(self.transformer.blocks):
+            ca = block.cross_attn
+            n, d = ca.tp_num_heads, ca.head_dim
+            k = v = None
+            if projected is not None:
+                k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d))
+                v = ca.v(projected).unflatten(2, (n, d))
+            k_img = v_img = None
+            if img_ctx is not None:
+                k_img = ca.norm_k_img(ca.k_img(img_ctx)).unflatten(2, (n, d))
+                v_img = ca.v_img(img_ctx).unflatten(2, (n, d))
+            s.kv_cache.write_cross_kv(i, is_negative, k, v, k_img, v_img)
+        if need_text:
+            s._cross_text_populated[is_negative] = True
+        if need_img:
+            s._cross_img_populated[is_negative] = True
+        logger.info(
+            "AR-Diffusion CROSS POPULATE [%s]: %d layers, text=%s img=%s",
+            "neg" if is_negative else "pos",
+            len(self.transformer.blocks),
+            "kept" if projected is None else tuple(context.shape),
+            None if img_ctx is None else tuple(img_ctx.shape),
+        )
+
+    def _kv_reset(self, state, *, clear_video_latents: bool = True):
+        """Reset the engine's pooled session window plus the model's non-KV state.
+
+        DreamZero resets at the attention-window boundary; the engine pool drops the
+        same window so the next forward starts fresh. ``clear_video_latents=False``
+        keeps the accumulated video latents for export.
+
+        ``clear_video_latents=False`` also marks a window ("inference") reset: the
+        prompt is unchanged, so the pool keeps the text cross-attn K/V and only the
+        image half repopulates on the restart forward.
+        """
+        state.reset(clear_video_latents=clear_video_latents)
+        self._ar_diffusion_kv_state.reset(keep_cross_text=not clear_video_latents)
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         """Initialize pipeline components.
 
         DreamZero root checkpoint layout (GEAR-Dreams/DreamZero-DROID):
-          config.json                     — root config (action_head_cfg, architectures, etc.)
-          model-*.safetensors             — all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
-          experiment_cfg/metadata.json    — per-embodiment action normalization stats
-          vae/                            — symlink to Wan2.1 VAE (diffusers-compatible)
+          config.json                     -- root config (action_head_cfg, architectures, etc.)
+          model-*.safetensors             -- all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*)
+          experiment_cfg/metadata.json    -- per-embodiment action normalization stats
+          vae/                            -- symlink to Wan2.1 VAE (diffusers-compatible)
 
         Components are instantiated from config (not from_pretrained), then filled
         by load_weights() which reads root safetensors and remaps key prefixes.
@@ -119,6 +204,18 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         """
         super().__init__()
 
+        # DreamZero is engine-only: every KV access in forward() routes through
+        # the AR-Diffusion engine's pool-backed state. Fail fast here — a stale
+        # or programmatic config that leaves engine_backend="default" would
+        # otherwise only crash mid-forward on the first KV access.
+        engine_backend = str(getattr(od_config, "engine_backend", "") or "")
+        if "ar_diffusion" not in engine_backend.lower().replace("-", "_"):
+            raise ValueError(
+                "DreamZeroPipeline requires the AR-Diffusion engine; set "
+                "engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine "
+                f"in the deploy config (got engine_backend={engine_backend!r})."
+            )
+
         model_path = od_config.model
         model_config = od_config.model_config
         local_files_only = os.path.exists(model_path)
@@ -229,8 +326,11 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         self.max_action_dim: int = ah_config["max_action_dim"]
 
         self.negative_prompt: str = model_config.get("negative_prompt", DEFAULT_NEGATIVE_PROMPT)
+        # The negative prompt is a model constant: encode it once, lazily, and
+        # reuse across every forward/session (UMT5 encode is deterministic).
+        self._negative_prompt_embeds_cache: torch.Tensor | None = None
 
-        # Embodiment name → numeric ID mapping (model knowledge)
+        # Embodiment name -> numeric ID mapping (model knowledge)
         self.embodiment_name_to_id: dict[str, int] = model_config.get(
             "embodiment_name_to_id",
             DEFAULT_EMBODIMENT_NAME_TO_ID,
@@ -345,11 +445,17 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             state=kwargs.get("state_features"),
             embodiment_id=kwargs.get("embodiment_id"),
         )
-        if kwargs.get("update_kv_cache", False) and updated_kv_caches:
-            state = kwargs.get("dreamzero_state", self.state)
+        if kwargs.get("update_kv_cache", False):
             is_neg = kwargs.get("is_negative", False)
-            for i, kv in enumerate(updated_kv_caches):
-                state.update_kv_cache(i, kv, is_negative=is_neg)
+            logger.debug(
+                "AR-Diffusion pipeline predict_noise -> commit paged context: "
+                "is_neg=%s seq_len=%s current_start_frame=%s layers=%d",
+                is_neg,
+                kwargs.get("seq_len"),
+                kwargs.get("current_start_frame"),
+                len(updated_kv_caches),
+            )
+            self._kv_commit(is_neg)
 
         return video_pred, action_pred
 
@@ -363,6 +469,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False``
         on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and
         compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop).
+        Incremental VAE encode (``_vae_encode_encoder_chunk``) stays eager because
+        Wan ``feat_cache`` mutation is incompatible with CUDAGraph capture.
         DiT blocks use per-block ``fullgraph=True``.
         """
         if not torch.cuda.is_available():
@@ -377,6 +485,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
         compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False}
         # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead.
+        # The AR-Diffusion paged self-attention is a registered custom op, so the
+        # block stays fullgraph even on that path.
         dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False}
 
         logger.info(
@@ -397,9 +507,9 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc)
 
         try:
-            self.vae.encode = torch.compile(self.vae.encode, **compile_ro)
+            self.vae._encode = torch.compile(self.vae._encode, **compile_ro)
         except Exception as exc:
-            logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc)
+            logger.warning("DreamZero: vae._encode compile failed (%s); skipping.", exc)
 
         compiled_blocks = 0
         for block in self.transformer.blocks:
@@ -438,10 +548,24 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
             try:
                 image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device)
-                self._encode_image(image, self.num_frames, 180, 320)
+                self._encode_image(image, self.num_frames, 180, 320, state=self.state)
             except Exception as exc:
                 logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc)
 
+            try:
+                self.state.reset_vae_encoder_stream()
+                dummy_video = torch.zeros(1, 3, 1, 180, 320, dtype=torch.bfloat16, device=device)
+                self._vae_stream_seed(self.state, dummy_video[:, :, :1])
+                for _ in range(4):
+                    self._vae_stream_append_frame(self.state, dummy_video[:, :, :1])
+                self._vae_stream_get_observation_latents(
+                    self.state,
+                    self.num_frame_per_block,
+                    dtype=torch.bfloat16,
+                )
+            except Exception as exc:
+                logger.warning("DreamZero compile warmup (vae encode stream) skipped: %s", exc)
+
             try:
                 latent_h, latent_w = 180 // 8, 320 // 8
                 dummy_latent = torch.zeros(
@@ -500,7 +624,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     # -----------------------------------------------------------------------
 
     def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor:
-        """uint8 [B,T,H,W,C] → bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
+        """uint8 [B,T,H,W,C] -> bfloat16 [B,C,T,H,W] normalized to [-1,1]."""
         videos = videos.permute(0, 4, 1, 2, 3)
         if videos.dtype == torch.uint8:
             videos = videos.float() / 255.0
@@ -547,6 +671,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         num_frames: int,
         height: int,
         width: int,
+        *,
+        state: DreamZeroState | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Encode first frame via CLIP + VAE.
         Returns: (clip_feas, ys, image_latent)
@@ -582,11 +708,147 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             new_image = y[:, :, 0:1]
             y = torch.concat([msk, y], dim=1)
 
+            if state is not None:
+                if not state.vae_stream_initialized:
+                    # Seed the AR streaming encoder after the compiled full encode above;
+                    # ``cudagraph_mark_step_begin`` isolates eager feat_cache work from CUDAGraph.
+                    self._vae_stream_seed(state, image_input[:, :, :1])
+                else:
+                    # Window ("inference") restart with a live stream: keep the real
+                    # frame history in the Wan feat_cache and append the restart
+                    # observation instead of reseeding from scratch.
+                    self._cudagraph_mark_step_begin()
+                    self._vae_stream_append_frame(state, image_input[:, :, :1])
+
         return clip_context, y, new_image
 
+    def _vae_patchify(self, videos: torch.Tensor) -> torch.Tensor:
+        if self.vae.config.patch_size is not None:
+            from diffusers.models.autoencoders.autoencoder_kl_wan import patchify
+
+            return patchify(videos, patch_size=self.vae.config.patch_size)
+        return videos
+
+    @staticmethod
+    def _vae_clone_feat_map(feat_map: list[torch.Tensor | None]) -> list[torch.Tensor | None]:
+        return [entry.clone() if isinstance(entry, torch.Tensor) else entry for entry in feat_map]
+
+    def _vae_init_enc_feat_map(self) -> list[torch.Tensor | None]:
+        self.vae.clear_cache()
+        return self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_encode_encoder_chunk(
+        self,
+        chunk: torch.Tensor,
+        feat_map: list[torch.Tensor | None],
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """Run one Wan encoder chunk while mutating ``feat_map`` in place.
+
+        Must stay eager (not ``torch.compile``): Wan causal ``feat_cache`` updates
+        conflict with ``reduce-overhead`` CUDAGraph capture.
+        """
+        self.vae._enc_feat_map = feat_map
+        self.vae._enc_conv_idx = [0]
+        out = self.vae.encoder(
+            chunk,
+            feat_cache=self.vae._enc_feat_map,
+            feat_idx=self.vae._enc_conv_idx,
+        )
+        return out, self._vae_clone_feat_map(self.vae._enc_feat_map)
+
+    def _vae_quantize_encoder_out(self, encoder_out: torch.Tensor) -> torch.Tensor:
+        enc = self.vae.quant_conv(encoder_out)
+        mu, _ = enc.chunk(2, dim=1)
+        mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
+        inv_std = self.vae_latents_inv_std.to(device=mu.device, dtype=mu.dtype)
+        return (mu - mean) * inv_std
+
+    def _vae_stream_seed(self, state: DreamZeroState, first_frame: torch.Tensor) -> None:
+        """Seed incremental VAE encode with the first observation frame."""
+        state.reset_vae_encoder_stream()
+        feat_map = self._vae_init_enc_feat_map()
+        chunk = self._vae_patchify(first_frame.to(dtype=self.vae.dtype))
+        self._cudagraph_mark_step_begin()
+        encoder_out, feat_map = self._vae_encode_encoder_chunk(chunk[:, :, :1], feat_map)
+        state.vae_enc_feat_map = feat_map
+        state.vae_encoder_out = encoder_out
+        state.vae_pending_body_frames = None
+        state.vae_stream_initialized = True
+
+    def _vae_stream_append_frame(self, state: DreamZeroState, new_frame: torch.Tensor) -> None:
+        """Append one pixel frame and encode a 4-frame body chunk when ready."""
+        if not state.vae_stream_initialized or state.vae_enc_feat_map is None or state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream is not initialized.")
+
+        frame = new_frame.to(dtype=self.vae.dtype)
+        if state.vae_pending_body_frames is None:
+            state.vae_pending_body_frames = frame
+        else:
+            state.vae_pending_body_frames = torch.cat([state.vae_pending_body_frames, frame], dim=2)
+
+        while state.vae_pending_body_frames is not None and state.vae_pending_body_frames.shape[2] >= 4:
+            body = state.vae_pending_body_frames[:, :, :4]
+            if state.vae_pending_body_frames.shape[2] > 4:
+                state.vae_pending_body_frames = state.vae_pending_body_frames[:, :, 4:]
+            else:
+                state.vae_pending_body_frames = None
+            chunk = self._vae_patchify(body)
+            encoder_chunk, feat_map = self._vae_encode_encoder_chunk(chunk, state.vae_enc_feat_map)
+            state.vae_encoder_out = torch.cat([state.vae_encoder_out, encoder_chunk], dim=2)
+            state.vae_enc_feat_map = feat_map
+
+    def _vae_stream_get_observation_latents(
+        self,
+        state: DreamZeroState,
+        num_latent_frames: int,
+        *,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if state.vae_encoder_out is None:
+            raise RuntimeError("VAE encoder stream has no accumulated encoder output.")
+        latents = self._vae_quantize_encoder_out(state.vae_encoder_out).to(dtype=dtype)
+        if latents.shape[2] >= num_latent_frames:
+            return latents[:, :, -num_latent_frames:]
+        pad_count = num_latent_frames - latents.shape[2]
+        pad = latents[:, :, -1:].expand(-1, -1, pad_count, -1, -1)
+        return torch.cat([pad, latents], dim=2)
+
+    def _preprocess_vae_observation_window(self, videos: torch.Tensor) -> torch.Tensor:
+        _, _, num_frames_raw, _, _ = videos.shape
+        if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
+            return videos
+        if num_frames_raw // 4 != self.num_frame_per_block:
+            repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
+            videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
+            first_frame = videos[:, :, 0:1]
+            return torch.cat([first_frame, videos], dim=2)
+        first_frame = videos[:, :, 0:1]
+        return torch.cat([first_frame, videos], dim=2)
+
+    def _encode_observation_latents(
+        self,
+        state: DreamZeroState,
+        videos: torch.Tensor,
+        *,
+        latent_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        """Encode current robot observation into normalized VAE latents."""
+        if state.vae_stream_initialized:
+            self._cudagraph_mark_step_begin()
+            self._vae_stream_append_frame(state, videos[:, :, -1:])
+            return self._vae_stream_get_observation_latents(
+                state,
+                self.num_frame_per_block,
+                dtype=latent_dtype,
+            )
+
+        videos = self._preprocess_vae_observation_window(videos)
+        return self._encode_vae_latents(videos).to(dtype=latent_dtype)
+
     def _encode_vae_latents(self, videos: torch.Tensor) -> torch.Tensor:
         """Encode videos into normalized VAE latents."""
         input_dtype = videos.dtype
+        self._cudagraph_mark_step_begin()
         hidden = self.vae._encode(videos.to(dtype=self.vae.dtype))
         mu, _ = hidden.chunk(2, dim=1)
         mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype)
@@ -636,7 +898,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
     ) -> None:
         """Prefill KV cache with first frame and/or current observation.
 
-        Uses predict_noise_maybe_with_cfg() for CFG parallel — same path as
+        Uses predict_noise_maybe_with_cfg() for CFG parallel -- same path as
         the denoise loop. The mixin handles rank dispatch automatically.
         KV cache update happens as a side effect inside predict_noise().
         """
@@ -647,7 +909,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         head_dim = self.transformer.dim // self.transformer.num_heads
 
         if state.current_start_frame == 0:
-            state.create_kv_caches(
+            self._kv_create(
+                state,
                 batch_size,
                 dtype,
                 device,
@@ -672,16 +935,16 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                kv_cache=self._kv_get(state, False, seq_len=frame_seqlen, update_kv_cache=True),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    kv_cache=self._kv_get(state, True, seq_len=frame_seqlen, update_kv_cache=True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -722,16 +985,16 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             )
             positive_kwargs = dict(
                 encoder_hidden_states=prompt_embeds,
-                kv_cache=state.get_kv_caches(False),
-                crossattn_cache=state.get_crossattn_caches(False),
+                kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=True),
+                crossattn_cache=self._kv_get_cross(state, False),
                 is_negative=False,
                 **common,
             )
             negative_kwargs = (
                 dict(
                     encoder_hidden_states=negative_prompt_embeds,
-                    kv_cache=state.get_kv_caches(True),
-                    crossattn_cache=state.get_crossattn_caches(True),
+                    kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=True),
+                    crossattn_cache=self._kv_get_cross(state, True),
                     is_negative=True,
                     **common,
                 )
@@ -764,8 +1027,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
         For each timestep:
           1. Build positive_kwargs / negative_kwargs
-          2. predict_noise_maybe_with_cfg()    → (video_pred, action_pred)
-          3. scheduler_step_maybe_with_cfg()   → VideoActionScheduler
+          2. predict_noise_maybe_with_cfg()    -> (video_pred, action_pred)
+          3. scheduler_step_maybe_with_cfg()   -> VideoActionScheduler
           4. _synchronize_cfg_parallel_step_output()
         """
         seq_len = kwargs["seq_len"]
@@ -824,8 +1087,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                     hidden_states=noisy_input.transpose(1, 2),
                     timestep_video=timestep,
                     encoder_hidden_states=prompt_embeds,
-                    kv_cache=state.get_kv_caches(False),
-                    crossattn_cache=state.get_crossattn_caches(False),
+                    kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=False),
+                    crossattn_cache=self._kv_get_cross(state, False),
                     y=y,
                     clip_feature=state.clip_feas,
                     action=noisy_input_action,
@@ -839,8 +1102,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                         hidden_states=noisy_input.transpose(1, 2),
                         timestep_video=timestep,
                         encoder_hidden_states=negative_prompt_embeds,
-                        kv_cache=state.get_kv_caches(True),
-                        crossattn_cache=state.get_crossattn_caches(True),
+                        kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=False),
+                        crossattn_cache=self._kv_get_cross(state, True),
                         y=y,
                         clip_feature=state.clip_feas,
                         action=noisy_input_action,
@@ -937,7 +1200,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             device=device,
         )
 
-        # State: raw from transform → pad to (B, state_horizon=1, max_state_dim)
+        # State: raw from transform -> pad to (B, state_horizon=1, max_state_dim)
         raw_state = unified_obs["state"]
         state_for_postprocess = None
         if raw_state is not None:
@@ -977,38 +1240,66 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         # Explicit reset from OpenPI serving is carried by `extra_args["reset"]`
         # on the next inference request after websocket reset/session switch.
         if extra_args.get("reset", False):
-            state.reset()
+            self._kv_reset(state)
         else:
+            # Auto-reset based on model state (before accumulation). Both a "session"
+            # reset and an "inference" (window-boundary) reset clear the model-local
+            # KV window, so route both through _kv_reset to drop the AR-Diffusion pool window
+            # too and stay parity-exact; "inference" keeps the accumulated video
+            # latents (reset_inference_state == reset(clear_video_latents=False)).
             reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size)
             if reset_reason == "session":
-                state.reset()
+                self._kv_reset(state)
             elif reset_reason == "inference":
-                state.reset_inference_state()
+                self._kv_reset(state, clear_video_latents=False)
         state.language = text_tokens
 
-        # Frame accumulation: stitched single frame → multi-frame video
+        # Frame accumulation: stitched single frame -> multi-frame video
         video_frames = state.accumulate_frames(stitched)  # (T, H, W, C)
         videos = torch.from_numpy(video_frames).unsqueeze(0).to(device)  # (B=1, T, H, W, C)
 
-        videos = self._preprocess_video(videos)  # → [B,C,T,H,W] bf16
+        videos = self._preprocess_video(videos)  # -> [B,C,T,H,W] bf16
         _, _, num_frames_raw, height, width = videos.shape
 
-        prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        # Optional phase timing (DZ_PHASE_TIMING=1): logs per-forward stage costs
+        # (text encode / obs VAE encode / KV prefill / denoise) at INFO. Each mark
+        # synchronizes CUDA, so leave it off for timed benchmark runs.
+        _pt = None
+        if os.environ.get("DZ_PHASE_TIMING"):
+            import time as _time
+
+            torch.accelerator.synchronize()
+            _pt = {"time": _time, "t0": _time.perf_counter(), "marks": []}
+
+        def _pt_mark(name: str) -> None:
+            if _pt is not None:
+                torch.accelerator.synchronize()
+                _pt["marks"].append((name, _pt["time"].perf_counter()))
+
+        # Prompt embeds are constant within a session (a prompt change triggers a
+        # "session" reset above, which clears this cache alongside state.language).
+        if state.prompt_embeds is None:
+            state.prompt_embeds = self._encode_text(text_tokens, attention_mask)
+        prompt_embeds = state.prompt_embeds
         # Negative prompt for CFG uncond branch (model constant)
         negative_prompt_embeds = None
         if self.cfg_scale > 1.0:
-            neg_inputs = self.tokenizer(
-                self.negative_prompt,
-                max_length=512,
-                padding="max_length",
-                truncation=True,
-                return_tensors="pt",
-                add_special_tokens=True,
-            )
-            negative_prompt_embeds = self._encode_text(
-                neg_inputs["input_ids"].to(device),
-                neg_inputs["attention_mask"].to(device),
-            )
+            if self._negative_prompt_embeds_cache is None:
+                neg_inputs = self.tokenizer(
+                    self.negative_prompt,
+                    max_length=512,
+                    padding="max_length",
+                    truncation=True,
+                    return_tensors="pt",
+                    add_special_tokens=True,
+                )
+                self._negative_prompt_embeds_cache = self._encode_text(
+                    neg_inputs["input_ids"].to(device),
+                    neg_inputs["attention_mask"].to(device),
+                )
+            negative_prompt_embeds = self._negative_prompt_embeds_cache
+
+        _pt_mark("text_encode")
 
         # Extract first/last frame for CLIP + VAE encoding
         if num_frames_raw == 4 or num_frames_raw == 9:
@@ -1022,27 +1313,24 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                 self.num_frames,
                 height,
                 width,
+                state=state,
             )
             state.clip_feas = clip_feas.to(dtype=image.dtype)
             state.ys = ys.to(dtype=image.dtype)
 
-        if state.current_start_frame != 0:
-            # Subsequent calls: encode current observation via VAE
-            if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
-                pass
-            elif num_frames_raw // 4 != self.num_frame_per_block:
-                repeat_factor = self.num_frame_per_block // (num_frames_raw // 4)
-                videos = torch.repeat_interleave(videos, repeat_factor, dim=2)
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
-            else:
-                first_frame = videos[:, :, 0:1]
-                videos = torch.cat([first_frame, videos], dim=2)
+            # Eager cross-attn population (AR-Diffusion only): cache text + image-token K/V
+            # into the pool now that the image is encoded (clip_feas available).
+            # Runs on the first forward of a session and after each window-boundary
+            # reset (current_start_frame returns to 0); the populate guards (text/img halves) gate re-entry.
+            self._kv_populate_cross(prompt_embeds, state.clip_feas, is_negative=False)
+            if negative_prompt_embeds is not None:
+                self._kv_populate_cross(negative_prompt_embeds, state.clip_feas, is_negative=True)
 
+        if state.current_start_frame != 0:
             latent_dtype = videos.dtype
             with torch.no_grad():
-                image = self._encode_vae_latents(videos)
-            image = image.to(dtype=latent_dtype)
+                image = self._encode_observation_latents(state, videos, latent_dtype=latent_dtype)
+        _pt_mark("obs_vae_encode")
 
         batch_size = image.shape[0]
         generator = torch.Generator(device=device).manual_seed(self.seed)
@@ -1083,6 +1371,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             do_true_cfg,
             state,
         )
+        _pt_mark("prefill_kv")
 
         sample_scheduler = copy.deepcopy(self.scheduler)
         sample_scheduler_action = copy.deepcopy(self.scheduler)
@@ -1124,6 +1413,19 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
             state_features=state_features,
             embodiment_id=embodiment_id,
         )
+        if _pt is not None:
+            _pt_mark("diffuse")
+            prev_t = _pt["t0"]
+            parts = []
+            for name, t in _pt["marks"]:
+                parts.append(f"{name}={1000 * (t - prev_t):.1f}ms")
+                prev_t = t
+            logger.info(
+                "DZ_PHASE_TIMING csf=%s total=%.1fms %s",
+                state.current_start_frame,
+                1000 * (prev_t - _pt["t0"]),
+                " ".join(parts),
+            )
 
         if state.current_start_frame == 1:
             video_out = torch.cat([image, video_out], dim=1)
@@ -1134,7 +1436,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         # q99 denorm: [-1,1] → real values
         action_out = self._denormalize_action(action_out.float(), embodiment_name)
 
-        # Relative → absolute: only for relative_action_keys (joint_position only)
+        # Relative -> absolute: only for relative_action_keys (joint_position only)
         # gripper_position is NOT relative, so don't add state back to it
         if self.relative_action and state_for_postprocess is not None:
             n_relative = self.relative_action_dim  # 7 for DROID (joint only)
@@ -1146,7 +1448,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
                 action_out[..., :n_relative] + last_state.unsqueeze(1)  # broadcast over horizon
             )
 
-        # Squeeze batch dim for output: (B, horizon, dim) → (horizon, dim)
+        # Squeeze batch dim for output: (B, horizon, dim) -> (horizon, dim)
         actions_np = action_out.squeeze(0).float().cpu().numpy()  # (horizon, max_action_dim)
         actions_np = transform.transform_action_output(actions_np)
 
@@ -1246,7 +1548,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
         q01 = stats["q01"].to(device=action.device, dtype=action.dtype)
         q99 = stats["q99"].to(device=action.device, dtype=action.dtype)
         # action shape: (B, horizon, action_dim) or (B, horizon, max_action_dim)
-        # q01/q99 shape: (actual_action_dim,) — only denorm actual dims
+        # q01/q99 shape: (actual_action_dim,) -- only denorm actual dims
         actual_dim = q01.shape[0]
         action_real = action.clone()
         action_real[..., :actual_dim] = (action[..., :actual_dim] + 1) / 2 * (q99 - q01) + q01

vllm_omni/diffusion/models/dreamzero/state_dreamzero.py

@@ -115,17 +115,24 @@ class DreamZeroState:
         # Concatenated along T before decode (matches upstream video_across_time).
         self.video_latents_across_time = saved_video_latents
 
-        # KV cache once robot-policy diffusion supports that integration.
-        self.kv_cache = None
-        self.kv_cache_neg = None
-        self.crossattn_cache = None
-        self.crossattn_cache_neg = None
         self.current_start_frame = 0
 
         self.clip_feas = None
         self.ys = None
         if clear_video_latents:
+            # Session reset: drop the prompt-embed cache and VAE encoder stream.
             self.language = None
+            self.prompt_embeds = None
+            self.reset_vae_encoder_stream()
+        # Window ("inference") resets keep both: the prompt is unchanged, and the
+        # Wan feat_cache history is independent of the DiT attention window.
+
+    def reset_vae_encoder_stream(self) -> None:
+        """Clear incremental Wan VAE encoder state used across AR steps."""
+        self.vae_stream_initialized = False
+        self.vae_enc_feat_map: list[torch.Tensor | None] | None = None
+        self.vae_encoder_out: torch.Tensor | None = None
+        self.vae_pending_body_frames: torch.Tensor | None = None
 
     def reset_inference_state(self) -> None:
         """Reset KV/frame state after local attention rolls without dropping video latents."""
@@ -166,57 +173,3 @@ class DreamZeroState:
     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()."""
         return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None
-
-    # ------------------------------------------------------------------
-    # 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, "k_img": None, "v_img": None} for _ in range(num_layers)
-        ]
-        self.crossattn_cache_neg = [
-            {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": 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/video_export_worker.py

@@ -9,6 +9,35 @@ import torch
 class DreamZeroVideoExportWorkerExtension:
     """DreamZero worker RPCs used by offline example video export."""
 
+    def gpu_mem_stats(self) -> dict:
+        """Peak GPU memory (GiB) for this worker's device, for profiling.
+
+        ``max_memory_reserved`` is the caching-allocator high-water mark (the best
+        in-process proxy for "peak VRAM"); ``max_memory_allocated`` is the live
+        tensor high-water mark. Both are monotonic from process start, so the
+        end-of-run values capture the whole rollout's peak.
+        """
+        dev = torch.accelerator.current_device_index()
+        return {
+            "device": int(dev),
+            "peak_reserved_gib": torch.accelerator.max_memory_reserved(dev) / (1024**3),
+            "peak_allocated_gib": torch.accelerator.max_memory_allocated(dev) / (1024**3),
+        }
+
+    def ar_diffusion_perf_stats(self, reset: bool = True) -> dict:
+        """Per-request server-side forward E2E times (s) recorded by the AR-Diffusion runner.
+
+        The worker-side compute time per ``execute_model`` — the analog of the
+        upstream server's per-request E2E — which excludes the engine<->worker IPC
+        that the client's ``omni.generate`` wall time includes. Returns ``[]`` when
+        the runner is the base (KV-disabled) runner that does not record timings.
+        """
+        runner = self.model_runner
+        times = list(getattr(runner, "_perf_e2e_times", []) or [])
+        if reset and hasattr(runner, "_perf_e2e_times"):
+            runner._perf_e2e_times = []
+        return {"server_e2e_s": times}
+
     @staticmethod
     def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor:
         decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous()

vllm_omni/diffusion/models/flux/pipeline_flux.py

@@ -502,38 +502,14 @@ class FluxPipeline(
             return False
         return True
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        prompt_2: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        negative_prompt_2: str | list[str] | None = None,
-        true_cfg_scale: float = 1.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 28,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 3.5,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.FloatTensor | None = None,
-        pooled_prompt_embeds: torch.FloatTensor | None = None,
-        negative_prompt_embeds: torch.FloatTensor | None = None,
-        negative_pooled_prompt_embeds: torch.FloatTensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ) -> list[DiffusionOutput]:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         """Forward pass for flux."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
         sampling_params_list = req.sampling_params_list
         common_sampling_params = sampling_params_list[0]
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
+
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
@@ -541,40 +517,39 @@ class FluxPipeline(
 
         height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps
-        sigmas = common_sampling_params.sigmas or sigmas
-        guidance_scale = (
-            common_sampling_params.guidance_scale
-            if common_sampling_params.guidance_scale is not None
-            else guidance_scale
-        )
-        true_cfg_scale = common_sampling_params.true_cfg_scale or true_cfg_scale
+        num_inference_steps = common_sampling_params.num_inference_steps or 28
+        sigmas = common_sampling_params.sigmas
+        guidance_scale = common_sampling_params.guidance_scale
+        true_cfg_scale = common_sampling_params.true_cfg_scale or 1.0
         num_images_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        generator = req.collate_request_generators(num_images_per_prompt, generator)
-        latents = req.collate_request_tensors("latents", latents)
+        max_sequence_length = common_sampling_params.max_sequence_length or 512
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        latents = req.collate_request_tensors("latents", None)
         prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
             req.prompts,
             {
-                "prompt_embeds": prompt_embeds,
-                "negative_prompt_embeds": negative_prompt_embeds,
-                "pooled_prompt_embeds": pooled_prompt_embeds,
-                "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds,
+                "prompt_embeds": None,
+                "negative_prompt_embeds": None,
+                "pooled_prompt_embeds": None,
+                "negative_pooled_prompt_embeds": None,
             },
         )
         prompt_embeds = prompt_fields["prompt_embeds"]
         negative_prompt_embeds = prompt_fields["negative_prompt_embeds"]
         pooled_prompt_embeds = prompt_fields["pooled_prompt_embeds"]
         negative_pooled_prompt_embeds = prompt_fields["negative_pooled_prompt_embeds"]
+        prompt_2 = None
+        negative_prompt_2 = None
         if prompt_embeds is not None:
             prompt = None
-            prompt_2 = None
         if negative_prompt_embeds is not None:
             negative_prompt = None
-            negative_prompt_2 = None
+
+        output_type = common_sampling_params.output_type or "pil"
+        joint_attention_kwargs: dict[str, Any] | None = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
@@ -598,9 +573,7 @@ class FluxPipeline(
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
+        if prompt is not None:
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]

vllm_omni/diffusion/models/flux2/flux2_transformer.py

@@ -1034,7 +1034,7 @@ class Flux2Transformer2DModel(nn.Module):
             name = original_name
             if name not in params_dict and ".to_out.0." in name:
                 name = name.replace(".to_out.0.", ".to_out.")
-            # Some GGUF checkpoints include quantized tensors for modules that
+            # Some quantized checkpoints include tensors for modules that
             # are intentionally left unquantized in this model.
             param = params_dict.get(name)
             if param is None:

vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py

@@ -1044,7 +1044,7 @@ class Flux2Transformer2DModel(nn.Module):
             name = original_name
             if name not in params_dict and ".to_out.0." in name:
                 name = name.replace(".to_out.0.", ".to_out.")
-            # Some GGUF checkpoints include quantized tensors for modules that
+            # Some quantized checkpoints include tensors for modules that
             # are intentionally left unquantized in this model.
             param = params_dict.get(name)
             if param is None:

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -750,108 +750,7 @@ class Flux2KleinPipeline(
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
-        reference_image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
-        mask_image: PIL.Image.Image | list[PIL.Image.Image] | None = None,
-        prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        strength: float = 1.0,
-        guidance_scale: float | None = 4.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-        text_encoder_out_layers: tuple[int, ...] = (9, 18, 27),
-        padding_mask_crop: int | None = None,
-    ) -> DiffusionOutput:
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list of these):
-                `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
-                numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
-                or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
-                list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
-                latents as `image`, but if passing latents directly it is not encoded again.
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            guidance_scale (`float`, *optional*, defaults to 4.0):
-                Guidance scale as defined in [Classifier-Free Diffusion
-                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
-                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
-                `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
-                the text `prompt`, usually at the expense of lower image quality. For step-wise distilled models,
-                `guidance_scale` is ignored.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`List[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Note that "" is used as the negative prompt in this pipeline.
-                If not provided, will be generated from "".
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
-            text_encoder_out_layers (`Tuple[int]`):
-                Layer indices to use in the `text_encoder` to derive the final prompt embeddings.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.flux2.Flux2PipelineOutput`] or `tuple`: [`~pipelines.flux2.Flux2PipelineOutput`] if
-            `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
-            generated images.
-        """
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         if len(req.prompts) > 1:
             logger.warning(
                 """This model only supports a single prompt, not a batched request.""",
@@ -878,39 +777,32 @@ class Flux2KleinPipeline(
         else:
             image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else cast(PIL.Image.Image, raw_image)
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        height = req.sampling_params.height
+        width = req.sampling_params.width
         num_inference_steps = (
-            req.sampling_params.num_inference_steps
-            if req.sampling_params.num_inference_steps is not None
-            else num_inference_steps
+            50 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
         )
-        sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
-        generator = req.sampling_params.generator or generator
+        sigmas = req.sampling_params.sigmas
+        guidance_scale = req.sampling_params.guidance_scale
+        generator = req.sampling_params.generator
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers)
+        max_sequence_length = req.sampling_params.max_sequence_length or 512
+        latents = req.sampling_params.latents
+        strength = req.sampling_params.strength if req.sampling_params.strength is not None else 1.0
 
-        req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts]
-        if any(p is not None for p in req_prompt_embeds):
-            # If at list one prompt is provided as an embedding,
-            # Then assume that the user wants to provide embeddings for all prompts, and enter this if block
-            # If the user in fact provides mixed input format, req_prompt_embeds will have some None's
-            # And `torch.stack` automatically raises an exception for us
-            prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore # intentionally expect TypeError
-
-        req_negative_prompt_embeds = [
-            p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts
-        ]
-        if any(p is not None for p in req_negative_prompt_embeds):
-            negative_prompt_embeds = torch.stack(req_negative_prompt_embeds)  # type: ignore # intentionally expect TypeError
+        text_encoder_out_layers: tuple[int, ...] = req.sampling_params.extra_args.get(
+            "text_encoder_out_layers", (9, 18, 27)
+        )
+        padding_mask_crop: int | None = req.sampling_params.extra_args.get("padding_mask_crop")
+        attention_kwargs: dict[str, Any] | None = None
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None
+        callback_on_step_end_tensor_inputs = ["latents"]
+        output_type = req.sampling_params.output_type or "pil"
 
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(
@@ -930,12 +822,7 @@ class Flux2KleinPipeline(
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1
 
         device = self._execution_device
 

vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py

@@ -19,6 +19,7 @@ from transformers.utils import cached_file
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig
+from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
 from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag
 from vllm_omni.diffusion.models.gr00t.dataio.state_action.state_action_processor import StateActionProcessor
 from vllm_omni.diffusion.models.gr00t.dataio.utils import parse_modality_configs, to_json_serializable
@@ -552,4 +553,9 @@ class Gr00tN1d7Processor(ProcessorMixin):
         return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs)
 
 
-AutoProcessor.register("Gr00tN1d7", Gr00tN1d7Processor)
+# transformers >= 5.5 requires the config *class* (not a model_type string) as the
+# first arg to AutoProcessor.register — it does `key.__module__` internally. Passing
+# the string "Gr00tN1d7" raises AttributeError: 'str' object has no attribute
+# '__module__'. Match the AutoModel.register(Gr00tN1d7Config, ...) call in
+# modeling/gr00t_n1d7.py.
+AutoProcessor.register(Gr00tN1d7Config, Gr00tN1d7Processor)

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py

@@ -3,6 +3,8 @@
 
 from typing import Any
 
+import torch
+import vllm.distributed.parallel_state as _vllm_ps
 import vllm.forward_context as _vllm_fc
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
@@ -24,6 +26,30 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None:
         forward_context.in_profile_run = False
 
 
+def _set_forward_context_dp_metadata(num_tokens: int) -> None:
+    """Populate vLLM MoE DP token metadata for diffusion-only DP mappings."""
+    if not _vllm_fc.is_forward_context_available():
+        return
+    forward_context = _vllm_fc.get_forward_context()
+    if getattr(forward_context, "dp_metadata", None) is not None:
+        return
+
+    dp_group = getattr(_vllm_ps, "_DP", None)
+    if dp_group is None or getattr(dp_group, "world_size", 1) <= 1:
+        return
+
+    gathered_num_tokens: list[int | None] = [None] * dp_group.world_size
+    torch.distributed.all_gather_object(
+        gathered_num_tokens,
+        int(num_tokens),
+        group=dp_group.cpu_group,
+    )
+    if any(count is None for count in gathered_num_tokens):
+        raise RuntimeError(f"Failed to gather MoE DP token counts: {gathered_num_tokens}")
+
+    forward_context.dp_metadata = _vllm_fc.DPMetadata(torch.tensor(gathered_num_tokens, dtype=torch.int64))
+
+
 class HunyuanFusedMoEDefault:
     """Adapter that configures the upstream FusedMoE ``MoERunner`` for HunyuanImage3.
 
@@ -62,6 +88,7 @@ class HunyuanFusedMoEDefault:
                 hidden_states = args[0]
             if hidden_states is not None:
                 _set_forward_context_num_tokens(hidden_states.shape[0])
+                _set_forward_context_dp_metadata(hidden_states.shape[0])
 
         moe_runner.register_forward_pre_hook(_num_tokens_pre_hook, with_kwargs=True)
 

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -1099,6 +1099,12 @@ class ImageKVCacheManager:
 
         neg_kv = (key.reshape(-1, num_kv_heads, head_dim), value.reshape(-1, num_kv_heads, head_dim))
         self._injected_ar_kv = [self._injected_ar_kv[0], neg_kv]
+        logger.debug(
+            "[AR KV Reuse] neg branch: reused shared_prefix_len=%d, neg_kv_total_len=%d (q_len_actual=%d)",
+            shared_prefix_len,
+            key.shape[1],
+            q_len_actual,
+        )
         return key, value
 
     def __call__(
@@ -2235,13 +2241,8 @@ class HunyuanImage3Model(nn.Module):
             # processed with quantization, LoRA, fine-tuning, etc.
             if self.config.tie_word_embeddings and "lm_head.weight" in name:
                 continue
-            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
-                # Loading kv cache scales for compressed-tensors quantization
-                param = params_dict[scale_name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                loaded_weight = loaded_weight[0]
-                weight_loader(param, loaded_weight)
-                continue
+            # KV-cache scales are renamed via maybe_remap_kv_scale_name below;
+            # quant_config.get_cache_scale was removed in vLLM v0.23.0 (see #4810).
 
             is_found = False
             for param_name, weight_name, shard_id in stacked_params_mapping:
@@ -2924,6 +2925,10 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
     ):
         ar_kv_data = model_kwargs.pop("ar_kv_data", None)
         if ar_kv_data is None:
+            logger.debug(
+                "[AR KV Reuse] cfg_rank=%s: no AR KV received, fallback to full recompute (reuse_len=0)",
+                cfg_rank,
+            )
             return input_ids, 0
 
         # 1. positive prefix len

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -2209,7 +2209,6 @@ class HunyuanImage3Pipeline(
         if output_type == "latent":
             return DiffusionOutput(
                 output=latents,
-                custom_output={},
                 stage_durations=getattr(self, "stage_durations", None),
             )
 
@@ -2232,12 +2231,14 @@ class HunyuanImage3Pipeline(
         )
 
         cot_text_list = state.extra.get(_STEP_COT_TEXT_LIST) or []
-        custom_output = {}
+        metadata = {}
         if any(text is not None for text in cot_text_list):
-            custom_output["ar_generated_text"] = cot_text_list[0]
+            metadata["text"] = {"ar_generated_text": cot_text_list[0]}
         return DiffusionOutput(
-            output=image[0],
-            custom_output=custom_output,
+            output={
+                "payload": {"image": image[0]},
+                "metadata": metadata,
+            },
             stage_durations=getattr(self, "stage_durations", None),
         )
 
@@ -2302,12 +2303,14 @@ class HunyuanImage3Pipeline(
 
         outputs = self._generate(**model_inputs, **kwargs)
         image = outputs[0]
-        custom_output = {}
+        metadata = {}
         if any(t is not None for t in cot_text_list):
-            custom_output["ar_generated_text"] = cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list
+            metadata["text"] = {"ar_generated_text": cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list}
         stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None
         return DiffusionOutput(
-            output=image,
-            custom_output=custom_output,
+            output={
+                "payload": {"image": image},
+                "metadata": metadata,
+            },
             stage_durations=stage_durations,
         )

vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py

@@ -247,11 +247,13 @@ class InternVLAA1Pipeline(nn.Module, DiffusionPipelineProfilerMixin):
             noise=extra_args.get("noise"),
             decode_image=bool(extra_args.get("decode_image", False)),
         )
-        custom_output: dict[str, Any] = {}
+        metadata: dict[str, Any] = {}
         if decoded is not None:
-            custom_output["decoded"] = decoded
+            metadata["actions"] = {"decoded": decoded}
         return DiffusionOutput(
-            output=output,
-            custom_output=custom_output,
+            output={
+                "payload": {"actions": output},
+                "metadata": metadata,
+            },
             post_process_func=get_internvla_a1_post_process_func(self.od_config),
         )

vllm_omni/diffusion/models/krea2/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Krea 2 diffusion model components."""
+
+from vllm_omni.diffusion.models.krea2.krea2_transformer import Krea2Transformer2DModel
+from vllm_omni.diffusion.models.krea2.pipeline_krea2 import (
+    Krea2Pipeline,
+    get_krea2_post_process_func,
+)
+
+__all__ = [
+    "Krea2Pipeline",
+    "Krea2Transformer2DModel",
+    "get_krea2_post_process_func",
+]

vllm_omni/diffusion/models/krea2/krea2_transformer.py

@@ -0,0 +1,571 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.embeddings import apply_rotary_emb, get_1d_rotary_pos_embed
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention
+
+if TYPE_CHECKING:
+    from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+logger = init_logger(__name__)
+
+# Devices that do not support float64; RoPE frequencies fall back to float32 there.
+_FP64_UNSUPPORTED_DEVICES = frozenset({"mps", "npu", "neuron"})
+
+
+def _rope_freqs_dtype(device: torch.device) -> torch.dtype:
+    return torch.float32 if device.type in _FP64_UNSUPPORTED_DEVICES else torch.float64
+
+
+def _join_prefix(prefix: str, name: str) -> str:
+    return f"{prefix}.{name}" if prefix else name
+
+
+def _linear(
+    in_features: int,
+    out_features: int,
+    bias: bool,
+    quant_config: "QuantizationConfig | None",
+    prefix: str,
+) -> ReplicatedLinear:
+    # ReplicatedLinear (not nn.Linear) so the diffusion LoRA manager can wrap these projections.
+    return ReplicatedLinear(
+        in_features,
+        out_features,
+        bias=bias,
+        quant_config=quant_config,
+        prefix=prefix,
+        return_bias=False,
+    )
+
+
+class Krea2RMSNorm(nn.Module):
+    """RMSNorm with a zero-centered scale: the effective multiplier is ``1 + weight``, matching the Krea 2 checkpoint
+    format. Activations are upcast so the normalization runs in float32."""
+
+    def __init__(self, dim: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        dtype = hidden_states.dtype
+        hidden_states = F.rms_norm(hidden_states.float(), (self.dim,), weight=self.weight + 1.0, eps=self.eps)
+        return hidden_states.to(dtype)
+
+
+class Krea2Attention(nn.Module):
+    """Self-attention with grouped-query projections, q/k RMSNorm, rotary embeddings and a sigmoid output gate.
+
+    Q/K/V layout is ``[B, seq, heads, head_dim]``; ``attention_mask`` is a 2D boolean key-padding mask
+    ``(batch, key_seq_len)`` that the backend broadcasts.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int | None = None,
+        eps: float = 1e-5,
+        role: str = "self",
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size={hidden_size} must be divisible by num_heads={num_heads}")
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.head_dim = hidden_size // num_heads
+
+        q_dim = self.head_dim * self.num_heads
+        kv_dim = self.head_dim * self.num_kv_heads
+        self.to_q = _linear(hidden_size, q_dim, False, quant_config, _join_prefix(prefix, "to_q"))
+        self.to_k = _linear(hidden_size, kv_dim, False, quant_config, _join_prefix(prefix, "to_k"))
+        self.to_v = _linear(hidden_size, kv_dim, False, quant_config, _join_prefix(prefix, "to_v"))
+        self.to_gate = _linear(hidden_size, hidden_size, False, quant_config, _join_prefix(prefix, "to_gate"))
+        self.norm_q = Krea2RMSNorm(self.head_dim, eps=eps)
+        self.norm_k = Krea2RMSNorm(self.head_dim, eps=eps)
+        self.to_out = nn.ModuleList(
+            [_linear(hidden_size, hidden_size, False, quant_config, _join_prefix(prefix, "to_out.0")), nn.Dropout(0.0)]
+        )
+
+        self.attn = Attention(
+            num_heads=self.num_heads,
+            head_size=self.head_dim,
+            softmax_scale=1.0 / (self.head_dim**0.5),
+            causal=False,
+            num_kv_heads=self.num_kv_heads,
+            role=role,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        query = self.to_q(hidden_states).unflatten(-1, (self.num_heads, self.head_dim))
+        key = self.to_k(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim))
+        value = self.to_v(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim))
+        gate = self.to_gate(hidden_states)
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+            key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+        query = query.to(value.dtype)
+        key = key.to(value.dtype)
+
+        attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
+        hidden_states = self.attn(query, key, value, attn_metadata)
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states * torch.sigmoid(gate)
+        return self.to_out[0](hidden_states)
+
+
+class Krea2SwiGLU(nn.Module):
+    """SwiGLU feed-forward network."""
+
+    def __init__(
+        self, dim: int, hidden_dim: int, quant_config: "QuantizationConfig | None" = None, prefix: str = ""
+    ) -> None:
+        super().__init__()
+        self.gate = _linear(dim, hidden_dim, False, quant_config, _join_prefix(prefix, "gate"))
+        self.up = _linear(dim, hidden_dim, False, quant_config, _join_prefix(prefix, "up"))
+        self.down = _linear(hidden_dim, dim, False, quant_config, _join_prefix(prefix, "down"))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(hidden_states)) * self.up(hidden_states))
+
+
+class Krea2TextFusionBlock(nn.Module):
+    """Pre-norm transformer block (no rotary embeddings, no time modulation) used by the text fusion stage."""
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+        intermediate_size: int,
+        eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.norm1 = Krea2RMSNorm(dim, eps=eps)
+        self.norm2 = Krea2RMSNorm(dim, eps=eps)
+        self.attn = Krea2Attention(
+            dim,
+            num_heads,
+            num_kv_heads,
+            eps=eps,
+            role="self",
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
+        )
+        self.ff = Krea2SwiGLU(dim, intermediate_size, quant_config=quant_config, prefix=_join_prefix(prefix, "ff"))
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.norm1(hidden_states), attention_mask=attention_mask)
+        hidden_states = hidden_states + self.ff(self.norm2(hidden_states))
+        return hidden_states
+
+
+class Krea2TextFusion(nn.Module):
+    """Fuses the stack of tapped text-encoder hidden states into a single sequence of text features.
+
+    Two ``layerwise_blocks`` attend across the ``num_text_layers`` axis independently for every token, a linear
+    ``projector`` collapses that axis, and two ``refiner_blocks`` attend across the token sequence.
+    """
+
+    def __init__(
+        self,
+        num_text_layers: int,
+        dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+        intermediate_size: int,
+        num_layerwise_blocks: int,
+        num_refiner_blocks: int,
+        eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.layerwise_blocks = nn.ModuleList(
+            [
+                Krea2TextFusionBlock(
+                    dim,
+                    num_heads,
+                    num_kv_heads,
+                    intermediate_size,
+                    eps,
+                    quant_config=quant_config,
+                    prefix=_join_prefix(prefix, f"layerwise_blocks.{i}"),
+                )
+                for i in range(num_layerwise_blocks)
+            ]
+        )
+        self.projector = _linear(num_text_layers, 1, False, quant_config, _join_prefix(prefix, "projector"))
+        self.refiner_blocks = nn.ModuleList(
+            [
+                Krea2TextFusionBlock(
+                    dim,
+                    num_heads,
+                    num_kv_heads,
+                    intermediate_size,
+                    eps,
+                    quant_config=quant_config,
+                    prefix=_join_prefix(prefix, f"refiner_blocks.{i}"),
+                )
+                for i in range(num_refiner_blocks)
+            ]
+        )
+
+    def forward(self, encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        batch_size, seq_len, num_text_layers, dim = encoder_hidden_states.shape
+
+        hidden_states = encoder_hidden_states.reshape(batch_size * seq_len, num_text_layers, dim)
+        for block in self.layerwise_blocks:
+            hidden_states = block(hidden_states.contiguous())
+
+        hidden_states = hidden_states.reshape(batch_size, seq_len, num_text_layers, dim).permute(0, 1, 3, 2)
+        hidden_states = self.projector(hidden_states).squeeze(-1)
+
+        for block in self.refiner_blocks:
+            hidden_states = block(hidden_states, attention_mask=attention_mask)
+
+        return hidden_states
+
+
+class Krea2TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        norm_eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.scale_shift_table = nn.Parameter(torch.zeros(6, hidden_size))
+        self.norm1 = Krea2RMSNorm(hidden_size, eps=norm_eps)
+        self.norm2 = Krea2RMSNorm(hidden_size, eps=norm_eps)
+        self.attn = Krea2Attention(
+            hidden_size,
+            num_heads,
+            num_kv_heads,
+            eps=norm_eps,
+            role="self",
+            quant_config=quant_config,
+            prefix=_join_prefix(prefix, "attn"),
+        )
+        self.ff = Krea2SwiGLU(
+            hidden_size, intermediate_size, quant_config=quant_config, prefix=_join_prefix(prefix, "ff")
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # temb: (B, 1, 6 * hidden_size), shared across all blocks; each block only learns an additive table.
+        modulation = temb.unflatten(-1, (6, -1)) + self.scale_shift_table
+        prescale, preshift, pregate, postscale, postshift, postgate = modulation.unbind(-2)
+
+        attn_out = self.attn(
+            (1.0 + prescale) * self.norm1(hidden_states) + preshift,
+            attention_mask=attention_mask,
+            image_rotary_emb=image_rotary_emb,
+        )
+        hidden_states = hidden_states + pregate * attn_out
+        ff_out = self.ff((1.0 + postscale) * self.norm2(hidden_states) + postshift)
+        hidden_states = hidden_states + postgate * ff_out
+        return hidden_states
+
+
+class Krea2TimestepEmbedding(nn.Module):
+    """Sinusoidal flow-time embedding (cos-first, input scaled by 1000) followed by a two-layer MLP.
+
+    Keeps the sequence dimension at size 1 so the per-block modulations broadcast over tokens.
+    """
+
+    def __init__(
+        self, embed_dim: int, hidden_size: int, quant_config: "QuantizationConfig | None" = None, prefix: str = ""
+    ) -> None:
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.linear_1 = _linear(embed_dim, hidden_size, True, quant_config, _join_prefix(prefix, "linear_1"))
+        self.linear_2 = _linear(hidden_size, hidden_size, True, quant_config, _join_prefix(prefix, "linear_2"))
+
+    def forward(self, timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+        half = self.embed_dim // 2
+        freqs = torch.exp(-math.log(1e4) * torch.arange(half, dtype=torch.float32, device=timestep.device) / half)
+        args = (timestep.float() * 1e3)[:, None, None] * freqs
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(dtype)
+        return self.linear_2(F.gelu(self.linear_1(emb), approximate="tanh"))
+
+
+class Krea2TextProjection(nn.Module):
+    """Projects the fused text features into the transformer width."""
+
+    def __init__(
+        self,
+        text_dim: int,
+        hidden_size: int,
+        eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.norm = Krea2RMSNorm(text_dim, eps=eps)
+        self.linear_1 = _linear(text_dim, hidden_size, True, quant_config, _join_prefix(prefix, "linear_1"))
+        self.linear_2 = _linear(hidden_size, hidden_size, True, quant_config, _join_prefix(prefix, "linear_2"))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.linear_1(self.norm(hidden_states))
+        return self.linear_2(F.gelu(hidden_states, approximate="tanh"))
+
+
+class Krea2FinalLayer(nn.Module):
+    """Final adaptive RMSNorm and output projection."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        out_channels: int,
+        eps: float,
+        quant_config: "QuantizationConfig | None" = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.scale_shift_table = nn.Parameter(torch.zeros(2, hidden_size))
+        self.norm = Krea2RMSNorm(hidden_size, eps=eps)
+        self.linear = _linear(hidden_size, out_channels, True, quant_config, _join_prefix(prefix, "linear"))
+
+    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        modulation = temb + self.scale_shift_table
+        scale, shift = modulation.chunk(2, dim=1)
+        hidden_states = (1.0 + scale) * self.norm(hidden_states) + shift
+        return self.linear(hidden_states)
+
+
+class Krea2RotaryPosEmbed(nn.Module):
+    """Multi-axis (t, h, w) rotary position embedding, following the Flux/Krea 2 convention."""
+
+    def __init__(self, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        n_axes = ids.shape[-1]
+        cos_out = []
+        sin_out = []
+        pos = ids.float()
+        freqs_dtype = _rope_freqs_dtype(ids.device)
+        for i in range(n_axes):
+            cos, sin = get_1d_rotary_pos_embed(
+                self.axes_dim[i],
+                pos[:, i],
+                theta=self.theta,
+                repeat_interleave_real=True,
+                use_real=True,
+                freqs_dtype=freqs_dtype,
+            )
+            cos_out.append(cos)
+            sin_out.append(sin)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+class Krea2Transformer2DModel(nn.Module):
+    r"""The single-stream MMDiT flow-matching backbone used by the Krea 2 pipeline (vLLM-Omni port).
+
+    Text conditioning enters as a stack of hidden states tapped from several layers of a multimodal text encoder. A
+    small text-fusion transformer collapses the layer axis and refines the token sequence; the result is concatenated
+    with the patchified image latents into a single ``[text, image]`` sequence processed by the transformer blocks.
+    The timestep conditions every block through one shared modulation vector plus per-block learned tables.
+    """
+
+    _repeated_blocks = ["Krea2TransformerBlock"]
+    _layerwise_offload_blocks_attrs = ["transformer_blocks"]
+
+    @staticmethod
+    def _is_transformer_block(name: str, module) -> bool:
+        return "transformer_blocks" in name and name.split(".")[-1].isdigit()
+
+    _hsdp_shard_conditions = [_is_transformer_block]
+
+    def __init__(
+        self,
+        in_channels: int = 64,
+        num_layers: int = 28,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 48,
+        num_key_value_heads: int = 12,
+        intermediate_size: int = 16384,
+        timestep_embed_dim: int = 256,
+        text_hidden_dim: int = 2560,
+        num_text_layers: int = 12,
+        text_num_attention_heads: int = 20,
+        text_num_key_value_heads: int = 20,
+        text_intermediate_size: int = 6912,
+        num_layerwise_text_blocks: int = 2,
+        num_refiner_text_blocks: int = 2,
+        axes_dims_rope: tuple[int, int, int] = (32, 48, 48),
+        rope_theta: float = 1000.0,
+        norm_eps: float = 1e-5,
+        od_config: "OmniDiffusionConfig | None" = None,
+        quant_config: "QuantizationConfig | None" = None,
+    ) -> None:
+        super().__init__()
+        # The pipeline reads ``self.dtype`` to cast inputs before the transformer forward.
+        self.dtype = od_config.dtype if od_config is not None else torch.get_default_dtype()
+        self.od_config = od_config
+
+        hidden_size = attention_head_dim * num_attention_heads
+        if sum(axes_dims_rope) != attention_head_dim:
+            raise ValueError(
+                f"sum(axes_dims_rope)={sum(axes_dims_rope)} must equal attention_head_dim={attention_head_dim}"
+            )
+
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+
+        self.img_in = _linear(in_channels, hidden_size, True, quant_config, "img_in")
+        self.time_embed = Krea2TimestepEmbedding(timestep_embed_dim, hidden_size, quant_config, "time_embed")
+        self.time_mod_proj = _linear(hidden_size, 6 * hidden_size, True, quant_config, "time_mod_proj")
+        self.text_fusion = Krea2TextFusion(
+            num_text_layers=num_text_layers,
+            dim=text_hidden_dim,
+            num_heads=text_num_attention_heads,
+            num_kv_heads=text_num_key_value_heads,
+            intermediate_size=text_intermediate_size,
+            num_layerwise_blocks=num_layerwise_text_blocks,
+            num_refiner_blocks=num_refiner_text_blocks,
+            eps=norm_eps,
+            quant_config=quant_config,
+            prefix="text_fusion",
+        )
+        self.txt_in = Krea2TextProjection(text_hidden_dim, hidden_size, norm_eps, quant_config, "txt_in")
+        self.rotary_emb = Krea2RotaryPosEmbed(theta=rope_theta, axes_dim=list(axes_dims_rope))
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                Krea2TransformerBlock(
+                    hidden_size=hidden_size,
+                    intermediate_size=intermediate_size,
+                    num_heads=num_attention_heads,
+                    num_kv_heads=num_key_value_heads,
+                    norm_eps=norm_eps,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{i}",
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        self.final_layer = Krea2FinalLayer(hidden_size, in_channels, norm_eps, quant_config, "final_layer")
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        position_ids: torch.Tensor,
+        encoder_attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        r"""Predict the flow-matching velocity for the image tokens.
+
+        Args:
+            hidden_states: Packed (patchified) noisy image latents ``(batch, image_seq_len, in_channels)``.
+            encoder_hidden_states: Tapped text-encoder hidden states
+                ``(batch, text_seq_len, num_text_layers, text_hidden_dim)``.
+            timestep: Flow-matching time in ``[0, 1]`` of shape ``(batch,)``.
+            position_ids: ``(t, h, w)`` rotary coordinates ``(text_seq_len + image_seq_len, 3)``.
+            encoder_attention_mask: Boolean mask marking valid text tokens ``(batch, text_seq_len)`` or ``None``.
+        """
+        if position_ids.ndim != 2 or position_ids.shape[-1] != 3:
+            raise ValueError(f"`position_ids` must have shape (sequence_length, 3), got {tuple(position_ids.shape)}.")
+
+        batch_size, image_seq_len, _ = hidden_states.shape
+
+        temb = self.time_embed(timestep, dtype=hidden_states.dtype)
+        temb_mod = self.time_mod_proj(F.gelu(temb, approximate="tanh"))
+
+        # 2D key-padding masks: padded text tokens are excluded as keys; their lanes are dropped at the output slice.
+        text_attention_mask = None
+        attention_mask = None
+        if encoder_attention_mask is not None:
+            text_attention_mask = encoder_attention_mask.bool()
+            image_mask = text_attention_mask.new_ones((batch_size, image_seq_len))
+            attention_mask = torch.cat([text_attention_mask, image_mask], dim=1)
+
+        encoder_hidden_states = self.text_fusion(encoder_hidden_states, attention_mask=text_attention_mask)
+        encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+        text_seq_len = encoder_hidden_states.shape[1]
+        hidden_states = self.img_in(hidden_states)
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        image_rotary_emb = self.rotary_emb(position_ids)
+
+        for block in self.transformer_blocks:
+            hidden_states = block(hidden_states, temb_mod, image_rotary_emb, attention_mask)
+
+        hidden_states = hidden_states[:, text_seq_len:]
+        return self.final_layer(hidden_states, temb)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        # Module/parameter names mirror the diffusers ``Krea2Transformer2DModel`` exactly, so weights map 1:1.
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if name not in params_dict:
+                continue
+            param = params_dict[name]
+            weight_loader = getattr(param, "weight_loader", default_weight_loader)
+            weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params

vllm_omni/diffusion/models/krea2/pipeline_krea2.py

@@ -0,0 +1,551 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import logging
+import os
+from collections.abc import Iterable
+from typing import ClassVar
+
+import numpy as np
+import torch
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers import AutoTokenizer, Qwen3VLModel
+from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.transformers_utils.config import get_hf_file_to_dict
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
+from vllm_omni.diffusion.distributed.utils import get_local_device
+from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.models.krea2.krea2_transformer import Krea2Transformer2DModel
+from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs
+
+logger = logging.getLogger(__name__)
+
+# Text-encoder layer taps used when model_index.json does not specify ``text_encoder_select_layers``.
+DEFAULT_TEXT_ENCODER_SELECT_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
+
+
+def get_krea2_post_process_func(od_config: OmniDiffusionConfig):
+    model_name = od_config.model
+    # Read only the VAE config (``get_hf_file_to_dict`` resolves a local dir via
+    # ``Path(model)/file`` and a hub repo via a single-file download) instead of
+    # pulling all weights just to compute ``vae_scale_factor``.
+    vae_config = get_hf_file_to_dict("vae/config.json", model_name) or {}
+    vae_scale_factor = 2 ** len(vae_config["temperal_downsample"]) if "temperal_downsample" in vae_config else 8
+
+    patch_size = 2
+    try:
+        model_index = get_hf_file_to_dict("model_index.json", model_name) or {}
+        patch_size = int(model_index.get("patch_size", patch_size))
+    except Exception:  # noqa: BLE001
+        pass
+
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * patch_size)
+
+    def post_process_func(images: torch.Tensor):
+        return image_processor.postprocess(images)
+
+    return post_process_func
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, int]:
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+class Krea2Pipeline(nn.Module, DiffusionPipelineProfilerMixin, ProgressBarMixin, SupportsComponentDiscovery):
+    """The Krea 2 text-to-image pipeline (vLLM-Omni port of ``diffusers.Krea2Pipeline``).
+
+    Components:
+    - ``scheduler``: ``FlowMatchEulerDiscreteScheduler`` (resolution-aware exponential time shift).
+    - ``vae``: the Qwen-Image VAE (f8, 16 latent channels).
+    - ``text_encoder``: a Qwen3-VL model; the pipeline consumes a stack of hidden states tapped from several
+      decoder layers rather than the last hidden state.
+    - ``tokenizer``: the tokenizer paired with the text encoder.
+    - ``transformer``: the Krea 2 single-stream MMDiT that predicts the flow-matching velocity.
+    """
+
+    _dit_modules: ClassVar[list[str]] = ["transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
+    _vae_modules: ClassVar[list[str]] = ["vae"]
+
+    _PROFILER_TARGETS: ClassVar[list[str]] = ["text_encoder.forward", "diffuse", "vae.decode"]
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config
+        # Only the transformer goes through the standard weights loader; text encoder and VAE load eagerly below.
+        self.weights_sources = [
+            DiffusersPipelineLoader.ComponentSource(
+                model_or_path=od_config.model,
+                subfolder="transformer",
+                revision=od_config.revision,
+                prefix="transformer.",
+                fall_back_to_pt=True,
+            )
+        ]
+
+        self.device = get_local_device()
+        model = od_config.model
+        local_files_only = os.path.isdir(model)
+
+        # Pipeline-level config lives in model_index.json, not the transformer config.
+        model_index = get_hf_file_to_dict("model_index.json", model) or {}
+        select_layers = model_index.get("text_encoder_select_layers") or DEFAULT_TEXT_ENCODER_SELECT_LAYERS
+        self.text_encoder_select_layers = tuple(int(i) for i in select_layers)
+        self.is_distilled = bool(model_index.get("is_distilled", False))
+        self.patch_size = int(model_index.get("patch_size", 2))
+
+        subfolders = ["scheduler", "text_encoder", "vae", "tokenizer"]
+        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
+        prefetch_subfolders(model, subfolders, local_files_only=local_files_only)
+
+        self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            model, subfolder="scheduler", local_files_only=local_files_only
+        )
+
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen3VLModel.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=od_config.dtype,
+        )
+        # Drop the unused Qwen3-VL vision tower before moving to GPU so it never consumes GPU memory.
+        if hasattr(self.text_encoder, "visual"):
+            del self.text_encoder.visual
+        else:
+            logger.warning("Krea2: vision tower not found on text encoder; skipping drop")
+        self.text_encoder = self.text_encoder.to(self.device)
+
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKLQwenImage.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=od_config.dtype,
+        ).to(self.device)
+
+        transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Krea2Transformer2DModel)
+        self.transformer = Krea2Transformer2DModel(
+            od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs
+        )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
+
+        self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * self.patch_size)
+        self.default_sample_size = 1024
+
+        # Chat template tokenized as a fixed-length block: prompt padded to a fixed length first, assistant suffix
+        # appended after the padding; the first ``prompt_template_encode_start_idx`` system tokens are later dropped.
+        self.prompt_template_encode_prefix = (
+            "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, "
+            "spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"
+        )
+        self.prompt_template_encode_suffix = "<|im_end|>\n<|im_start|>assistant\n"
+        self.prompt_template_encode_start_idx = 34
+        self.prompt_template_encode_num_suffix_tokens = 5
+
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+        )
+
+    def get_text_hidden_states(
+        self,
+        prompt: str | list[str],
+        max_sequence_length: int = 512,
+        device: torch.device | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Tokenize ``prompt`` into the fixed-length Krea 2 layout and tap the selected encoder hidden states.
+
+        Returns ``(hidden_states, attention_mask)`` of shapes ``(batch, text_seq_len, num_text_layers,
+        text_hidden_dim)`` and ``(batch, text_seq_len)`` (bool).
+        """
+        device = device or self.device
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prefix_idx = self.prompt_template_encode_start_idx
+        text = [self.prompt_template_encode_prefix + e for e in prompt]
+        text_tokens = self.tokenizer(
+            text,
+            truncation=True,
+            padding="max_length",
+            max_length=max_sequence_length + prefix_idx - self.prompt_template_encode_num_suffix_tokens,
+            return_tensors="pt",
+        ).to(device)
+        suffix_tokens = self.tokenizer([self.prompt_template_encode_suffix] * len(text), return_tensors="pt").to(device)
+
+        input_ids = torch.cat([text_tokens.input_ids, suffix_tokens.input_ids], dim=1)
+        attention_mask = torch.cat([text_tokens.attention_mask, suffix_tokens.attention_mask], dim=1).bool()
+
+        # Padding sits in the middle (``[prefix | prompt | PAD | suffix]``), so positions must count only real tokens
+        # (padding consumes no position). Broadcast across the 3 mRoPE axes (T/H/W equal for text).
+        position_ids = (attention_mask.long().cumsum(dim=-1) - 1).clamp(min=0)
+        position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
+
+        outputs = self.text_encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_hidden_states=True,
+        )
+        hidden_states = torch.stack([outputs.hidden_states[i] for i in self.text_encoder_select_layers], dim=2)
+
+        hidden_states = hidden_states[:, prefix_idx:]
+        attention_mask = attention_mask[:, prefix_idx:]
+        return hidden_states, attention_mask
+
+    def encode_prompt(
+        self,
+        prompt: str | list[str],
+        device: torch.device | None = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_embeds_mask: torch.Tensor | None = None,
+        max_sequence_length: int = 512,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        device = device or self.device
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self.get_text_hidden_states(prompt, max_sequence_length, device)
+
+        batch_size, seq_len, num_text_layers, dim = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, num_text_layers, dim)
+        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt)
+        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+        return prompt_embeds, prompt_embeds_mask
+
+    def _pack_latents(self, latents, batch_size, num_channels_latents, height, width):
+        p = self.patch_size
+        latents = latents.view(batch_size, num_channels_latents, height // p, p, width // p, p)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // p) * (width // p), num_channels_latents * p * p)
+        return latents
+
+    def _unpack_latents(self, latents, height, width):
+        batch_size, _, channels = latents.shape
+        p = self.patch_size
+
+        height = p * (int(height) // (self.vae_scale_factor * p))
+        width = p * (int(width) // (self.vae_scale_factor * p))
+
+        latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+        latents = latents.reshape(batch_size, channels // (p * p), 1, height, width)
+        return latents
+
+    @staticmethod
+    def prepare_position_ids(text_seq_len: int, grid_height: int, grid_width: int, device: torch.device):
+        """Build the ``(text_seq_len + grid_height * grid_width, 3)`` rotary coordinates: text tokens sit at the
+        origin, image tokens carry their ``(0, h, w)`` latent-grid coordinates."""
+        text_ids = torch.zeros(text_seq_len, 3, device=device)
+        image_ids = torch.zeros(grid_height, grid_width, 3, device=device)
+        image_ids[..., 1] = torch.arange(grid_height, device=device)[:, None]
+        image_ids[..., 2] = torch.arange(grid_width, device=device)[None, :]
+        image_ids = image_ids.reshape(grid_height * grid_width, 3)
+        return torch.cat([text_ids, image_ids], dim=0)
+
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        latent_height = height // self.vae_scale_factor
+        latent_width = width // self.vae_scale_factor
+        shape = (batch_size, num_channels_latents, latent_height, latent_width)
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = self._pack_latents(latents, batch_size, num_channels_latents, latent_height, latent_width)
+        return latents
+
+    def _extract_prompts(self, prompts):
+        """Extract prompt and negative_prompt from the request's OmniPromptType list."""
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
+        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
+            negative_prompt = None
+        elif prompts:
+            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
+        else:
+            negative_prompt = None
+        return prompt, negative_prompt
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 0
+
+    def diffuse(
+        self,
+        latents: torch.Tensor,
+        timesteps: torch.Tensor,
+        position_ids: torch.Tensor,
+        prompt_embeds: torch.Tensor,
+        prompt_embeds_mask: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor | None,
+        negative_prompt_embeds_mask: torch.Tensor | None,
+        guidance_scale: float,
+    ) -> torch.Tensor:
+        """Run the flow-matching denoising loop and return the final latents."""
+        self.scheduler.set_begin_index(0)
+        transformer_dtype = self.transformer.dtype
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for _, t in enumerate(timesteps):
+                timestep = (t / self.scheduler.config.num_train_timesteps).expand(latents.shape[0])
+                timestep = timestep.to(transformer_dtype)
+                latent_input = latents.to(transformer_dtype)
+
+                noise_pred = self.transformer(
+                    hidden_states=latent_input,
+                    encoder_hidden_states=prompt_embeds.to(transformer_dtype),
+                    timestep=timestep,
+                    position_ids=position_ids,
+                    encoder_attention_mask=prompt_embeds_mask,
+                )
+
+                if self.do_classifier_free_guidance:
+                    neg_noise_pred = self.transformer(
+                        hidden_states=latent_input,
+                        encoder_hidden_states=negative_prompt_embeds.to(transformer_dtype),
+                        timestep=timestep,
+                        position_ids=position_ids,
+                        encoder_attention_mask=negative_prompt_embeds_mask,
+                    )
+                    noise_pred = noise_pred + guidance_scale * (noise_pred - neg_noise_pred)
+
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                progress_bar.update()
+
+        return latents
+
+    def forward(
+        self,
+        req: OmniDiffusionRequest,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: int = 28,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 4.5,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        output_type: str | None = "pil",
+        max_sequence_length: int = 512,
+    ) -> DiffusionOutput:
+        extracted_prompt, extracted_negative = self._extract_prompts(req.prompts)
+        prompt = extracted_prompt or prompt
+        if extracted_negative is not None:
+            negative_prompt = extracted_negative
+
+        sp = req.sampling_params
+        height = sp.height or height
+        width = sp.width or width
+        num_inference_steps = sp.num_inference_steps or num_inference_steps
+        sigmas = sp.sigmas or sigmas
+        max_sequence_length = sp.max_sequence_length or max_sequence_length
+        generator = sp.generator or generator
+        if sp.guidance_scale_provided:
+            guidance_scale = sp.guidance_scale
+        else:
+            # Distilled checkpoint is trained guidance-free (and the request layer coerces an explicit
+            # guidance_scale=0 to "not provided"), so default it to no CFG.
+            guidance_scale = 0.0 if self.is_distilled else guidance_scale
+        num_images_per_prompt = sp.num_outputs_per_prompt if sp.num_outputs_per_prompt > 0 else num_images_per_prompt
+
+        multiple = self.vae_scale_factor * self.patch_size
+        if height % multiple != 0 or width % multiple != 0:
+            rounded_height = ((height + multiple - 1) // multiple) * multiple
+            rounded_width = ((width + multiple - 1) // multiple) * multiple
+            logger.warning(
+                "`height` and `width` must be multiples of %d; rounding up from %dx%d to %dx%d.",
+                multiple,
+                height,
+                width,
+                rounded_height,
+                rounded_width,
+            )
+            height, width = rounded_height, rounded_width
+
+        device = self.device
+        self._guidance_scale = guidance_scale
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = 1
+
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        negative_prompt_embeds = None
+        negative_prompt_embeds_mask = None
+        if self.do_classifier_free_guidance:
+            if negative_prompt is None:
+                negative_prompt = ""
+            if isinstance(negative_prompt, str):
+                negative_prompt = [negative_prompt] * batch_size
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                prompt=negative_prompt,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+
+        num_channels_latents = self.transformer.in_channels // (self.patch_size**2)
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        grid_height = height // (self.vae_scale_factor * self.patch_size)
+        grid_width = width // (self.vae_scale_factor * self.patch_size)
+        position_ids = self.prepare_position_ids(prompt_embeds.shape[1], grid_height, grid_width, device)
+
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        if self.is_distilled:
+            mu = 1.15
+        else:
+            mu = calculate_shift(
+                image_seq_len,
+                self.scheduler.config.get("base_image_seq_len", 256),
+                self.scheduler.config.get("max_image_seq_len", 6400),
+                self.scheduler.config.get("base_shift", 0.5),
+                self.scheduler.config.get("max_shift", 1.15),
+            )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu
+        )
+
+        latents = self.diffuse(
+            latents=latents,
+            timesteps=timesteps,
+            position_ids=position_ids,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds=negative_prompt_embeds,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            guidance_scale=guidance_scale,
+        )
+
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width)
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+
+        return DiffusionOutput(
+            output=image,
+            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        loaded_weights = loader.load_weights(weights)
+        # VAE and text encoder load eagerly in __init__; record their params so strict load checks pass.
+        loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
+        if self.text_encoder is not None:
+            loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
+        return loaded_weights

vllm_omni/diffusion/models/lance/pipeline_lance.py

@@ -836,8 +836,10 @@ class LancePipeline(BagelPipeline):
         frames = [Image.fromarray(f) for f in frames_np]
         logger.info("Lance t2v: decoded %d frames at %dx%d", len(frames), frames[0].width, frames[0].height)
         return DiffusionOutput(
-            output=frames[0],
-            custom_output={"video_frames": frames, "video_shape": video_shape},
+            output={
+                "payload": {"video": frames},
+                "metadata": {"video": {"shape": video_shape}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1069,13 +1071,11 @@ class LancePipeline(BagelPipeline):
             )
 
         img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape)
-        # Put the PIL image into ``custom_output`` too — the orchestrator
-        # serializes ``custom_output`` across the IPC boundary but strips
-        # the bare ``output`` field, so consumers that only see
-        # ``outputs[0].output`` (e.g. gradio_demo) get ``None`` otherwise.
         return DiffusionOutput(
-            output=img,
-            custom_output={"image": img, "image_shape": image_shape},
+            output={
+                "payload": {"image": img},
+                "metadata": {"image": {"shape": image_shape}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1361,8 +1361,10 @@ class LancePipeline(BagelPipeline):
         frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], out_shape)
         frames = [Image.fromarray(f) for f in frames_np]
         return DiffusionOutput(
-            output=frames[0],
-            custom_output={"video_frames": frames, "video_shape": out_shape},
+            output={
+                "payload": {"video": frames},
+                "metadata": {"video": {"shape": out_shape}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1622,8 +1624,10 @@ class LancePipeline(BagelPipeline):
         frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape)
         frames = [Image.fromarray(f) for f in frames_np]
         return DiffusionOutput(
-            output=frames[0],
-            custom_output={"video_frames": frames, "video_shape": video_shape},
+            output={
+                "payload": {"video": frames},
+                "metadata": {"video": {"shape": video_shape}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1705,8 +1709,10 @@ class LancePipeline(BagelPipeline):
         text_output = text_output.lstrip("\n")
         logger.info("Lance x2t_video: generated %d tokens", token_ids.shape[0])
         return DiffusionOutput(
-            output=text_output,
-            custom_output={"text_output": text_output},
+            output={
+                "payload": {"text": text_output},
+                "metadata": {"text": {"text_output": text_output}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 
@@ -1801,8 +1807,10 @@ class LancePipeline(BagelPipeline):
         text_output = text_output.lstrip("\n")
         logger.info("Lance x2t_image: generated %d tokens", token_ids.shape[0])
         return DiffusionOutput(
-            output=text_output,
-            custom_output={"text_output": text_output},
+            output={
+                "payload": {"text": text_output},
+                "metadata": {"text": {"text_output": text_output}},
+            },
             stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
         )
 

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py

@@ -495,66 +495,35 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
         noise_pred = comb_pred * scale
         return noise_pred
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 4.5,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        enable_cfg_renorm: bool | None = True,
-        cfg_renorm_min: float | None = 0.0,
-        enable_prompt_rewrite: bool | None = True,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
-        height = req.sampling_params.height or height or self.default_sample_size * self.vae_scale_factor
-        width = req.sampling_params.width or width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        generator = req.sampling_params.generator or generator
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
+        width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        generator = req.sampling_params.generator
+        guidance_scale = req.sampling_params.guidance_scale
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt is not None
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        enable_prompt_rewrite = req.sampling_params.extra_args.get("enable_prompt_rewrite", enable_prompt_rewrite)
-        enable_cfg_renorm = req.sampling_params.extra_args.get("enable_cfg_renorm", enable_cfg_renorm)
-        cfg_renorm_min = req.sampling_params.extra_args.get("cfg_renorm_min", cfg_renorm_min)
+        latents = req.sampling_params.latents
 
-        req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts]
-        if any(p is not None for p in req_prompt_embeds):
-            # If at list one prompt is provided as an embedding,
-            # Then assume that the user wants to provide embeddings for all prompts, and enter this if block
-            # If the user in fact provides mixed input format, req_prompt_embeds will have some None's
-            # And `torch.stack` automatically raises an exception for us
-            prompt_embeds = torch.stack(req_prompt_embeds)  # type: ignore # intentionally expect TypeError
-
-        req_negative_prompt_embeds = [
-            p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts
-        ]
-        if any(p is not None for p in req_negative_prompt_embeds):
-            negative_prompt_embeds = torch.stack(req_negative_prompt_embeds)  # type: ignore # intentionally expect TypeError
+        extra_args = getattr(req.sampling_params, "extra_args", {}) or {}
+        enable_prompt_rewrite: bool = extra_args.get("enable_prompt_rewrite", True)
+        enable_cfg_renorm: bool = extra_args.get("enable_cfg_renorm", True)
+        cfg_renorm_min: float = extra_args.get("cfg_renorm_min", 0.0)
+        joint_attention_kwargs: dict[str, Any] | None = None
+        output_type = req.sampling_params.output_type or "pil"
 
         self.check_inputs(
             prompt,
@@ -571,12 +540,7 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = len(prompt)
 
         device = self.device
         if enable_prompt_rewrite and prompt is not None:

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py

@@ -6,7 +6,7 @@ import math
 import os
 import re
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -542,24 +542,7 @@ class LongCatImageEditPipeline(
                 f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
             )
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 3.5,
-        num_images_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
         if len(req.prompts) > 1:
@@ -570,27 +553,22 @@ class LongCatImageEditPipeline(
         first_prompt = req.prompts[0]
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
         negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
-        prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("prompt_embeds")
-        negative_prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt_embeds")  # type: ignore # Why it is list[torch.Tensor] in OmniTokenInputs or OmniEmbedsPrompt? Doesn't make sense
 
-        sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        sigmas = req.sampling_params.sigmas
+        guidance_scale = req.sampling_params.guidance_scale
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt is not None
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        generator = req.sampling_params.generator or generator
+        generator = req.sampling_params.generator
         height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        output_type = req.sampling_params.output_type or "pil"
 
-        if prompt is not None:
-            batch_size = 1 if isinstance(prompt, str) else len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
 
         if not isinstance(first_prompt, str) and "preprocessed_image" in (
             additional_information := first_prompt.get("additional_information", {})
@@ -600,13 +578,7 @@ class LongCatImageEditPipeline(
             calculated_height = additional_information.get("calculated_height", height)
             calculated_width = additional_information.get("calculated_width", width)
         else:
-            image_size = image[0].size if isinstance(image, list) else image.size
-            calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1])
-
-            if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
-                image = self.image_processor.resize(image, calculated_height, calculated_width)
-                prompt_image = self.image_processor.resize(image, calculated_height // 2, calculated_width // 2)
-                image = self.image_processor.preprocess(image, calculated_height, calculated_width)
+            raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
         self.check_inputs(
             prompt,
@@ -621,7 +593,8 @@ class LongCatImageEditPipeline(
             prompt=prompt, image=prompt_image, prompt_embeds=prompt_embeds, num_images_per_prompt=num_images_per_prompt
         )
 
-        if guidance_scale > 1:
+        do_true_cfg = guidance_scale > 1 and negative_prompt is not None
+        if do_true_cfg:
             (negative_prompt_embeds, negative_text_ids) = self.encode_prompt(
                 prompt=negative_prompt,
                 image=prompt_image,
@@ -680,7 +653,6 @@ class LongCatImageEditPipeline(
                 latent_model_input = torch.cat([latents, image_latents], dim=1)
 
             timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype)
-            do_true_cfg = guidance_scale > 1
             positive_kwargs = {
                 "hidden_states": latent_model_input,
                 "timestep": timestep / 1000,

vllm_omni/diffusion/models/ltx2/__init__.py

@@ -10,10 +10,8 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     get_ltx2_post_process_func,
     load_transformer_config,
 )
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import (
-    LTX23ImageToVideoPipeline,
-    LTX23Pipeline,
-)
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import (
     LTX2I2VDMD2Pipeline,
     LTX2ImageToVideoPipeline,

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py

@@ -1,16 +1,452 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Re-exports for LTX-2.3 I2V pipeline variants.
+"""LTX-2.3 image-to-video pipeline."""
 
-The registry loads pipeline classes by (mod_folder, mod_relname, cls_name).
-This module exposes the I2V class names so the registry can find them.
-"""
+from __future__ import annotations
+
+from typing import Any
+
+import PIL.Image
+import torch
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
 from .pipeline_ltx2_3 import (
-    LTX23ImageToVideoPipeline,
-    get_ltx2_post_process_func,  # noqa: F401 - loaded by registry via getattr
+    LTX23Pipeline,
+    _LTX23DenoiseContext,
+    _LTX23ForwardContext,
+    _LTX23PromptContext,
+    _LTX23RequestInputs,
+    get_ltx2_post_process_func,
 )
+from .pipeline_ltx2_image2video import LTX2ImageToVideoPipeline, _I2VVideoAudioScheduler
+
+
+class LTX23ImageToVideoPipeline(LTX23Pipeline):
+    """LTX-2.3 image-to-video pipeline.
+
+    This keeps the LTX-2.3 prompt connector, x0-space CFG, sigma prompt
+    modulation, and audio branch semantics from ``LTX23Pipeline`` while
+    reusing the existing LTX image-conditioning contract: the first video
+    latent frame is encoded from the input image and remains fixed during
+    denoising.
+    """
+
+    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
+        super().__init__(od_config=od_config, prefix=prefix)
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear")
+
+    support_image_input = True
+
+    _normalize_latents = staticmethod(LTX2ImageToVideoPipeline._normalize_latents)
+    _create_noised_state = staticmethod(LTX2ImageToVideoPipeline._create_noised_state)
+
+    @staticmethod
+    def _resolve_single_prompt_image(raw_image: Any) -> Any:
+        if isinstance(raw_image, list):
+            if len(raw_image) != 1:
+                raise ValueError(
+                    "LTX-2.3 I2V prompt dictionaries support exactly one image per prompt. "
+                    "Pass one image per prompt for batched I2V requests."
+                )
+            return raw_image[0]
+        return raw_image
+
+    @staticmethod
+    def _resolve_additional_image(additional: dict[str, Any]) -> Any:
+        raw_image = additional.get("preprocessed_image")
+        if raw_image is None:
+            raw_image = additional.get("pixel_values")
+        if raw_image is None:
+            raw_image = additional.get("image")
+        return raw_image
+
+    def prepare_latents(
+        self,
+        image: torch.Tensor | None = None,
+        batch_size: int = 1,
+        num_channels_latents: int = 128,
+        height: int = 512,
+        width: int = 768,
+        num_frames: int = 121,
+        noise_scale: float = 0.0,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Prepare I2V latents and the first-frame conditioning mask.
+
+        If caller-provided latents are used without an image, the latents must
+        already represent the full video state including the conditioning first
+        frame. Packed 3D latents are assumed to be in transformer token layout.
+        """
+        height = height // self.vae_spatial_compression_ratio
+        width = width // self.vae_spatial_compression_ratio
+        num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
+
+        shape = (batch_size, num_channels_latents, num_frames, height, width)
+        mask_shape = (batch_size, 1, num_frames, height, width)
+
+        if latents is not None:
+            if latents.ndim == 5:
+                batch_size, _, num_frames, height, width = latents.shape
+                mask_shape = (batch_size, 1, num_frames, height, width)
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
+
+                latents = self._normalize_latents(
+                    latents,
+                    self.vae.latents_mean,
+                    self.vae.latents_std,
+                    self.vae.config.scaling_factor,
+                )
+                latents = self._create_noised_state(latents, noise_scale * (1 - conditioning_mask), generator)
+                latents = self._pack_latents(
+                    latents,
+                    self.transformer_spatial_patch_size,
+                    self.transformer_temporal_patch_size,
+                )
+            else:
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
+
+            conditioning_mask = self._pack_latents(
+                conditioning_mask,
+                self.transformer_spatial_patch_size,
+                self.transformer_temporal_patch_size,
+            ).squeeze(-1)
+            if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
+                raise ValueError(
+                    "Provided `latents` tensor has shape"
+                    f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
+                )
+            return latents.to(device=device, dtype=dtype), conditioning_mask
+
+        if image is None:
+            raise ValueError("`image` must be provided when `latents` is None.")
+
+        image_batch_size = image.shape[0]
+        if image_batch_size == 0:
+            raise ValueError("`image` batch is empty.")
+        if batch_size % image_batch_size != 0:
+            raise ValueError(
+                f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) "
+                "for image-to-video outputs."
+            )
+        num_videos_per_prompt = batch_size // image_batch_size
+
+        if isinstance(generator, list):
+            if len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective"
+                    f" batch size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+            image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)]
+            init_latents = [
+                retrieve_latents(self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), image_generators[i], "argmax")
+                for i in range(image_batch_size)
+            ]
+        else:
+            init_latents = [
+                retrieve_latents(self.vae.encode(img.unsqueeze(0).unsqueeze(2)), generator, "argmax") for img in image
+            ]
+
+        init_latents = torch.cat(init_latents, dim=0).to(dtype)
+        if num_videos_per_prompt > 1:
+            init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0)
+        init_latents = self._normalize_latents(
+            init_latents,
+            self.vae.latents_mean,
+            self.vae.latents_std,
+        )
+        init_latents = init_latents.repeat(1, 1, num_frames, 1, 1)
+
+        conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype)
+        conditioning_mask[:, :, 0] = 1.0
+
+        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)
+
+        conditioning_mask = self._pack_latents(
+            conditioning_mask,
+            self.transformer_spatial_patch_size,
+            self.transformer_temporal_patch_size,
+        ).squeeze(-1)
+        latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size)
+
+        return latents, conditioning_mask
+
+    def check_inputs(
+        self,
+        image,
+        height,
+        width,
+        prompt,
+        latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+    ):
+        if image is None and latents is None:
+            raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.")
+        super().check_inputs(
+            prompt=prompt,
+            height=height,
+            width=width,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+        )
+
+    _step_video_latents_i2v = LTX2ImageToVideoPipeline._step_video_latents_i2v
+
+    def _resolve_request_image(
+        self,
+        req: DiffusionRequestBatch,
+        image: PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None,
+        request_inputs: _LTX23RequestInputs,
+    ) -> PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None:
+        if image is not None or not req.prompts:
+            return image
+
+        raw_images = []
+        for prompt_item in req.prompts:
+            if isinstance(prompt_item, str):
+                raw_image = None
+            else:
+                multi_modal_data = prompt_item.get("multi_modal_data") or {}
+                raw_image = multi_modal_data.get("image")
+                if raw_image is None:
+                    additional = prompt_item.get("additional_information") or {}
+                    raw_image = self._resolve_additional_image(additional)
+            raw_image = self._resolve_single_prompt_image(raw_image)
+            if isinstance(raw_image, str):
+                raw_image = PIL.Image.open(raw_image).convert("RGB")
+            raw_images.append(raw_image)
+
+        if any(img is None for img in raw_images) and request_inputs.latents is None:
+            raise ValueError("Image is required for LTX-2.3 I2V generation.")
+        if len(raw_images) == 1:
+            return raw_images[0]
+        if raw_images:
+            return raw_images
+        return image
+
+    def _check_forward_inputs(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        image: Any | None = None,
+    ) -> None:
+        self.check_inputs(
+            image=image,
+            height=request_inputs.height,
+            width=request_inputs.width,
+            prompt=request_inputs.prompt,
+            latents=request_inputs.latents,
+            prompt_embeds=request_inputs.prompt_embeds,
+            negative_prompt_embeds=request_inputs.negative_prompt_embeds,
+            prompt_attention_mask=request_inputs.prompt_attention_mask,
+            negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask,
+        )
+
+    def _resolve_video_latent_dimensions(
+        self,
+        request_inputs: _LTX23RequestInputs,
+    ) -> tuple[int, int, int]:
+        latent_num_frames = (request_inputs.num_frames - 1) // self.vae_temporal_compression_ratio + 1
+        latent_height = request_inputs.height // self.vae_spatial_compression_ratio
+        latent_width = request_inputs.width // self.vae_spatial_compression_ratio
+        if request_inputs.latents is not None:
+            if request_inputs.latents.ndim == 5:
+                _, _, latent_num_frames, latent_height, latent_width = request_inputs.latents.shape
+            elif request_inputs.latents.ndim != 3:
+                raise ValueError(
+                    f"Provided `latents` tensor has shape {request_inputs.latents.shape}, but the expected shape is "
+                    "either [batch_size, seq_len, num_features] or "
+                    "[batch_size, latent_dim, latent_frames, latent_height, latent_width]."
+                )
+        return latent_num_frames, latent_height, latent_width
+
+    def _prepare_video_latents_stage(
+        self,
+        request_inputs: _LTX23RequestInputs,
+        prompt_context: _LTX23PromptContext,
+        *,
+        device: torch.device,
+        noise_scale: float,
+        image: Any | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if request_inputs.latents is None:
+            if isinstance(image, torch.Tensor):
+                if image.ndim == 3:
+                    image = image.unsqueeze(0)
+            elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor):
+                image = torch.stack(image, dim=0)
+            else:
+                image = self.video_processor.preprocess(
+                    image,
+                    height=request_inputs.height,
+                    width=request_inputs.width,
+                )
+            image = image.to(device=device, dtype=prompt_context.positive_connector_prompt_embeds.dtype)
+
+        num_channels_latents = self.transformer.config.in_channels
+        return self.prepare_latents(
+            image,
+            prompt_context.batch_size * request_inputs.num_videos_per_prompt,
+            num_channels_latents,
+            request_inputs.height,
+            request_inputs.width,
+            request_inputs.num_frames,
+            noise_scale,
+            torch.float32,
+            device,
+            request_inputs.generator,
+            request_inputs.latents,
+        )
+
+    def _make_video_audio_scheduler(
+        self,
+        audio_scheduler: Any,
+        latent_num_frames: int,
+        latent_height: int,
+        latent_width: int,
+    ) -> Any:
+        return _I2VVideoAudioScheduler(
+            self,
+            audio_scheduler,
+            latent_num_frames,
+            latent_height,
+            latent_width,
+        )
+
+    def _prepare_denoise_context_for_cfg(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+    ) -> _LTX23DenoiseContext:
+        denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
+        if denoise_ctx.conditioning_mask is None:
+            raise ValueError("I2V denoising requires a conditioning mask.")
+        if self.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready:
+            denoise_ctx.conditioning_mask_for_model = torch.cat(
+                [denoise_ctx.conditioning_mask, denoise_ctx.conditioning_mask]
+            )
+        else:
+            denoise_ctx.conditioning_mask_for_model = denoise_ctx.conditioning_mask
+        return denoise_ctx
+
+    def _denoise_timestep_kwargs(
+        self,
+        ts: torch.Tensor,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+    ) -> dict[str, torch.Tensor]:
+        conditioning_mask = (
+            denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model
+        )
+        if conditioning_mask is None:
+            raise ValueError("I2V denoising requires a conditioning mask.")
+        return {
+            "timestep": ts.unsqueeze(-1) * (1 - conditioning_mask),
+            "audio_timestep": ts,
+            "sigma": ts,
+        }
+
+    def _step_denoised_latents(
+        self,
+        forward_ctx: _LTX23ForwardContext,
+        denoise_ctx: _LTX23DenoiseContext,
+        noise_pred_video: torch.Tensor,
+        noise_pred_audio: torch.Tensor,
+        t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        latents, audio_latents = self.scheduler_step_maybe_with_cfg(
+            (noise_pred_video, noise_pred_audio),
+            (t, t),
+            (denoise_ctx.latents, denoise_ctx.audio_latents),
+            do_true_cfg=self.do_classifier_free_guidance,
+            per_request_scheduler=forward_ctx.video_audio_scheduler,
+        )
+        return self._synchronize_cfg_parallel_step_output(
+            (latents, audio_latents),
+            do_true_cfg=self.do_classifier_free_guidance,
+        )
+
+    @torch.no_grad()
+    def forward(
+        self,
+        req: DiffusionRequestBatch,
+        image: PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None = None,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_frames: int | None = None,
+        frame_rate: float | None = None,
+        num_inference_steps: int | None = None,
+        sigmas: list[float] | None = None,
+        timesteps: list[int] | None = None,
+        guidance_scale: float = 4.0,
+        noise_scale: float = 0.0,
+        num_videos_per_prompt: int | None = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        audio_latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        negative_prompt_attention_mask: torch.Tensor | None = None,
+        decode_timestep: float | list[float] = 0.0,
+        decode_noise_scale: float | list[float] | None = None,
+        output_type: str = "np",
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        max_sequence_length: int | None = None,
+    ) -> list[DiffusionOutput]:
+        request_inputs = self._resolve_request_inputs(
+            req,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            num_inference_steps=num_inference_steps,
+            timesteps=timesteps,
+            guidance_scale=guidance_scale,
+            num_videos_per_prompt=num_videos_per_prompt,
+            generator=generator,
+            latents=latents,
+            audio_latents=audio_latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            decode_timestep=decode_timestep,
+            decode_noise_scale=decode_noise_scale,
+            output_type=output_type,
+            max_sequence_length=max_sequence_length,
+        )
+        image = self._resolve_request_image(req, image, request_inputs)
+        return self._forward_impl(
+            req,
+            request_inputs,
+            noise_scale=noise_scale,
+            sigmas=sigmas,
+            timesteps=timesteps,
+            attention_kwargs=attention_kwargs,
+            image=image,
+        )
+
 
 __all__ = [
     "LTX23ImageToVideoPipeline",

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -9,6 +9,7 @@ from diffusers.utils.torch_utils import randn_tensor
 from torch import nn
 from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
 from vllm.config import VllmConfig
+from vllm.logger import init_logger
 from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
@@ -19,6 +20,8 @@ from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DMode
 from .rope_real import RotaryPosEmbedReal
 from .schedulers import FlowMatchEulerDiscreteScheduler
 
+logger = init_logger(__name__)
+
 
 class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
     """
@@ -60,14 +63,14 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config)
 
         # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
-        llm_hidden_size = int(
-            getattr(
-                getattr(self.config.llm_config, "text_config", None),
-                "hidden_size",
-                0,
-            )
-            or 0
-        )
+        llm_hidden_size = 0
+        text_config = self.config.get_text_config()
+        if text_config is None:
+            logger.warning("No text config; failed to infer llm_hidden_size.")
+        elif not hasattr(text_config, "hidden_size"):
+            logger.warning("Text config exists, but has no hidden_size attribute; failed to infer llm_hidden_size.")
+        else:
+            llm_hidden_size = int(text_config.hidden_size or 0)
         if llm_hidden_size <= 0:
             raise ValueError(
                 "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size"

vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py

@@ -24,7 +24,6 @@ from __future__ import annotations
 
 import logging
 import os
-from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
 
@@ -49,7 +48,9 @@ from vllm_omni.diffusion.models.ming_flash_omni.ming_zimage_transformer import (
     MingZImageTransformer2DModel,
 )
 from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
@@ -58,38 +59,6 @@ from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenCon
 logger = logging.getLogger(__name__)
 
 
-@dataclass
-class _ZPipelineSamplingParams:
-    """Typed shim satisfying the attributes ZImagePipeline.forward reads
-    from ``req.sampling_params``.  Unlike SimpleNamespace this fails loudly
-    at construction if a required field is missing."""
-
-    height: int
-    width: int
-    num_inference_steps: int
-    guidance_scale: float
-    generator: torch.Generator | None = None
-    strength: float | None = None
-    sigmas: list[float] | None = None
-    max_sequence_length: int = 512
-    guidance_rescale: float | None = None
-    num_outputs_per_prompt: int = 1
-
-
-@dataclass
-class _ZPipelineRequest:
-    """Typed shim satisfying the attributes ZImagePipeline.forward reads
-    from ``req``."""
-
-    request_id: str
-    sampling_params: _ZPipelineSamplingParams
-    prompts: list[dict[str, Any]] = field(default_factory=lambda: [{"prompt": "", "negative_prompt": ""}])
-
-    @property
-    def num_reqs(self) -> int:
-        return 1
-
-
 class MingImagePipeline(ZImagePipeline):
     """Ming-flash-omni-2.0 text-to-image diffusion pipeline.
 
@@ -123,6 +92,10 @@ class MingImagePipeline(ZImagePipeline):
         self.device = self._execution_device  # Ming convention alias
         self._dtype = dtype
 
+        # Request-scoped conditioning handed to the inherited encode_prompt override
+        self._pending_prompt_embeds: list[torch.Tensor] | None = None
+        self._pending_negative_prompt_embeds: list[torch.Tensor] | None = None
+
         # Ming's per-checkpoint image-gen configuration. We cannot rely on
         # ``od_config.hf_config.image_gen_config`` because the diffusion
         # stage is started with ``hf_config_name: thinker_config`` (the
@@ -273,6 +246,21 @@ class MingImagePipeline(ZImagePipeline):
         latent = self.vae.encode(ref).latent_dist.mode()
         return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor
 
+    def encode_prompt(self, *args, **kwargs):  # noqa: ARG002
+        """Return Ming's precomputed conditioning instead of encoding text.
+
+        NOTE: Ming has no Z-Image text_encoder; its conditioning (cap_feats, optionally ByT5-augmented)
+        is computed in forward and stashed on `self._pending_*` immediately before
+        the `super().forward` call, so we simply hand it back here.
+        """
+        if self._pending_prompt_embeds is None:
+            raise RuntimeError(
+                "MingImagePipeline.encode_prompt called without pending "
+                "conditioning; it must run within MingImagePipeline.forward."
+            )
+
+        return self._pending_prompt_embeds, self._pending_negative_prompt_embeds
+
     # ------------------------------------------------------------------
     # Forward
     # ------------------------------------------------------------------
@@ -416,17 +404,27 @@ class MingImagePipeline(ZImagePipeline):
             negative_prompt_embeds = [negative_cap_feats[i] for i in range(negative_cap_feats.shape[0])]
         else:
             negative_prompt_embeds = [self.condition_encoder.zero_negative(e) for e in prompt_embeds]
+        self._pending_prompt_embeds = prompt_embeds
+        self._pending_negative_prompt_embeds = negative_prompt_embeds
 
-        z_sp = _ZPipelineSamplingParams(
+        # Build a real single-request DiffusionRequestBatch;
+        # the prompt text itself is a neutral placeholder here.
+        z_sp = OmniDiffusionSamplingParams(
             height=height,
             width=width,
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
             generator=generator,
+            output_type="pt",
         )
-        z_req = _ZPipelineRequest(
-            request_id=req.request_id or "ming-imagegen",
-            sampling_params=z_sp,
+        z_req = DiffusionRequestBatch(
+            requests=[
+                OmniDiffusionRequest(
+                    prompt={"prompt": ""},
+                    sampling_params=z_sp,
+                    request_id=req.request_id or "ming-imagegen",
+                )
+            ]
         )
 
         # Reference image (img2img) → VAE-encoded latent published on the
@@ -446,29 +444,14 @@ class MingImagePipeline(ZImagePipeline):
             None if ref_latent is None else tuple(ref_latent.shape),
         )
         try:
-            outputs = super().forward(
-                z_req,
-                prompt=None,
-                height=height,
-                width=width,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale,
-                prompt_embeds=prompt_embeds,
-                negative_prompt_embeds=negative_prompt_embeds,
-                generator=generator,
-                output_type="pt",
-                return_dict=True,
-            )
+            outputs: DiffusionOutput = super().forward(z_req)
         finally:
             set_forward_context_ref_latent(None)
+            # Drop request-scoped conditioning so we don't retain GPU tensors.
+            self._pending_prompt_embeds = None
+            self._pending_negative_prompt_embeds = None
 
-        output = outputs[0]
-        if hasattr(output, "output") and output.output is not None:
-            raw = output.output
-        elif hasattr(output, "images"):
-            raw = output.images
-        else:
-            raw = output
+        raw = outputs.output
         if not isinstance(raw, torch.Tensor):
             raise RuntimeError(f"ZImagePipeline returned non-tensor output: {type(raw).__name__}")
         if logger.isEnabledFor(logging.DEBUG):
@@ -478,7 +461,7 @@ class MingImagePipeline(ZImagePipeline):
                 raw.float().min().item(),
                 raw.float().max().item(),
             )
-        return DiffusionOutput(output=raw)
+        return outputs
 
 
 # ----------------------------------------------------------------------

vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py

@@ -22,6 +22,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,
+    SequenceParallelOutput,
+)
 from vllm_omni.platforms import current_omni_platform
 
 logger = logging.getLogger(__name__)
@@ -808,6 +812,36 @@ class OmniGen2TransformerBlock(nn.Module):
         return hidden_states
 
 
+class _OmniGen2SPInputBoundary(nn.Module):
+    """SP input-sharding boundary marker for ``OmniGen2Transformer2DModel._sp_plan``.
+
+    Resets ``sp_original_seq_len`` on every call: OmniGen2 runs the transformer
+    once per CFG branch and the unconditional / conditional branches have
+    different joint sequence lengths (text prefix differs). The SP framework's
+    auto-pad records the length only on the first call, so without this reset
+    a stale length from an earlier branch truncates the gather of later ones.
+    """
+
+    def forward(self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        from vllm_omni.diffusion.forward_context import (
+            get_forward_context,
+            is_forward_context_available,
+        )
+
+        if is_forward_context_available():
+            ctx = get_forward_context()
+            ctx.sp_original_seq_len = None
+            ctx.sp_padding_size = 0
+        return hidden_states, rotary_emb
+
+
+class _OmniGen2SPOutputBoundary(nn.Module):
+    """SP output-gathering boundary marker for ``OmniGen2Transformer2DModel._sp_plan``."""
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states
+
+
 class OmniGen2Transformer2DModel(nn.Module):
     """
     OmniGen2 Transformer 2D Model.
@@ -836,6 +870,17 @@ class OmniGen2Transformer2DModel(nn.Module):
         timestep_scale: Scale factor for timestep embeddings
     """
 
+    # ``attention_mask`` is intentionally omitted from the plan: OmniGen2Attention
+    # accepts but never reads it. ``auto_pad=True`` requires Ulysses SP — Ring SP
+    # cannot use attention masks.
+    _sp_plan = {
+        "sp_input_boundary": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+            1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True),
+        },
+        "sp_output_boundary": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+
     def __init__(
         self,
         patch_size: int = 2,
@@ -983,6 +1028,9 @@ class OmniGen2Transformer2DModel(nn.Module):
         # Add learnable embeddings to distinguish different images
         self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size))  # support max 5 ref images
 
+        self.sp_input_boundary = _OmniGen2SPInputBoundary()
+        self.sp_output_boundary = _OmniGen2SPOutputBoundary()
+
     def img_patch_embed_and_refine(
         self,
         hidden_states,
@@ -1251,9 +1299,13 @@ class OmniGen2Transformer2DModel(nn.Module):
 
         hidden_states = joint_hidden_states
 
+        hidden_states, rotary_emb = self.sp_input_boundary(hidden_states, rotary_emb)
+
         for layer in self.layers:
             hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
 
+        hidden_states = self.sp_output_boundary(hidden_states)
+
         # 4. Output norm & projection
         hidden_states = self.norm_out(hidden_states, temb)
 

vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py

@@ -18,7 +18,7 @@
 import inspect
 import json
 import os
-from collections.abc import Callable, Iterable
+from collections.abc import Iterable
 from typing import Any, ClassVar
 
 import numpy as np
@@ -551,117 +551,33 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        guidance_scale: float = 5.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        num_images_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: torch.FloatTensor | None = None,
-        negative_prompt_embeds: torch.FloatTensor | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 256,
-    ) -> DiffusionOutput:
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            negative_prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                not greater than `1`).
-            guidance_scale (`float`, *optional*, defaults to 1.0):
-                True classifier-free guidance (guidance scale) is enabled when `guidance_scale` > 1 and
-                `negative_prompt` is provided.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`list[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.FloatTensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
-            joint_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`list`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.ovis_image.OvisImagePipelineOutput`] or `tuple`:
-            [`~pipelines.ovis_image.OvisImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
-            returning a tuple, the first element is a list with the generated images.
-        """
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
         height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
-        )
-        generator = req.sampling_params.generator or generator
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        guidance_scale = req.sampling_params.guidance_scale
+        generator = req.sampling_params.generator
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        max_sequence_length = req.sampling_params.max_sequence_length or 256
+
+        joint_attention_kwargs: dict[str, Any] | None = None
+        output_type = req.sampling_params.output_type or "pil"
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # Steps:
         # 1. Check Inputs
@@ -691,12 +607,7 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi
         self._interrupt = False
 
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = len(prompt)
 
         do_classifier_free_guidance = guidance_scale > 1.0
         prompt_embeds, text_ids = self.encode_prompt(

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -73,8 +73,19 @@ def get_qwen_image_post_process_func(
     image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
 
     def post_process_func(
-        images: torch.Tensor,
+        images: torch.Tensor | dict[str, Any],
     ):
+        if isinstance(images, dict) and isinstance(images.get("payload"), dict):
+            payload = dict(images["payload"])
+            image_payload = payload.get("image")
+            if image_payload is None:
+                raise ValueError("Qwen-Image postprocess expected payload['image'] in output envelope.")
+            payload["image"] = image_processor.postprocess(image_payload)
+            metadata = images.get("metadata") or {}
+            return {
+                "payload": payload,
+                "metadata": metadata if isinstance(metadata, dict) else {},
+            }
         return image_processor.postprocess(images)
 
     return post_process_func
@@ -978,58 +989,34 @@ class QwenImagePipeline(
 
         height = state.sampling.height or self.default_sample_size * self.vae_scale_factor
         width = state.sampling.width or self.default_sample_size * self.vae_scale_factor
-        output_type = kwargs.get("output_type", "pil")
+        output_type = kwargs.get("output_type") or state.sampling.output_type or "pil"
 
         return self._decode_latents(state.latents, height, width, output_type)
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        true_cfg_scale: float = 4.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 1.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 1024,
-    ) -> list[DiffusionOutput]:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         sampling_params_list = req.sampling_params_list
         common_sampling_params = sampling_params_list[0]
         extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
-        prompt = extracted_prompt or prompt
+        prompt = extracted_prompt
 
         height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor
         height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps
-        sigmas = common_sampling_params.sigmas or sigmas
-        max_sequence_length = common_sampling_params.max_sequence_length or max_sequence_length
+        num_inference_steps = common_sampling_params.num_inference_steps or 50
+        sigmas = common_sampling_params.sigmas
+        max_sequence_length = common_sampling_params.max_sequence_length or 1024
         num_images_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        generator = req.collate_request_generators(num_images_per_prompt, generator)
-        latents = req.collate_request_tensors("latents", latents)
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        latents = req.collate_request_tensors("latents", None)
         prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
             req.prompts,
             {
-                "prompt_embeds": prompt_embeds,
-                "prompt_embeds_mask": prompt_embeds_mask,
-                "negative_prompt_embeds": negative_prompt_embeds,
-                "negative_prompt_embeds_mask": negative_prompt_embeds_mask,
+                "prompt_embeds": None,
+                "prompt_embeds_mask": None,
+                "negative_prompt_embeds": None,
+                "negative_prompt_embeds_mask": None,
             },
         )
         prompt_embeds = prompt_fields["prompt_embeds"]
@@ -1040,9 +1027,16 @@ class QwenImagePipeline(
             prompt = None
         if negative_prompt_embeds is not None:
             negative_prompt = None
-        true_cfg_scale = common_sampling_params.true_cfg_scale or true_cfg_scale
+        true_cfg_scale = common_sampling_params.true_cfg_scale or 4.0
         if common_sampling_params.guidance_scale_provided:
             guidance_scale = common_sampling_params.guidance_scale
+        else:
+            guidance_scale = 1.0
+
+        latents = req.collate_request_tensors("latents", None)
+        output_type = common_sampling_params.output_type or "pil"
+        attention_kwargs = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         ctx = self._prepare_generation_context(
             prompt=prompt,
@@ -1087,6 +1081,7 @@ class QwenImagePipeline(
         )
 
         self._current_timestep = None
+
         result = self._decode_latents(latents, height, width, output_type)
         return split_diffusion_output_by_request(
             result,

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -677,30 +677,7 @@ class QwenImageEditPipeline(
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        true_cfg_scale: float = 4.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 1.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 1024,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """Forward pass for image editing."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
@@ -731,32 +708,28 @@ class QwenImageEditPipeline(
             height = req.sampling_params.height
             width = req.sampling_params.width
         else:
-            # fallback to run pre-processing in pipeline (debug only)
-            image_size = image[0].size if isinstance(image, list) else image.size
-            calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
-            height = height or calculated_height
-            width = width or calculated_width
+            raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
-            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-
-            if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
-                image = self.image_processor.resize(image, calculated_height, calculated_width)
-                prompt_image = image
-                image = self.image_processor.preprocess(image, calculated_height, calculated_width)
-                image = image.unsqueeze(2)
-
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or 1024
+        generator = req.sampling_params.generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 1.0
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        prompt_embeds_mask = None
+        negative_prompt_embeds_mask = None
+        output_type = req.sampling_params.output_type or "pil"
+        attention_kwargs = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # 1. check inputs
         # 2. encode prompts
@@ -783,16 +756,9 @@ class QwenImageEditPipeline(
         self._current_timestep = None
         self._interrupt = False
 
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1
 
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
+        has_neg_prompt = negative_prompt is not None
 
         do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
         self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -5,7 +5,7 @@ import json
 import logging
 import os
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -623,30 +623,7 @@ class QwenImageEditPlusPipeline(
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor | None = None,
-        true_cfg_scale: float = 4.0,
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 1.0,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 1024,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """Forward pass for image editing with support for multiple images."""
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
@@ -674,55 +651,35 @@ class QwenImageEditPlusPipeline(
         ):
             condition_images = additional_information.get("condition_images")
             vae_images = additional_information.get("vae_images")
-            condition_image_sizes = additional_information.get("condition_image_sizes")
             vae_image_sizes = additional_information.get("vae_image_sizes")
             calculated_height = additional_information.get("calculated_height")
             calculated_width = additional_information.get("calculated_width")
-            height = req.sampling_params.height
-            width = req.sampling_params.width
+            height = req.sampling_params.height or calculated_height
+            width = req.sampling_params.width or calculated_width
+            image = None
         else:
-            # fallback to run pre-processing in pipeline (debug only)
-            if image is None:
-                raise ValueError("Image is required for QwenImageEditPlusPipeline")
+            raise RuntimeError("Missing preprocess images that should have been created by the preprocess function.")
 
-            if not isinstance(image, list):
-                image = [image]
-
-            image_size = image[0].size
-            calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1])
-            height = height or calculated_height
-            width = width or calculated_width
-
-            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-
-            condition_images = []
-            vae_images = []
-            condition_image_sizes = []
-            vae_image_sizes = []
-
-            for img in image:
-                image_width, image_height = img.size
-                condition_width, condition_height = calculate_dimensions(
-                    CONDITION_IMAGE_SIZE, image_width / image_height
-                )
-                vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height)
-                condition_image_sizes.append((condition_width, condition_height))
-                vae_image_sizes.append((vae_width, vae_height))
-                condition_images.append(self.image_processor.resize(img, condition_height, condition_width))
-                vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
-
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or 1024
+        generator = req.sampling_params.generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 1.0
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        prompt_embeds_mask = None
+        negative_prompt_embeds_mask = None
+        output_type = req.sampling_params.output_type or "pil"
+        attention_kwargs = None
+        callback_on_step_end_tensor_inputs = ["latents"]
 
         # 1. check inputs
         # 2. encode prompts
@@ -749,16 +706,9 @@ class QwenImageEditPlusPipeline(
         self._current_timestep = None
         self._interrupt = False
 
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1
 
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
+        has_neg_prompt = negative_prompt is not None
 
         do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
         self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -7,7 +7,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any, ClassVar, cast
+from typing import ClassVar, cast
 
 import numpy as np
 import PIL.Image
@@ -648,31 +648,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        true_cfg_scale: float = 4.0,
-        layers: int | None = 4,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float | None = None,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        prompt_embeds_mask: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds_mask: torch.Tensor | None = None,
-        output_type: str | None = "pil",
-        attention_kwargs: dict[str, Any] | None = None,
-        max_sequence_length: int = 1024,
-        resolution: int = 640,
-        cfg_normalize: bool = False,
-        use_en_prompt: bool = False,
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """Forward pass for image layered."""
 
         # 1. Get preprocessed image from request (pre-processing is done in DiffusionEngine)
@@ -688,26 +664,27 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
         negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
-        layers = req.sampling_params.layers if req.sampling_params.layers is not None else layers
-        resolution = req.sampling_params.resolution if req.sampling_params.resolution is not None else resolution
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        cfg_normalize = (
-            req.sampling_params.cfg_normalize if req.sampling_params.cfg_normalize is not None else cfg_normalize
-        )
-        use_en_prompt = (
-            req.sampling_params.use_en_prompt if req.sampling_params.use_en_prompt is not None else use_en_prompt
-        )
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        sigmas = req.sampling_params.sigmas or sigmas
-        generator = req.sampling_params.generator or generator
-        true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale
+        layers = req.sampling_params.layers
+        max_sequence_length = req.sampling_params.max_sequence_length or 1024
+        cfg_normalize = req.sampling_params.cfg_normalize
+        use_en_prompt = req.sampling_params.use_en_prompt
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
+        sigmas = req.sampling_params.sigmas
+        generator = req.sampling_params.generator
+        true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0
+        guidance_scale = None
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        prompt_embeds_mask = None
+        negative_prompt_embeds_mask = None
+        output_type = req.sampling_params.output_type or "pil"
+        attention_kwargs = None
 
         if not isinstance(first_prompt, str) and "preprocessed_image" in (
             additional_information := first_prompt.get("additional_information", {})
@@ -720,25 +697,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
             height = req.sampling_params.height
             width = req.sampling_params.width
         else:
-            # fallback to run pre-processing in pipeline (debug only)
-            if isinstance(image, PIL.Image.Image) and image.mode != "RGBA":
-                image = image.convert("RGBA")
-            image_size = image[0].size if isinstance(image, list) else image.size
-            assert resolution in [640, 1024], f"resolution must be either 640 or 1024, but got {resolution}"
-            calculated_width, calculated_height = calculate_dimensions(
-                resolution * resolution, image_size[0] / image_size[1]
-            )
-            height = calculated_height
-            width = calculated_width
-
-            height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
-
-            if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
-                image = self.image_processor.resize(image, calculated_height, calculated_width)
-                prompt_image = image
-                image = self.image_processor.preprocess(image, calculated_height, calculated_width)
-                image = image.unsqueeze(2)
-                image = image.to(dtype=self.text_encoder.dtype)
+            raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.")
 
         # 2. check inputs
         self.check_inputs(
@@ -761,16 +720,9 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as
         # 3. encode prompot & negative prompt
         if prompt is None or prompt == "" or prompt == " ":
             prompt = self.get_image_caption(prompt_image, use_en_prompt=use_en_prompt, device=self.device)
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = 1
 
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
+        has_neg_prompt = negative_prompt is not None
 
         if true_cfg_scale > 1 and not has_neg_prompt:
             logger.warning(

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -17,7 +17,6 @@ from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
 from vllm.logger import init_logger
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     QKVParallelLinear,
@@ -555,8 +554,8 @@ class QwenImageCrossAttention(nn.Module):
         self.query_num_heads = self.to_qkv.num_heads
         self.kv_num_heads = self.to_qkv.num_kv_heads
 
-        self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
-        self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_q = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
+        self.norm_k = nn.RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         self.inner_dim = out_dim if out_dim is not None else head_dim * self.total_num_heads
 
@@ -593,8 +592,8 @@ class QwenImageCrossAttention(nn.Module):
             prefix=f"{prefix}.to_out.0",
         )
 
-        self.norm_added_q = RMSNorm(head_dim, eps=eps)
-        self.norm_added_k = RMSNorm(head_dim, eps=eps)
+        self.norm_added_q = nn.RMSNorm(head_dim, eps=eps)
+        self.norm_added_k = nn.RMSNorm(head_dim, eps=eps)
 
         self.attn = Attention(
             num_heads=self.query_num_heads,
@@ -1007,7 +1006,7 @@ class QwenImageTransformer2DModel(CachedTransformer):
             quant_config=quant_config,
         )
 
-        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+        self.txt_norm = nn.RMSNorm(joint_attention_dim, eps=1e-6)
 
         # Entry projections (image/text) are kept full precision —
         # small sensitive layers at the network boundary (see #2728).

vllm_omni/diffusion/models/schedulers/__init__.py

@@ -2,11 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from vllm_omni.diffusion.models.schedulers.scheduling_dmd2_euler import DMD2EulerScheduler
+from vllm_omni.diffusion.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
 from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import (
     FlowUniPCMultistepScheduler,
 )
 
 __all__ = [
     "DMD2EulerScheduler",
+    "FlowMatchEulerDiscreteScheduler",
     "FlowUniPCMultistepScheduler",
 ]

vllm_omni/diffusion/models/schedulers/scheduling_flow_match_euler_discrete.py

@@ -0,0 +1,660 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# Originally from https://github.com/huggingface/diffusers/blob/v0.39.0/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py
+# Copyright 2025 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Flow Match Euler scheduler copied from Diffusers v0.39.0.
+
+Modified by vLLM-Omni contributors for local integration and RNG compatibility.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput, is_scipy_available, logging
+from diffusers.utils.torch_utils import randn_tensor
+
+if is_scipy_available():
+    import scipy.stats
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
+            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
+            denoising loop.
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin):
+    """
+    Euler scheduler.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        shift (`float`, defaults to 1.0):
+            The shift value for the timestep schedule.
+        use_dynamic_shifting (`bool`, defaults to False):
+            Whether to apply timestep shifting on-the-fly based on the image resolution.
+        base_shift (`float`, defaults to 0.5):
+            Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent
+            with desired output.
+        max_shift (`float`, defaults to 1.15):
+            Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be
+            more exaggerated or stylized.
+        base_image_seq_len (`int`, defaults to 256):
+            The base image sequence length.
+        max_image_seq_len (`int`, defaults to 4096):
+            The maximum image sequence length.
+        invert_sigmas (`bool`, defaults to False):
+            Whether to invert the sigmas.
+        shift_terminal (`float`, defaults to None):
+            The end value of the shifted timestep schedule.
+        use_karras_sigmas (`bool`, defaults to False):
+            Whether to use Karras sigmas for step sizes in the noise schedule during sampling.
+        use_exponential_sigmas (`bool`, defaults to False):
+            Whether to use exponential sigmas for step sizes in the noise schedule during sampling.
+        use_beta_sigmas (`bool`, defaults to False):
+            Whether to use beta sigmas for step sizes in the noise schedule during sampling.
+        time_shift_type (`str`, defaults to "exponential"):
+            The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".
+        stochastic_sampling (`bool`, defaults to False):
+            Whether to use stochastic sampling.
+    """
+
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 1.0,
+        use_dynamic_shifting: bool = False,
+        base_shift: float | None = 0.5,
+        max_shift: float | None = 1.15,
+        base_image_seq_len: int = 256,
+        max_image_seq_len: int = 4096,
+        invert_sigmas: bool = False,
+        shift_terminal: float = None,
+        use_karras_sigmas: bool = False,
+        use_exponential_sigmas: bool = False,
+        use_beta_sigmas: bool = False,
+        time_shift_type: Literal["exponential", "linear"] = "exponential",
+        stochastic_sampling: bool = False,
+    ):
+        if self.config.use_beta_sigmas and not is_scipy_available():
+            raise ImportError("Make sure to install scipy if you want to use beta sigmas.")
+        if (
+            sum(
+                [
+                    self.config.use_beta_sigmas,
+                    self.config.use_exponential_sigmas,
+                    self.config.use_karras_sigmas,
+                ]
+            )
+            > 1
+        ):
+            raise ValueError(
+                "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, "
+                "`config.use_karras_sigmas` can be used."
+            )
+        if time_shift_type not in {"exponential", "linear"}:
+            raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.")
+
+        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)
+
+        sigmas = timesteps / num_train_timesteps
+        if not use_dynamic_shifting:
+            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        self.timesteps = sigmas * num_train_timesteps
+
+        self._step_index = None
+        self._begin_index = None
+
+        self._shift = shift
+
+        self.sigmas = sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def shift(self):
+        """
+        The value used for shifting.
+        """
+        return self._shift
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`, defaults to `0`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def set_shift(self, shift: float):
+        """
+        Sets the shift value for the scheduler.
+
+        Args:
+            shift (`float`):
+                The shift value to be set.
+        """
+        self._shift = shift
+
+    def scale_noise(
+        self,
+        sample: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        noise: torch.FloatTensor | None = None,
+    ) -> torch.FloatTensor:
+        """
+        Forward process in flow-matching
+
+        Args:
+            sample (`torch.FloatTensor`):
+                The input sample.
+            timestep (`torch.FloatTensor`):
+                The current timestep in the diffusion chain.
+            noise (`torch.FloatTensor`):
+                The noise tensor.
+
+        Returns:
+            `torch.FloatTensor`:
+                A scaled input sample.
+        """
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)
+
+        if sample.device.type == "mps" and torch.is_floating_point(timestep):
+            # mps does not support float64
+            schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32)
+            timestep = timestep.to(sample.device, dtype=torch.float32)
+        else:
+            schedule_timesteps = self.timesteps.to(sample.device)
+            timestep = timestep.to(sample.device)
+
+        # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
+        if self.begin_index is None:
+            step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep]
+        elif self.step_index is not None:
+            # add_noise is called after first denoising step (for inpainting)
+            step_indices = [self.step_index] * timestep.shape[0]
+        else:
+            # add noise is called before first denoising step to create initial latent(img2img)
+            step_indices = [self.begin_index] * timestep.shape[0]
+
+        sigma = sigmas[step_indices].flatten()
+        while len(sigma.shape) < len(sample.shape):
+            sigma = sigma.unsqueeze(-1)
+
+        sample = sigma * noise + (1.0 - sigma) * sample
+
+        return sample
+
+    def _sigma_to_t(self, sigma) -> float:
+        return sigma * self.config.num_train_timesteps
+
+    def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+        """
+        Apply time shifting to the sigmas.
+
+        Args:
+            mu (`float`):
+                The mu parameter for the time shift.
+            sigma (`float`):
+                The sigma parameter for the time shift.
+            t (`torch.Tensor`):
+                The input timesteps.
+
+        Returns:
+            `torch.Tensor`:
+                The time-shifted timesteps.
+        """
+        if self.config.time_shift_type == "exponential":
+            return self._time_shift_exponential(mu, sigma, t)
+        elif self.config.time_shift_type == "linear":
+            return self._time_shift_linear(mu, sigma, t)
+
+    def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
+        r"""
+        Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config
+        value.
+
+        Reference:
+        https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51
+
+        Args:
+            t (`torch.Tensor`):
+                A tensor of timesteps to be stretched and shifted.
+
+        Returns:
+            `torch.Tensor`:
+                A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`.
+        """
+        one_minus_z = 1 - t
+        scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal)
+        stretched_t = 1 - (one_minus_z / scale_factor)
+        return stretched_t
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int | None = None,
+        device: str | torch.device = None,
+        sigmas: list[float] | None = None,
+        mu: float | None = None,
+        timesteps: list[float] | None = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`, *optional*):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            sigmas (`list[float]`, *optional*):
+                Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
+                automatically.
+            mu (`float`, *optional*):
+                Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
+                shifting.
+            timesteps (`list[float]`, *optional*):
+                Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
+                automatically.
+        """
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`")
+
+        if sigmas is not None and timesteps is not None:
+            if len(sigmas) != len(timesteps):
+                raise ValueError("`sigmas` and `timesteps` should have the same length")
+
+        if num_inference_steps is not None:
+            if (sigmas is not None and len(sigmas) != num_inference_steps) or (
+                timesteps is not None and len(timesteps) != num_inference_steps
+            ):
+                raise ValueError(
+                    "`sigmas` and `timesteps` should have the same length as num_inference_steps, "
+                    "if `num_inference_steps` is provided"
+                )
+        else:
+            num_inference_steps = len(sigmas) if sigmas is not None else len(timesteps)
+
+        self.num_inference_steps = num_inference_steps
+
+        # 1. Prepare default sigmas
+        is_timesteps_provided = timesteps is not None
+
+        if is_timesteps_provided:
+            timesteps = np.array(timesteps).astype(np.float32)
+
+        if sigmas is None:
+            if timesteps is None:
+                timesteps = np.linspace(
+                    self._sigma_to_t(self.sigma_max),
+                    self._sigma_to_t(self.sigma_min),
+                    num_inference_steps,
+                )
+            sigmas = timesteps / self.config.num_train_timesteps
+        else:
+            sigmas = np.array(sigmas).astype(np.float32)
+            num_inference_steps = len(sigmas)
+
+        # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of
+        #    "exponential" or "linear" type is applied
+        if self.config.use_dynamic_shifting:
+            sigmas = self.time_shift(mu, 1.0, sigmas)
+        else:
+            sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)
+
+        # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
+        if self.config.shift_terminal:
+            sigmas = self.stretch_shift_to_terminal(sigmas)
+
+        # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
+        if self.config.use_karras_sigmas:
+            sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+        elif self.config.use_exponential_sigmas:
+            sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+        elif self.config.use_beta_sigmas:
+            sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+
+        # 5. Convert sigmas and timesteps to tensors and move to specified device
+        sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device)
+        if not is_timesteps_provided:
+            timesteps = sigmas * self.config.num_train_timesteps
+        else:
+            timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device)
+
+        # 6. Append the terminal sigma value.
+        #    If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
+        #    `invert_sigmas` flag can be set to `True`. This case is only required in Mochi
+        if self.config.invert_sigmas:
+            sigmas = 1.0 - sigmas
+            timesteps = sigmas * self.config.num_train_timesteps
+            sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)])
+        else:
+            sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+
+        self.timesteps = timesteps
+        self.sigmas = sigmas
+        self._step_index = None
+        self._begin_index = None
+
+    def index_for_timestep(
+        self,
+        timestep: float | torch.FloatTensor,
+        schedule_timesteps: torch.FloatTensor | None = None,
+    ) -> int:
+        """
+        Get the index for the given timestep.
+
+        Args:
+            timestep (`float` or `torch.FloatTensor`):
+                The timestep to find the index for.
+            schedule_timesteps (`torch.FloatTensor`, *optional*):
+                The schedule timesteps to validate against. If `None`, the scheduler's timesteps are used.
+
+        Returns:
+            `int`:
+                The index of the timestep.
+        """
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep: float | torch.FloatTensor) -> None:
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: torch.Generator | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        return_dict: bool = True,
+    ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
+        process from the learned model outputs (most often the predicted noise).
+
+        Args:
+            model_output (`torch.FloatTensor`):
+                The direct output from learned diffusion model.
+            timestep (`float`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.FloatTensor`):
+                A current instance of a sample created by the diffusion process.
+            s_churn (`float`):
+            s_tmin  (`float`):
+            s_tmax  (`float`):
+            s_noise (`float`, defaults to 1.0):
+                Scaling factor for noise added to the sample.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            per_token_timesteps (`torch.Tensor`, *optional*):
+                The timesteps for each token in the sample.
+            return_dict (`bool`, defaults to `True`):
+                Whether or not to return a
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.
+
+        Returns:
+            [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`:
+                If return_dict is `True`,
+                [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned,
+                otherwise a tuple is returned where the first element is the sample tensor.
+        """
+
+        if isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, torch.LongTensor):
+            raise ValueError(
+                (
+                    "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
+                    " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
+                    " one of the `scheduler.timesteps` as a timestep."
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+
+        if per_token_timesteps is not None:
+            per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps
+
+            sigmas = self.sigmas[:, None, None]
+            lower_mask = sigmas < per_token_sigmas[None] - 1e-6
+            lower_sigmas = lower_mask * sigmas
+            lower_sigmas, _ = lower_sigmas.max(dim=0)
+
+            current_sigma = per_token_sigmas[..., None]
+            next_sigma = lower_sigmas[..., None]
+            dt = current_sigma - next_sigma
+        else:
+            sigma_idx = self.step_index
+            sigma = self.sigmas[sigma_idx]
+            sigma_next = self.sigmas[sigma_idx + 1]
+
+            current_sigma = sigma
+            next_sigma = sigma_next
+            dt = sigma_next - sigma
+
+        if self.config.stochastic_sampling:
+            x0 = sample - current_sigma * model_output
+            noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype)
+            prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
+        else:
+            prev_sample = sample + dt * model_output
+
+        # upon completion increase step index by one
+        self._step_index += 1
+        if per_token_timesteps is None:
+            # Cast sample back to model compatible dtype
+            prev_sample = prev_sample.to(model_output.dtype)
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+    def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
+        """
+        Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative
+        Models](https://huggingface.co/papers/2206.00364).
+
+        Args:
+            in_sigmas (`torch.Tensor`):
+                The input sigma values to be converted.
+            num_inference_steps (`int`):
+                The number of inference steps to generate the noise schedule for.
+
+        Returns:
+            `torch.Tensor`:
+                The converted sigma values following the Karras noise schedule.
+        """
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = sigma_min ** (1 / rho)
+        max_inv_rho = sigma_max ** (1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
+    def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
+        """
+        Construct an exponential noise schedule.
+
+        Args:
+            in_sigmas (`torch.Tensor`):
+                The input sigma values to be converted.
+            num_inference_steps (`int`):
+                The number of inference steps to generate the noise schedule for.
+
+        Returns:
+            `torch.Tensor`:
+                The converted sigma values following an exponential schedule.
+        """
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps))
+        return sigmas
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
+    def _convert_to_beta(
+        self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6
+    ) -> torch.Tensor:
+        """
+        Construct a beta noise schedule as proposed in [Beta Sampling is All You
+        Need](https://huggingface.co/papers/2407.12173).
+
+        Args:
+            in_sigmas (`torch.Tensor`):
+                The input sigma values to be converted.
+            num_inference_steps (`int`):
+                The number of inference steps to generate the noise schedule for.
+            alpha (`float`, *optional*, defaults to `0.6`):
+                The alpha parameter for the beta distribution.
+            beta (`float`, *optional*, defaults to `0.6`):
+                The beta parameter for the beta distribution.
+
+        Returns:
+            `torch.Tensor`:
+                The converted sigma values following a beta distribution schedule.
+        """
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        sigmas = np.array(
+            [
+                sigma_min + (ppf * (sigma_max - sigma_min))
+                for ppf in [
+                    scipy.stats.beta.ppf(timestep, alpha, beta)
+                    for timestep in 1 - np.linspace(0, 1, num_inference_steps)
+                ]
+            ]
+        )
+        return sigmas
+
+    def _time_shift_exponential(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+        return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+    def _time_shift_linear(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
+        return mu / (mu + (1 / t - 1) ** sigma)
+
+    def __len__(self) -> int:
+        return self.config.num_train_timesteps

vllm_omni/diffusion/models/sd3/pipeline_sd3.py

@@ -628,57 +628,39 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
 
         return latents
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] = "",
-        prompt_2: str | list[str] = "",
-        prompt_3: str | list[str] = "",
-        negative_prompt: str | list[str] = "",
-        negative_prompt_2: str | list[str] = "",
-        negative_prompt_3: str | list[str] = "",
-        height: int | None = None,
-        width: int | None = None,
-        num_inference_steps: int = 28,
-        sigmas: list[float] | None = None,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        pooled_prompt_embeds: torch.Tensor | None = None,
-        negative_pooled_prompt_embeds: torch.Tensor | None = None,
-        max_sequence_length: int = 256,
-    ) -> list[DiffusionOutput]:
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
         sampling_params_list = req.sampling_params_list
         common_sampling_params = sampling_params_list[0]
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
-        negative_prompt = [
-            "" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts
-        ] or negative_prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
+        negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
+        prompt_2 = ""
+        prompt_3 = ""
+        negative_prompt_2 = ""
+        negative_prompt_3 = ""
+        prompt_embeds = None
+        negative_prompt_embeds = None
 
         height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor
         width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor
-        sigmas = common_sampling_params.sigmas or sigmas
-        max_sequence_length = common_sampling_params.max_sequence_length or max_sequence_length
-        num_inference_steps = common_sampling_params.num_inference_steps or num_inference_steps
+        sigmas = common_sampling_params.sigmas
+        max_sequence_length = common_sampling_params.max_sequence_length or 256
+        num_inference_steps = common_sampling_params.num_inference_steps or 28
         num_images_per_prompt = (
-            common_sampling_params.num_outputs_per_prompt
-            if common_sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1
         )
-        if generator is None:
-            generator = req.collate_request_generators(num_images_per_prompt, generator)
-        latents = req.collate_request_tensors("latents", latents)
+        generator = req.collate_request_generators(num_images_per_prompt, None)
+        latents = req.collate_request_tensors("latents", None)
+        output_type = common_sampling_params.output_type or "pil"
+
         prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
             req.prompts,
             {
                 "prompt_embeds": prompt_embeds,
                 "negative_prompt_embeds": negative_prompt_embeds,
-                "pooled_prompt_embeds": pooled_prompt_embeds,
-                "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds,
+                "pooled_prompt_embeds": None,
+                "negative_pooled_prompt_embeds": None,
             },
         )
         prompt_embeds = prompt_fields["prompt_embeds"]
@@ -717,9 +699,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         self._current_timestep = None
         self._interrupt = False
 
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
+        if prompt is not None:
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]
@@ -771,7 +751,7 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         )
 
         self._current_timestep = None
-        if self.output_type == "latent":
+        if output_type == "latent":
             image = latents
         else:
             # Ensure the latents are the same dtype as the VAE for decode

vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py

@@ -938,8 +938,10 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         )
         logger.info("Text generation: %d chars", len(text_output))
         return DiffusionOutput(
-            output=text_output,
-            custom_output={"text_output": text_output},
+            output={
+                "payload": {"text": text_output},
+                "metadata": {"text": {"text_output": text_output}},
+            },
         )
 
     # -----------------------------------------------------------------------
@@ -1457,10 +1459,15 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
 
         images = _to_pil(image_prediction)
         img = images[0] if images else None
-        custom = {}
+        metadata = {}
         if think_text:
-            custom["think_text"] = think_text
-        return DiffusionOutput(output=img, custom_output=custom)
+            metadata["text"] = {"think_text": think_text}
+        return DiffusionOutput(
+            output={
+                "payload": {"image": img},
+                "metadata": metadata,
+            }
+        )
 
     # -----------------------------------------------------------------------
     # Weight loading

vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py

@@ -164,11 +164,11 @@ class SoulXPreprocessPipeline(nn.Module, SupportAudioInput, SupportsComponentDis
                 sample_rate=self.target_sr,
             )
 
-        # Non-None placeholder so DiffusionEngine forwards custom_output to stage-1
-        # (output=None triggers an early return that drops handoff).
         return DiffusionOutput(
-            output=torch.zeros(0),
-            custom_output={SOULX_PREPROCESSED_KEY: payload},
+            output={
+                "payload": {},
+                "metadata": {SOULX_PREPROCESSED_KEY: payload},
+            },
             to_cpu=True,
         )
 

vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py

@@ -4,6 +4,7 @@ from collections.abc import Callable
 from contextlib import contextmanager
 from typing import Any, ClassVar
 
+import numpy as np
 import torch
 import torch.nn as nn
 from vllm.logger import init_logger
@@ -30,6 +31,21 @@ _SOULX_PROFILER_TARGETS = [
 ]
 
 
+def convert_soulx_audio_output_to_numpy(audio: torch.Tensor | dict[str, Any]) -> np.ndarray | dict[str, Any]:
+    if isinstance(audio, dict) and isinstance(audio.get("payload"), dict):
+        payload = dict(audio["payload"])
+        audio_payload = payload.get("audio")
+        if isinstance(audio_payload, torch.Tensor):
+            payload["audio"] = audio_payload.detach().cpu().float().numpy()
+        return {
+            "payload": payload,
+            "metadata": audio.get("metadata") or {},
+        }
+    if isinstance(audio, torch.Tensor):
+        return audio.detach().cpu().float().numpy()
+    return audio
+
+
 class FlowMatchingAudioPipeline(
     nn.Module,
     CFGParallelMixin,
@@ -326,7 +342,7 @@ class FlowMatchingAudioPipeline(
         req: OmniDiffusionRequest,
         *,
         kind: str,
-        custom_output_key: str,
+        metadata_key: str,
         infer_batch_fn: Callable[..., tuple[torch.Tensor, int]],
         prepare_extra_args: Callable[[dict[str, Any], Any], dict[str, Any]] | None = None,
     ) -> DiffusionOutput:
@@ -353,7 +369,9 @@ class FlowMatchingAudioPipeline(
             )
 
         return DiffusionOutput(
-            output=audio,
-            custom_output={custom_output_key: pitch_shift},
+            output={
+                "payload": {"audio": audio},
+                "metadata": {"audio": {metadata_key: pitch_shift}},
+            },
             stage_durations=self._profiler_stage_durations() or {},
         )

vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py

@@ -19,6 +19,7 @@ from vllm_omni.diffusion.models.soulx_singer.modules import (
 )
 from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import (
     FlowMatchingAudioPipeline,
+    convert_soulx_audio_output_to_numpy,
 )
 from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import (
     SOULX_PREPROCESSED_KEY,
@@ -96,7 +97,7 @@ def get_soulxsinger_svc_pre_process_func(od_config: OmniDiffusionConfig):
 
 def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig):
     def post_process_func(audio: torch.Tensor):
-        return audio.detach().cpu().float().numpy()
+        return convert_soulx_audio_output_to_numpy(audio)
 
     return post_process_func
 
@@ -400,7 +401,7 @@ class PipelineSoulXSingerSVC(FlowMatchingAudioPipeline):
         return self._forward_batch_from_request(
             req,
             kind="svc",
-            custom_output_key="pitch_shift",
+            metadata_key="pitch_shift",
             infer_batch_fn=self.infer_svc_batch,
         )
 

vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py

@@ -15,6 +15,7 @@ from vllm_omni.diffusion.models.soulx_singer.modules import (
 )
 from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import (
     FlowMatchingAudioPipeline,
+    convert_soulx_audio_output_to_numpy,
 )
 from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import (
     SOULX_PREPROCESSED_KEY,
@@ -97,7 +98,7 @@ def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig):
     """Convert pipeline audio tensor output for offline consumers."""
 
     def post_process_func(audio: torch.Tensor):
-        return audio.detach().cpu().float().numpy()
+        return convert_soulx_audio_output_to_numpy(audio)
 
     return post_process_func
 
@@ -343,7 +344,7 @@ class PipelineSoulXSingerSVS(FlowMatchingAudioPipeline):
         return self._forward_batch_from_request(
             req,
             kind="svs",
-            custom_output_key="f0_shift",
+            metadata_key="f0_shift",
             infer_batch_fn=self.infer_svs_batch,
             prepare_extra_args=lambda extra_args, _sampling: normalize_svs_control_extra_args(extra_args),
         )

vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py

@@ -13,7 +13,7 @@ SOULX_PRECOMPUTED_KEYS_BY_KIND: dict[str, tuple[str, ...]] = {
 
 
 # Key under which the SoulX preprocess payload is stored in
-# prompt["additional_information"] or DiffusionOutput.custom_output.
+# prompt["additional_information"] or diffusion output metadata.
 SOULX_PREPROCESSED_KEY = "soulx_preprocessed"
 
 

vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py

@@ -385,37 +385,21 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         latents = latents * self.scheduler.init_noise_sigma
         return latents
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        audio_end_in_s: float | None = None,
-        audio_start_in_s: float = 0.0,
-        guidance_scale: float = 7.0,
-        num_waveforms_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        output_type: str = "np",
-    ) -> DiffusionOutput:
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         """
         Generate audio from text prompt.
 
         Args:
-            req: OmniDiffusionRequest containing generation parameters
-            prompt: Text prompt for audio generation
-            negative_prompt: Negative prompt for CFG
-            audio_end_in_s: Audio end time in seconds (max ~47s for stable-audio-open-1.0)
-            audio_start_in_s: Audio start time in seconds
-            guidance_scale: CFG scale
-            num_waveforms_per_prompt: Number of audio outputs per prompt
-            generator: Random generator for reproducibility
-            latents: Pre-generated latents
-            prompt_embeds: Pre-computed prompt embeddings
-            negative_prompt_embeds: Pre-computed negative prompt embeddings
-            output_type: Output format ("np", "pt", or "latent")
+            req: OmniDiffusionRequest containing generation parameters.
+                The `req.sampling_params.extra_args` can include the following keys:
+                - audio_start_in_s (`float`, *optional*, defaults to 0.0):
+                    Start time of the audio in seconds.
+                - audio_end_in_s (`float`, *optional*):
+                    End time of the audio in seconds.
+                - num_waveforms_per_prompt (`int`, *optional*, defaults to 1):
+                    Number of audio outputs per prompt.
+                - output_type (`str`, *optional*, defaults to "np"):
+                    Output format ("np", "pt", or "latent").
 
         Returns:
             DiffusionOutput containing generated audio
@@ -423,28 +407,33 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         # Extract from request
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
-        num_inference_steps = req.sampling_params.num_inference_steps
-        if num_inference_steps is None:
-            num_inference_steps = 100  # Default steps
-        num_inference_steps = req.sampling_params.num_inference_steps
-        if num_inference_steps is None:
-            num_inference_steps = 50  # Default steps
+
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 7.0
 
-        if generator is None:
-            generator = req.sampling_params.generator
+        generator = req.sampling_params.generator
         if generator is None and req.sampling_params.seed is not None:
             generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed)
+        latents = req.sampling_params.latents
 
         # Get audio duration from request extra params or defaults
-        audio_start_in_s = req.sampling_params.extra_args.get("audio_start_in_s", audio_start_in_s)
-        audio_end_in_s = req.sampling_params.extra_args.get("audio_end_in_s", audio_end_in_s)
+        audio_start_in_s: float = req.sampling_params.extra_args.get("audio_start_in_s", 0.0)
+        audio_end_in_s: float | None = req.sampling_params.extra_args.get("audio_end_in_s", None)
+        num_waveforms_per_prompt: int = req.sampling_params.extra_args.get(
+            "num_waveforms_per_prompt", req.sampling_params.num_outputs_per_prompt or 1
+        )
+        output_type = req.sampling_params.output_type or "np"
 
         # Calculate audio length
         downsample_ratio = self.vae.hop_length
@@ -474,12 +463,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         )
 
         # Determine batch size
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
+        batch_size = len(prompt)
 
         device = self.device
         do_classifier_free_guidance = guidance_scale > 1.0
@@ -500,7 +484,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
             audio_start_in_s,
             audio_end_in_s,
             device,
-            do_classifier_free_guidance and (negative_prompt is not None or negative_prompt_embeds is not None),
+            do_classifier_free_guidance and negative_prompt is not None,
             batch_size,
         )
 
@@ -515,7 +499,7 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscov
         )
 
         # Handle CFG without negative prompt
-        if do_classifier_free_guidance and negative_prompt_embeds is None and negative_prompt is None:
+        if do_classifier_free_guidance and negative_prompt is None:
             negative_text_audio_duration_embeds = torch.zeros_like(text_audio_duration_embeds)
             text_audio_duration_embeds = torch.cat(
                 [negative_text_audio_duration_embeds, text_audio_duration_embeds],

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -187,7 +187,7 @@ def get_wan22_post_process_func(
     ):
         if output_type == "latent":
             return video
-        custom_output = {}
+        video_metadata = {}
         if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
             video, multiplier = interpolate_video_tensor(
                 video,
@@ -195,10 +195,10 @@ def get_wan22_post_process_func(
                 scale=sampling_params.frame_interpolation_scale,
                 model_path=sampling_params.frame_interpolation_model_path,
             )
-            custom_output["video_fps_multiplier"] = multiplier
+            video_metadata["video_fps_multiplier"] = multiplier
         return {
-            "video": video_processor.postprocess_video(video, output_type=output_type),
-            "custom_output": custom_output,
+            "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)},
+            "metadata": {"video": video_metadata} if video_metadata else {},
         }
 
     return post_process_func
@@ -540,38 +540,27 @@ class Wan22Pipeline(
 
         return latents
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | None = None,
-        negative_prompt: str | None = None,
-        height: int = 480,
-        width: int = 832,
-        num_inference_steps: int = 40,
-        guidance_scale: float | tuple[float, float] = 4.0,
-        frame_num: int = 81,
-        output_type: str | None = "np",
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        attention_kwargs: dict | None = None,
-        **kwargs,
-    ) -> DiffusionOutput:
-        # Get parameters from request or arguments
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        prompt: str | None = None
+        negative_prompt: str | None = None
+        prompt_embeds: torch.Tensor | None = None
+        negative_prompt_embeds: torch.Tensor | None = None
         if len(req.prompts) > 1:
             raise ValueError(
                 """This model only supports a single prompt, not a batched request.""",
                 """Please pass in a single prompt object or string, or a single-item list.""",
             )
-        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
-            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
-            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
-        if prompt is None and prompt_embeds is None:
-            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
+        if len(req.prompts) == 1:
+            first_prompt = req.prompts[0]
+            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+            negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num
+        if not prompt:
+            raise ValueError("Prompt is required for Wan2.2 generation.")
+
+        height = req.sampling_params.height or 480
+        width = req.sampling_params.width or 832
+        num_frames = req.sampling_params.num_frames or 81
 
         # Ensure dimensions are compatible with VAE and patch size
         # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size)
@@ -579,11 +568,16 @@ class Wan22Pipeline(
         mod_value = self.vae_scale_factor_spatial * patch_size[1]  # 16*2=32 for TI2V, 8*2=16 for I2V
         height = (height // mod_value) * mod_value
         width = (width // mod_value) * mod_value
-        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
         # Respect per-request guidance_scale when explicitly provided.
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 4.0
+
+        output_type = req.sampling_params.output_type or "np"
+        attention_kwargs: dict | None = None
 
         guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
         guidance_high = (
@@ -634,8 +628,7 @@ class Wan22Pipeline(
             dtype = self.text_encoder.dtype
 
         # Seed / generator
-        if generator is None:
-            generator = req.sampling_params.generator
+        generator = req.sampling_params.generator
         if generator is None and req.sampling_params.seed is not None:
             generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
@@ -644,24 +637,16 @@ class Wan22Pipeline(
             current_omni_platform.synchronize()
             _t_pipeline_start = time.perf_counter()
             _t_text_enc_start = _t_pipeline_start
-        if prompt_embeds is None:
-            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
-                max_sequence_length=req.sampling_params.max_sequence_length or 512,
-                device=device,
-                dtype=dtype,
-            )
-        else:
-            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
-            if negative_prompt_embeds is not None:
-                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
-            elif guidance_low > 1.0 or guidance_high > 1.0:
-                raise ValueError(
-                    "negative_prompt_embeds must be provided when prompt_embeds are given and guidance > 1."
-                )
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
+            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+            max_sequence_length=req.sampling_params.max_sequence_length or 512,
+            device=device,
+            dtype=dtype,
+        )
+
         if DEBUG_PERF:
             current_omni_platform.synchronize()
             _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -65,7 +65,7 @@ def get_wan22_i2v_post_process_func(
     ):
         if output_type == "latent":
             return video
-        custom_output = {}
+        video_metadata = {}
         if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
             video, multiplier = interpolate_video_tensor(
                 video,
@@ -73,10 +73,10 @@ def get_wan22_i2v_post_process_func(
                 scale=sampling_params.frame_interpolation_scale,
                 model_path=sampling_params.frame_interpolation_model_path,
             )
-            custom_output["video_fps_multiplier"] = multiplier
+            video_metadata["video_fps_multiplier"] = multiplier
         return {
-            "video": video_processor.postprocess_video(video, output_type=output_type),
-            "custom_output": custom_output,
+            "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)},
+            "metadata": {"video": video_metadata} if video_metadata else {},
         }
 
     return post_process_func
@@ -429,66 +429,59 @@ class Wan22I2VPipeline(
         quant_config = getattr(self.od_config, "quantization_config", None)
         return create_transformer_from_config(config, quant_config=quant_config)
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | None = None,
-        negative_prompt: str | None = None,
-        image: PIL.Image.Image | torch.Tensor | None = None,
-        height: int = 480,
-        width: int = 832,
-        num_inference_steps: int = 40,
-        guidance_scale: float | tuple[float, float] = 5.0,
-        frame_num: int = 81,
-        output_type: str | None = "np",
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        image_embeds: torch.Tensor | None = None,
-        last_image: PIL.Image.Image | torch.Tensor | None = None,
-        attention_kwargs: dict | None = None,
-        **kwargs,
-    ) -> DiffusionOutput:
-        # Get parameters from request or arguments
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
+        prompt: str | None = None
+        negative_prompt: str | None = None
+        prompt_embeds: torch.Tensor | None = None
+        negative_prompt_embeds: torch.Tensor | None = None
+        image_embeds: torch.Tensor | None = None
         if len(req.prompts) > 1:
             raise ValueError(
                 """This model only supports a single prompt, not a batched request.""",
                 """Please pass in a single prompt object or string, or a single-item list.""",
             )
-        if len(req.prompts) == 1:  # If req.prompt is empty, default to prompt & neg_prompt in param list
-            prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
-            negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
-        if prompt is None and prompt_embeds is None:
-            raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
+        if len(req.prompts) == 1:
+            first_prompt = req.prompts[0]
+            prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "")
+            negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt")
+
+        if not prompt:
+            raise ValueError("Prompt is required for Wan2.2 generation.")
 
         # Get image from request
-        if image is None:
-            multi_modal_data = (
-                req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
-            )
-            raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
-            if raw_image is None:
-                raise ValueError("Image is required for I2V generation.")
-            if isinstance(raw_image, list):
-                if len(raw_image) > 1:
-                    logger.warning(
-                        """Received a list of image. Only a single image is supported by this model."""
-                        """Taking only the first image for now."""
-                    )
-                raw_image = raw_image[0]
-            if isinstance(raw_image, str):
-                image = PIL.Image.open(raw_image)
-            else:
-                image = cast(PIL.Image.Image | torch.Tensor, raw_image)
+        multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
+        raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
+        if raw_image is None:
+            raise ValueError("Image is required for I2V generation.")
+        if isinstance(raw_image, list):
+            if len(raw_image) > 1:
+                logger.warning(
+                    """Received a list of image. Only a single image is supported by this model."""
+                    """Taking only the first image for now."""
+                )
+            raw_image = raw_image[0]
+        if isinstance(raw_image, str):
+            image = PIL.Image.open(raw_image)
+        else:
+            image = cast(PIL.Image.Image | torch.Tensor, raw_image)
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_frames = req.sampling_params.num_frames or frame_num
-        num_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        last_image: PIL.Image.Image | torch.Tensor | None = None
+        if multi_modal_data is not None:
+            last_image = multi_modal_data.get("last_image", None)
+
+        height = req.sampling_params.height or 480
+        width = req.sampling_params.width or 832
+        num_frames = req.sampling_params.num_frames or 81
+        num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps
 
         # Respect per-request guidance_scale when explicitly provided.
         if req.sampling_params.guidance_scale_provided:
             guidance_scale = req.sampling_params.guidance_scale
+        else:
+            guidance_scale = 5.0
+
+        output_type = req.sampling_params.output_type or "np"
+        attention_kwargs: dict | None = None
 
         # Handle guidance scales
         guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0]
@@ -533,8 +526,7 @@ class Wan22I2VPipeline(
         dtype = self.transformer.dtype
 
         # Generator setup
-        if generator is None:
-            generator = req.sampling_params.generator
+        generator = req.sampling_params.generator
         if generator is None and req.sampling_params.seed is not None:
             generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
 
@@ -544,20 +536,15 @@ class Wan22I2VPipeline(
             _t_pipeline_start = time.perf_counter()
             _t_text_enc_start = _t_pipeline_start
 
-        if prompt_embeds is None:
-            prompt_embeds, negative_prompt_embeds = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
-                num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
-                max_sequence_length=req.sampling_params.max_sequence_length or 512,
-                device=device,
-                dtype=dtype,
-            )
-        else:
-            prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
-            if negative_prompt_embeds is not None:
-                negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0,
+            num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
+            max_sequence_length=req.sampling_params.max_sequence_length or 512,
+            device=device,
+            dtype=dtype,
+        )
 
         if DEBUG_PERF:
             current_omni_platform.synchronize()

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -24,7 +24,7 @@ from typing import Any, ClassVar
 import PIL.Image
 import torch
 import torch.nn as nn
-from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.image_processor import VaeImageProcessor
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import logging
 from diffusers.utils.torch_utils import randn_tensor
@@ -410,124 +410,21 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
     def interrupt(self):
         return self._interrupt
 
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        image: PipelineImageInput = None,
-        strength: float = 0.6,
-        height: int = 1024,
-        width: int = 1024,
-        num_inference_steps: int = 50,
-        sigmas: list[float] | None = None,
-        guidance_scale: float = 5.0,
-        cfg_normalization: bool = False,
-        cfg_truncation: float = 1.0,
-        negative_prompt: str | list[str] | None = None,
-        num_images_per_prompt: int = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.FloatTensor | None = None,
-        prompt_embeds: list[torch.FloatTensor] | None = None,
-        negative_prompt_embeds: list[torch.FloatTensor] | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: dict[str, Any] | None = None,
-        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
-        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ) -> DiffusionOutput:
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            image (`PipelineImageInput`, *optional*):
-                The image to use for img2img generation. If provided, the pipeline
-                will perform img2img instead of text-to-image.
-            strength (`float`, *optional*, defaults to 0.6):
-                Indicates extent to transform the reference `image`. Must be between 0 and 1.
-            height (`int`, *optional*, defaults to 1024):
-                The height in pixels of the generated image.
-            width (`int`, *optional*, defaults to 1024):
-                The width in pixels of the generated image.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`list[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            guidance_scale (`float`, *optional*, defaults to 5.0):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                0`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            cfg_normalization (`bool`, *optional*, defaults to False):
-                Whether to apply configuration normalization.
-            cfg_truncation (`float`, *optional*, defaults to 1.0):
-                The truncation value for configuration.
-            negative_prompt (`str` or `list[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
-                less than or equal to `0`).
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.FloatTensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`list[torch.FloatTensor]`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a plain
-                tuple.
-            joint_attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`list`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int`, *optional*, defaults to 512):
-                Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if
-            `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
-            generated images.
-        """
+    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
         # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "")
         # TODO: May be some data formatting operations on the API side. Hack for now.
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
+        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts]
+
         if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts):
             negative_prompt = None
         elif req.prompts:
             negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts]
 
-        # Handle img2img: extract image from request
-        if image is None and req.prompts:
+        prompt_embeds = None
+        negative_prompt_embeds = None
+
+        image = None
+        if req.prompts:
             if len(req.prompts) > 1:
                 logger.warning(
                     "This model only supports a single prompt for img2img, not a batched request. "
@@ -542,10 +439,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
                     else:
                         image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else raw_image
 
-        # strength is currently only applicable for Z-Image I2I; other pipelines ignore this parameter
         explicit_strength = req.sampling_params.strength is not None
-        if explicit_strength:
-            strength = req.sampling_params.strength
+        strength = req.sampling_params.strength if explicit_strength else 0.6
         if explicit_strength and image is None:
             logger.warning(
                 "strength parameter (%.2f) is only applicable for image-to-image (I2I) generation. "
@@ -556,20 +451,24 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
         if image is not None and strength is not None and (strength < 0 or strength > 1):
             raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}")
 
-        height = req.sampling_params.height or height
-        width = req.sampling_params.width or width
-        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
+        height = req.sampling_params.height or 1024
+        width = req.sampling_params.width or 1024
+        num_inference_steps = req.sampling_params.num_inference_steps or 50
         generator = req.sampling_params.generator
-        sigmas = req.sampling_params.sigmas or sigmas
-        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
-        guidance_scale = (
-            req.sampling_params.guidance_scale if req.sampling_params.guidance_rescale is not None else guidance_scale
-        )
+        sigmas = req.sampling_params.sigmas
+        max_sequence_length = req.sampling_params.max_sequence_length or 512
+        guidance_scale = req.sampling_params.guidance_scale
         num_images_per_prompt = (
-            req.sampling_params.num_outputs_per_prompt
-            if req.sampling_params.num_outputs_per_prompt > 0
-            else num_images_per_prompt
+            req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1
         )
+        latents = req.sampling_params.latents
+
+        cfg_normalization = req.sampling_params.cfg_normalize
+        cfg_truncation = req.sampling_params.extra_args.get("cfg_truncation", 1.0)
+        joint_attention_kwargs: dict[str, Any] | None = None
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None
+        callback_on_step_end_tensor_inputs = ["latents"]
+        output_type = req.sampling_params.output_type or "pil"
 
         vae_scale = self.vae_scale_factor * 2
         if height % vae_scale != 0:
@@ -591,33 +490,20 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponen
         self._cfg_normalization = cfg_normalization
         self._cfg_truncation = cfg_truncation
         # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = len(prompt_embeds)
+        batch_size = len(prompt)
 
-        # If prompt_embeds is provided and prompt is None, skip encoding
-        if prompt_embeds is not None and prompt is None:
-            if self.do_classifier_free_guidance and negative_prompt_embeds is None:
-                raise ValueError(
-                    "When `prompt_embeds` is provided without `prompt`, "
-                    "`negative_prompt_embeds` must also be provided for classifier-free guidance."
-                )
-        else:
-            (
-                prompt_embeds,
-                negative_prompt_embeds,
-            ) = self.encode_prompt(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                do_classifier_free_guidance=self.do_classifier_free_guidance,
-                prompt_embeds=prompt_embeds,
-                negative_prompt_embeds=negative_prompt_embeds,
-                device=device,
-                max_sequence_length=max_sequence_length,
-            )
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            device=device,
+            max_sequence_length=max_sequence_length,
+        )
 
         # 4. Prepare latent variables
         num_channels_latents = self.transformer.in_channels