vllm-omni/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Cosmos3 text/image/video/sound/action pipeline for vllm-omni.

One pipeline class serves the Cosmos3 family modes. Output modality is selected
mainly by ``prompt["modalities"]``:

* ``"image"`` selects T2I (text-to-image) and forces a single visual frame.
* ``"video"`` or omitted modalities select video generation.
* ``"audio"`` is accepted for compatibility but does not request sound by
  itself; sound is enabled with ``generate_sound`` or ``sound_gen``.

Video generation is further specialized by inputs and extra args:

* no image/video input: T2V (text-to-video).
* ``multi_modal_data["image"]``: I2V (image-to-video).
* ``multi_modal_data["video"]`` with no action/transfer mode: V2V
  (video-to-video).
* transfer hints (``edge``, ``blur``, ``depth``, ``seg``, or ``wsm``): control
  transfer video generation.
* ``action_mode``: action-capable video generation. RoboLab/OpenPI observation
  payloads in ``extra_args["robot_obs"]`` or ``extra_args["observation"]``
  bypass normal video output and return an action-only payload/metadata
  envelope.

Generated sound is video-only, cannot be combined with action or transfer, and
is produced from sound latents rather than from ``multi_modal_data["audio"]``.
"""

from __future__ import annotations

import json
import math
import os
import time
from collections.abc import Iterable, Mapping
from contextlib import nullcontext
from dataclasses import fields
from typing import Any, ClassVar

import numpy as np
import PIL.Image
import torch
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from torch import nn
from transformers import AutoTokenizer
from vllm.logger import init_logger
from vllm.model_executor.models.utils import AutoWeightsLoader

from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
    get_classifier_free_guidance_rank,
    get_classifier_free_guidance_world_size,
)
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.interface import (
    ReferenceVideoDecodeSpec,
    SupportImageInput,
    SupportsComponentDiscovery,
)
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
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,
)
from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.utils.video_decode import decode_path_video_frames, is_video_file_path
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.entrypoints.openai.video_api_utils import positive_float
from vllm_omni.experimental.world_models.adapters.state_cosmos3_adapter import (
    Cosmos3StateAdapter,
)
from vllm_omni.experimental.world_models.session_state import (
    SessionStateManager,
    resolve_session_state_config,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

from .action import (
    ACTION_MODE_FORWARD_DYNAMICS,
    ACTION_MODE_INVERSE_DYNAMICS,
    ACTION_MODE_POLICY,
    action_start_frame_offset,
    build_action_condition_mask,
    build_vision_condition_mask,
    find_closest_target_size,
    load_action_tensor,
    normalize_action_mode,
    pad_action_to_dim,
    resolve_domain_id,
    vision_condition_indexes,
)
from .transfer import (
    Cosmos3TransferConfig,
    has_transfer_hints,
    load_or_compute_control_frames,
    media_hw,
    media_to_uint8_cthw,
    normalized_video_to_uint8_cthw,
    pad_temporal_frames,
    resize_center_crop_uint8_cthw,
    resolve_transfer_config,
    transfer_max_frames_from_extra_args,
    uint8_cthw_to_normalized_5d,
)
from .transformer_cosmos3 import Cosmos3VFMTransformer, _tf_config_get, resolve_sound_gen
from .transformer_cosmos3_edge import COSMOS3_EDGE_BACKBONE_TYPE, Cosmos3EdgeVFMTransformer
from .utils import (
    COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION,
    COSMOS3_VAE_TEMPORAL_COMPRESSION,
    ROBOLAB_CONCAT_VIEW_DESCRIPTION,
    ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE,
    ROBOLAB_DEFAULT_ACTION_SPACE,
    ROBOLAB_DEFAULT_CONDITIONING_FPS,
    ROBOLAB_DEFAULT_DOMAIN_NAME,
    ROBOLAB_DEFAULT_FLOW_SHIFT,
    ROBOLAB_DEFAULT_GUIDANCE_SCALE,
    ROBOLAB_DEFAULT_IMAGE_HEIGHT,
    ROBOLAB_DEFAULT_IMAGE_WIDTH,
    ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS,
    ROBOLAB_DEFAULT_RAW_ACTION_DIM,
    ROBOLAB_DEFAULT_RESOLUTION,
    ROBOLAB_MIDTRAIN_RAW_ACTION_DIM,
    VIDEO_RES_SIZE_INFO,
    RoboLabActionPostprocessInputs,
    RoboLabPolicyInputs,
    build_abs_pose_from_components,
    build_robolab_unipc_scheduler,
    condition_pixel_frame_count,
    convert_midtrain_rotation,
    ensure_2d_float_array,
    ensure_gripper_array,
    extract_robolab_image,
    extract_robolab_prompt_image,
    get_robolab_domain_id,
    lazy_action_transform_pipeline,
    make_robolab_action_postprocess_inputs,
    next_robolab_seed,
    normalize_condition_frame_indexes_vision,
    normalize_condition_video_keep,
    normalize_robolab_action_space,
    pose_abs_to_rel,
    postprocess_robolab_action,
    resize_rgb_uint8,
)

logger = init_logger(__name__)


COSMOS3_DEFAULT_CONDITION_PIXEL_FRAMES = (
    max(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION) * COSMOS3_VAE_TEMPORAL_COMPRESSION + 1
)
COSMOS3_V2V_DEFAULT_FLOW_SHIFT = 10.0
COSMOS3_DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS."
COSMOS3_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution."
COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolution."
COSMOS3_INVERSE_DURATION_TEMPLATE = "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS."
COSMOS3_INVERSE_RESOLUTION_TEMPLATE = "This video is not of {height}x{width} resolution."
COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE = "This image is not of {height}x{width} resolution."
# NOTE: Intentional typo in "give" instead of "given" to match training setup.
COSMOS3_SYSTEM_PROMPT = "You are a helpful assistant who will generate videos from a give prompt."
COSMOS3_T2I_SYSTEM_PROMPT = "You are a helpful assistant who will generate images from a give prompt."
COSMOS3_TRANSFER_SYSTEM_PROMPT = (
    "You are a helpful assistant that generates images or videos following the user's instructions"
    " and control signals (edge maps, blur, depth, or segmentation)."
)
COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE = (
    "Follow the {hint_names} control video precisely: shape, contour, silhouette, position, and motion of every "
    "visible structure must align with the {hint_names} signal at every frame."
)

COSMOS3_T2V_DEFAULT_HEIGHT = 720
COSMOS3_T2V_DEFAULT_WIDTH = 1280
COSMOS3_T2V_DEFAULT_NUM_FRAMES = 189
COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS = 35
COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE = 6.0
COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT = 10.0

COSMOS3_T2I_DEFAULT_HEIGHT = 1024
COSMOS3_T2I_DEFAULT_WIDTH = 1024
COSMOS3_T2I_DEFAULT_NUM_INFERENCE_STEPS = 50
COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE = 7.0
COSMOS3_T2I_DEFAULT_FLOW_SHIFT = 3.0
COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0)

COSMOS3_EDGE_T2V_DEFAULT_HEIGHT = 480
COSMOS3_EDGE_T2V_DEFAULT_WIDTH = 832
COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE = 5.0
COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT = 3.0
COSMOS3_EDGE_T2I_DEFAULT_HEIGHT = 640
COSMOS3_EDGE_T2I_DEFAULT_WIDTH = 640

COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS = "FlowMatchEulerDiscreteScheduler"

# Truncation cap on the prompt token count (shared by T2I and T2V).  Prompts
# are tokenized to their natural length (no padding); this only bounds the
# UND pathway / GEN cross-attention cost for pathologically long prompts.
COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096
COSMOS3_TRANSFER_REQUESTED_SIZE_KEY = "_cosmos3_transfer_requested_size"


def _ceil_video_num_frames(num_frames: int, temporal_compression_factor: int) -> int:
    """Round a video length up to the causal VAE's ``factor * k + 1`` grid."""
    if temporal_compression_factor <= 0:
        raise ValueError(f"Cosmos3 temporal_compression_factor must be positive, got {temporal_compression_factor}.")
    if num_frames <= 1:
        return num_frames
    latent_frames = math.ceil((num_frames - 1) / temporal_compression_factor)
    return latent_frames * temporal_compression_factor + 1


def _format_json_object_prompt(
    prompt: str,
    *,
    num_frames: int,
    frame_rate: float,
    height: int,
    width: int,
    aspect_ratio: str | None,
) -> str | None:
    """Match Imaginaire4 metadata injection for JSON-object prompts."""
    try:
        prompt_obj = json.loads(prompt)
    except (json.JSONDecodeError, TypeError, ValueError):
        return None
    if not isinstance(prompt_obj, dict):
        return None

    metadata: dict[str, Any] = {}
    if num_frames > 1:
        duration_seconds = int(num_frames / frame_rate) if frame_rate > 0 else 0
        metadata.update({"duration": f"{duration_seconds}s", "fps": float(frame_rate)})
    else:
        prompt_obj.pop("duration", None)
        prompt_obj.pop("fps", None)
    metadata["resolution"] = {"H": int(height), "W": int(width)}
    if aspect_ratio is not None:
        metadata["aspect_ratio"] = aspect_ratio

    prompt_obj.update(metadata)
    return json.dumps(prompt_obj)


def _json_object_aspect_ratio(prompt: str) -> Any | None:
    """Return an aspect-ratio field from a JSON-object prompt, if present."""
    try:
        prompt_obj = json.loads(prompt)
    except (json.JSONDecodeError, TypeError, ValueError):
        return None
    if not isinstance(prompt_obj, dict):
        return None
    return prompt_obj.get("aspect_ratio")


def _normalize_aspect_ratio(value: Any) -> str | None:
    """Normalize supported ratio spellings for comparison and prompt metadata."""
    if value is None:
        return None
    parts = str(value).strip().replace(":", ",").split(",")
    if len(parts) != 2:
        return None
    try:
        width, height = (int(part.strip()) for part in parts)
    except ValueError:
        return None
    if width <= 0 or height <= 0:
        return None
    divisor = math.gcd(width, height)
    return f"{width // divisor},{height // divisor}"


def _aspect_ratio_for_dimensions(width: int, height: int) -> str:
    """Return the nearest supported Cosmos aspect-ratio label."""
    if width <= 0 or height <= 0:
        raise ValueError(f"Cosmos3 canvas dimensions must be positive, got {width}x{height}.")
    canvas_ratio = width / height

    def ratio_value(label: str) -> float:
        ratio_width, ratio_height = (int(part) for part in label.split(","))
        return ratio_width / ratio_height

    return min(
        (label for sizes in VIDEO_RES_SIZE_INFO.values() for label in sizes),
        key=lambda label: abs(canvas_ratio - ratio_value(label)),
    )


def resolve_cosmos3_transformer_cls(model_config: Any) -> type[Cosmos3VFMTransformer]:
    """Select the Cosmos3 transformer implementation from transformer/config.json."""
    backbone_type = _tf_config_get(model_config, "backbone_type", None)
    if backbone_type is None:
        return Cosmos3VFMTransformer
    if backbone_type == COSMOS3_EDGE_BACKBONE_TYPE:
        return Cosmos3EdgeVFMTransformer
    raise ValueError(f"Unsupported Cosmos3 transformer backbone_type={backbone_type!r}.")


# ---------------------------------------------------------------------------
# Post-process function (registered in registry.py)
# ---------------------------------------------------------------------------
def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
    """Build the request preprocessor for Cosmos3 image/video inputs.

    For plain T2V (no image or video in ``multi_modal_data``), the request is
    returned unchanged after the optional guardrail check. For I2V, the
    conditioning image is loaded, aspect-resized, center-cropped, and stored as
    ``additional_information.preprocessed_image``. For V2V, source frames are
    cropped to the target size and stored as
    ``additional_information.preprocessed_video``.

    Action modes reuse image/video preprocessing but use action-specific resize
    and padding rules. Transfer requests store
    ``additional_information.preprocessed_transfer_video`` for optional input
    video conditioning. Cosmos3 sound generation is not driven by
    ``multi_modal_data["audio"]``; it is enabled later from sampling params.
    """
    from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled

    video_processor = VideoProcessor(vae_scale_factor=16)
    is_edge_model = (
        _tf_config_get(getattr(od_config, "tf_model_config", None), "backbone_type", None) == COSMOS3_EDGE_BACKBONE_TYPE
    )
    # Eager-load guardrail models at pipeline build time when the server-level
    # gate is on. Per-request overrides only decide whether the loaded models
    # are *invoked* — they cannot turn checks on without a server-side preload.
    if is_guardrails_enabled(od_config):
        ensure_initialized(od_config)

    def _extra_args(request: OmniDiffusionRequest) -> dict[str, Any]:
        extra = getattr(getattr(request, "sampling_params", None), "extra_args", None)
        return extra if isinstance(extra, dict) else {}

    def _request_action_mode(request: OmniDiffusionRequest) -> str | None:
        return normalize_action_mode(_extra_args(request).get("action_mode"))

    def _set_transfer_size_from_image(request: OmniDiffusionRequest, image: PIL.Image.Image) -> tuple[int, int]:
        extra = _extra_args(request)
        resolution = extra.get("resolution", extra.get("image_size", 720))
        target_w, target_h = find_closest_target_size(image.height, image.width, resolution)
        request.sampling_params.height = target_h
        request.sampling_params.width = target_w
        return int(target_h), int(target_w)

    def _set_action_size_from_image(request: OmniDiffusionRequest, image: PIL.Image.Image) -> tuple[int, int]:
        sp = request.sampling_params
        if sp.height is not None and sp.width is not None:
            return int(sp.height), int(sp.width)

        extra = _extra_args(request)
        resolution = extra.get("resolution", extra.get("image_size", 480))
        target_w, target_h = find_closest_target_size(image.height, image.width, resolution)
        if sp.height is None:
            sp.height = target_h
        if sp.width is None:
            sp.width = target_w
        return int(sp.height), int(sp.width)

    def _pil_to_rgb(value: Any) -> PIL.Image.Image:
        if isinstance(value, str):
            return PIL.Image.open(value).convert("RGB")
        if isinstance(value, PIL.Image.Image):
            return value.convert("RGB")
        if isinstance(value, np.ndarray):
            array = value
            if array.ndim == 3 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4):
                array = np.transpose(array, (1, 2, 0))
            if np.issubdtype(array.dtype, np.floating):
                if array.min() < 0.0 or array.max() > 1.0:
                    array = np.clip(array, -1.0, 1.0) * 0.5 + 0.5
                array = (np.clip(array, 0.0, 1.0) * 255.0).round().astype(np.uint8)
            return PIL.Image.fromarray(array).convert("RGB")
        if isinstance(value, torch.Tensor):
            tensor = value.detach().cpu()
            if tensor.ndim == 3 and tensor.shape[0] in (3, 4):
                tensor = tensor.permute(1, 2, 0)
            if tensor.is_floating_point():
                if tensor.min().item() < 0.0 or tensor.max().item() > 1.0:
                    tensor = tensor.clamp(-1.0, 1.0) * 0.5 + 0.5
                tensor = (tensor.clamp(0.0, 1.0) * 255.0).round().to(torch.uint8)
            return PIL.Image.fromarray(tensor.numpy()).convert("RGB")
        raise TypeError(
            f"Cosmos3 preprocessing expected PIL image, numpy array, torch tensor, or path, got {type(value)!r}."
        )

    def _resize_and_pad_action_image(image: PIL.Image.Image, target_h: int, target_w: int) -> PIL.Image.Image:
        scale = min(target_w / image.width, target_h / image.height, 1.0)
        resize_w = max(1, int(scale * image.width + 0.5))
        resize_h = max(1, int(scale * image.height + 0.5))
        if (resize_w, resize_h) != image.size:
            image = image.resize((resize_w, resize_h), PIL.Image.Resampling.BICUBIC)

        array = np.asarray(image)
        pad_h = target_h - resize_h
        pad_w = target_w - resize_w
        if pad_h < 0 or pad_w < 0:
            raise ValueError(
                f"Cosmos3 action image resize exceeded target size: resized={(resize_h, resize_w)}, "
                f"target={(target_h, target_w)}."
            )
        if pad_h == 0 and pad_w == 0:
            return image
        pad_mode = "reflect" if pad_h < resize_h and pad_w < resize_w else "edge"
        padded = np.pad(array, ((0, pad_h), (0, pad_w), (0, 0)), mode=pad_mode)
        return PIL.Image.fromarray(padded)

    def _preprocess_action_image(image: PIL.Image.Image, target_h: int, target_w: int) -> torch.Tensor:
        image = _resize_and_pad_action_image(image, target_h, target_w)
        return video_processor.preprocess(image, height=target_h, width=target_w)

    def _preprocess_action_video(frames: list[Any], target_h: int, target_w: int) -> torch.Tensor:
        if not frames:
            raise ValueError("Cosmos3 action video input must contain at least one frame.")
        processed = [_preprocess_action_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in frames]
        return torch.stack(processed, dim=1).unsqueeze(0).contiguous()

    def _preprocess_condition_image(image: PIL.Image.Image, target_h: int, target_w: int) -> torch.Tensor:
        scale = max(target_w / image.width, target_h / image.height)
        resize_w = int(np.ceil(scale * image.width))
        resize_h = int(np.ceil(scale * image.height))
        image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
        left = (resize_w - target_w) // 2
        top = (resize_h - target_h) // 2
        image = image.crop((left, top, left + target_w, top + target_h))
        return video_processor.preprocess(image, height=target_h, width=target_w)

    def _video_payload_value(video: Any, key: str) -> Any:
        if isinstance(video, Mapping):
            return video.get(key)
        return getattr(video, key, None)

    def _video_payload_fps(video: Any) -> float | None:
        for key in ("fps", "frame_rate", "source_fps", "input_fps", "avg_fps", "average_fps"):
            fps = positive_float(_video_payload_value(video, key))
            if fps is not None:
                return fps
        for key in ("metadata", "info"):
            metadata = _video_payload_value(video, key)
            if metadata is None or metadata is video:
                continue
            fps = _video_payload_fps(metadata)
            if fps is not None:
                return fps
        if isinstance(video, Mapping):
            for key in ("frames", "data", "video"):
                nested = video.get(key)
                if nested is None or nested is video:
                    continue
                fps = _video_payload_fps(nested)
                if fps is not None:
                    return fps
        return None

    def _unwrap_video_payload(video: Any) -> Any:
        if isinstance(video, Mapping):
            for key in ("frames", "data", "video"):
                nested = video.get(key)
                if nested is not None:
                    return nested
        return video

    def _decode_video_file_to_frames(
        path: str | os.PathLike,
        *,
        max_frames: int | None,
        keep: str,
    ) -> list[PIL.Image.Image]:
        rgb_frames = decode_path_video_frames(path, max_frames=max_frames, keep=keep)
        return [PIL.Image.fromarray(frame).convert("RGB") for frame in rgb_frames]

    def _expand_video_payload_item(
        item: Any,
        *,
        max_frames: int | None,
        keep: str,
    ) -> list[Any]:
        if is_video_file_path(item):
            return _decode_video_file_to_frames(item, max_frames=max_frames, keep=keep)
        return [item]

    def _video_payload_to_frames(
        video: Any,
        *,
        max_frames: int | None,
        keep: str,
    ) -> list[Any]:
        video = _unwrap_video_payload(video)
        if is_video_file_path(video):
            return _decode_video_file_to_frames(video, max_frames=max_frames, keep=keep)
        if isinstance(video, list):
            frames: list[Any] = []
            for item in video:
                frames.extend(_expand_video_payload_item(item, max_frames=max_frames, keep=keep))
            return frames
        if isinstance(video, torch.Tensor):
            tensor = video.detach().cpu()
            if tensor.ndim == 5:
                if tensor.shape[0] != 1:
                    raise TypeError("Cosmos3 video preprocessing supports only batch size 1.")
                tensor = tensor[0]
            if tensor.ndim == 4 and tensor.shape[0] in (3, 4) and tensor.shape[-1] not in (3, 4):
                return [tensor[:, i] for i in range(tensor.shape[1])]
            if tensor.ndim == 4 and tensor.shape[-1] in (3, 4):
                return [tensor[i] for i in range(tensor.shape[0])]
        if isinstance(video, np.ndarray):
            array = video
            if array.ndim == 5:
                if array.shape[0] != 1:
                    raise TypeError("Cosmos3 video preprocessing supports only batch size 1.")
                array = array[0]
            if array.ndim == 4 and array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4):
                return [array[:, i] for i in range(array.shape[1])]
            if array.ndim == 4 and array.shape[-1] in (3, 4):
                return [array[i] for i in range(array.shape[0])]
        raise TypeError("Cosmos3 video input must be a non-empty list of frames or a single video tensor/array.")

    def _select_video_frames(frames: list[Any], max_frames: int, keep: str) -> list[Any]:
        if not frames:
            raise ValueError("Cosmos3 video input must contain at least one frame.")
        if keep == "last":
            return frames[-max_frames:]
        return frames[:max_frames]

    def _preprocess_condition_video(
        frames: list[Any],
        target_h: int,
        target_w: int,
        max_frames: int,
        keep: str,
    ) -> torch.Tensor:
        selected = _select_video_frames(frames, max_frames, keep)
        processed = [
            _preprocess_condition_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in selected
        ]
        return torch.stack(processed, dim=1).unsqueeze(0).contiguous()

    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
        action_mode = _request_action_mode(request)
        prompt = request.prompt
        if is_guardrails_enabled(od_config, request.sampling_params):
            text = prompt if isinstance(prompt, str) else prompt.get("prompt", "")
            check_text_safety(text)

        if isinstance(prompt, str):
            return request
        multi_modal_data = prompt.get("multi_modal_data", {}) or {}
        raw_image = multi_modal_data.get("image")
        raw_video = multi_modal_data.get("video")
        if raw_image is None and raw_video is None:
            return request
        if raw_image is not None and raw_video is not None and action_mode is None:
            raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")

        if "additional_information" not in prompt:
            prompt["additional_information"] = {}

        extra = _extra_args(request)
        transfer_requested = action_mode is None and has_transfer_hints(extra)
        condition_frame_indexes_vision = normalize_condition_frame_indexes_vision(
            extra.get("condition_frame_indexes_vision", prompt.get("condition_frame_indexes_vision"))
        )
        condition_video_keep = normalize_condition_video_keep(
            extra.get("condition_video_keep", prompt.get("condition_video_keep"))
        )

        raw_video_frames: list[Any] | None = None
        transfer_input_fps: float | None = None
        if raw_video is not None:
            transfer_input_fps = _video_payload_fps(raw_video)
            decode_extra = dict(extra)
            decode_extra["condition_frame_indexes_vision"] = list(condition_frame_indexes_vision)
            decode_extra["condition_video_keep"] = condition_video_keep
            decode_spec = Cosmos3OmniDiffusersPipeline.reference_video_decode_spec(
                num_frames=getattr(request.sampling_params, "num_frames", None),
                extra_args=decode_extra,
            )
            raw_video_frames = _video_payload_to_frames(
                raw_video,
                max_frames=decode_spec.max_frames,
                keep=decode_spec.keep,
            )
            if not raw_video_frames:
                raise TypeError("Cosmos3 video input must be a non-empty list of PIL images or image paths.")

        if raw_image is None:
            assert raw_video_frames is not None  # raw_image and raw_video can't both be None here
            image = _pil_to_rgb(raw_video_frames[0])
        else:
            image = _pil_to_rgb(raw_image)

        # Resolve missing H/W.
        if transfer_requested:
            requested_width = request.sampling_params.width
            requested_height = request.sampling_params.height
            if (
                requested_width is not None
                and requested_height is not None
                and COSMOS3_TRANSFER_REQUESTED_SIZE_KEY not in extra
            ):
                extra[COSMOS3_TRANSFER_REQUESTED_SIZE_KEY] = {
                    "width": int(requested_width),
                    "height": int(requested_height),
                }
            _set_transfer_size_from_image(request, image)
        elif request.sampling_params.height is None or request.sampling_params.width is None:
            if action_mode is not None:
                _set_action_size_from_image(request, image)
            elif is_edge_model:
                if request.sampling_params.height is None:
                    request.sampling_params.height = COSMOS3_EDGE_T2V_DEFAULT_HEIGHT
                if request.sampling_params.width is None:
                    request.sampling_params.width = COSMOS3_EDGE_T2V_DEFAULT_WIDTH
            else:
                max_area = 720 * 1280
                aspect_ratio = image.height / image.width
                mod_value = 16
                height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
                width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
                if request.sampling_params.height is None:
                    request.sampling_params.height = height
                if request.sampling_params.width is None:
                    request.sampling_params.width = width

        target_w = request.sampling_params.width
        target_h = request.sampling_params.height
        if action_mode is not None:
            prompt["additional_information"]["preprocessed_image"] = _preprocess_action_image(
                image,
                int(target_h),
                int(target_w),
            )
        elif raw_video is None:
            prompt["additional_information"]["preprocessed_image"] = _preprocess_condition_image(
                image,
                int(target_h),
                int(target_w),
            )
        else:
            assert raw_video_frames is not None
            if transfer_requested:
                if transfer_input_fps is not None:
                    prompt["additional_information"]["transfer_input_fps"] = transfer_input_fps
                transfer_frames = media_to_uint8_cthw(
                    raw_video_frames,
                    height=int(target_h),
                    width=int(target_w),
                    max_frames=transfer_max_frames_from_extra_args(extra),
                )
                prompt["additional_information"]["preprocessed_transfer_video"] = uint8_cthw_to_normalized_5d(
                    transfer_frames,
                    dtype=torch.float32,
                )
            else:
                max_frames = condition_pixel_frame_count(condition_frame_indexes_vision)
                prompt["additional_information"]["preprocessed_video"] = _preprocess_condition_video(
                    raw_video_frames,
                    int(target_h),
                    int(target_w),
                    max_frames,
                    condition_video_keep,
                )
                prompt["additional_information"]["condition_frame_indexes_vision"] = list(
                    condition_frame_indexes_vision
                )
        if action_mode is not None and raw_video_frames is not None:
            prompt["additional_information"]["preprocessed_video"] = _preprocess_action_video(
                raw_video_frames,
                int(target_h),
                int(target_w),
            )
        request.prompt = prompt

        return request

    return pre_process_func


def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig):
    """Build the postprocessor for Cosmos3 image, video, and video+audio output.

    The pipeline returns image payloads as ``{"image": tensor}`` and video
    payloads as ``{"video": tensor}``. Sound-enabled video returns the same
    video payload plus ``audio`` and ``audio_sample_rate``. Image output with
    audio is rejected because Cosmos3 sound generation is video-only.
    """
    from .guardrails import check_video_safety, is_guardrails_enabled

    video_processor = VideoProcessor(vae_scale_factor=16)

    def _sampling_param(sampling_params, key: str, default=None):
        extra = getattr(sampling_params, "extra_args", None)
        if isinstance(extra, dict) and extra.get(key) is not None:
            return extra[key]
        value = getattr(sampling_params, key, None)
        return default if value is None else value

    def _resolve_output_fps(sampling_params):
        fps = (
            _sampling_param(sampling_params, "resolved_frame_rate")
            or _sampling_param(sampling_params, "frame_rate")
            or _sampling_param(sampling_params, "fps")
            or 24.0
        )
        try:
            fps_value = float(fps)
        except (TypeError, ValueError):
            fps_value = 24.0
        if fps_value <= 0:
            fps_value = 24.0
        return int(fps_value) if fps_value.is_integer() else fps_value

    def post_process_func(
        output: torch.Tensor | dict[str, torch.Tensor] | tuple,
        output_type: str = "np",
        sampling_params=None,
    ):
        if output_type == "latent":
            return output

        def _postprocess_action(action: Any, metadata: dict[str, Any]) -> Any:
            internal_metadata = metadata.get("internal")
            inputs = (
                internal_metadata.get("robolab_action_postprocess") if isinstance(internal_metadata, dict) else None
            )
            if isinstance(inputs, RoboLabActionPostprocessInputs):
                return postprocess_robolab_action(action, inputs)
            return action

        pending_action = None
        pending_action_metadata: dict[str, Any] = {}
        envelope_public_metadata: dict[str, Any] = {}
        if isinstance(output, dict) and isinstance(output.get("payload"), dict):
            envelope_payload = dict(output.get("payload") or {})
            metadata = output.get("metadata") or {}
            envelope_metadata = metadata if isinstance(metadata, dict) else {}
            envelope_public_metadata = {key: value for key, value in envelope_metadata.items() if key != "internal"}
            action = envelope_payload.pop("actions", None)
            if action is not None:
                pending_action = _postprocess_action(action, envelope_metadata)
                pending_action_metadata = envelope_public_metadata
                if not envelope_payload:
                    return {
                        "payload": {
                            "video": [],
                            "actions": pending_action,
                        },
                        "metadata": pending_action_metadata,
                    }
            output = envelope_payload

        audio = None
        audio_sample_rate = None
        if isinstance(output, dict):
            if "image" in output and "video" in output:
                raise ValueError("Cosmos3 output cannot contain both image and video payloads.")
            if "image" in output:
                video = output["image"]
            elif "video" in output:
                video = output["video"]
            else:
                raise ValueError("Cosmos3 postprocess expected an 'image' or 'video' output payload.")
            audio = output.get("audio")
            audio_sample_rate = output.get("audio_sample_rate")
        elif isinstance(output, tuple):
            if len(output) == 3:
                video, audio, audio_sample_rate = output
            elif len(output) == 2:
                video, audio = output
            else:
                raise ValueError(
                    "Cosmos3 postprocess expects output tensor, output dict, or (video, audio[, sample_rate]) tuple."
                )
        else:
            video = output

        if isinstance(output, dict) and "image" in output:
            if audio is not None:
                raise ValueError("Cosmos3 text-to-image postprocess does not support audio output.")
            if video.ndim != 5 or video.shape[2] != 1:
                raise ValueError(
                    "Cosmos3 text-to-image postprocess expects decoded output "
                    f"with shape [B, C, 1, H, W], got {tuple(video.shape)}."
                )
            image = video.squeeze(2)  # [B, 3, H, W]
            if is_guardrails_enabled(od_config, sampling_params):
                # check_video_safety expects a 5D tensor; re-add T axis.
                checked = check_video_safety(image.unsqueeze(2))
                image = checked.squeeze(2)
            processed_image = video_processor.postprocess(image, output_type="pil")
            if envelope_public_metadata:
                return {
                    "payload": {"image": processed_image},
                    "metadata": envelope_public_metadata,
                }
            return processed_image
        guardrails_enabled = is_guardrails_enabled(od_config, sampling_params)
        if guardrails_enabled:
            video = check_video_safety(video)
        processed_video = video_processor.postprocess_video(video, output_type=output_type)
        if audio is None:
            if pending_action is not None:
                return {
                    "payload": {
                        "video": processed_video,
                        "actions": pending_action,
                    },
                    "metadata": pending_action_metadata,
                }
            if envelope_public_metadata:
                return {
                    "payload": {"video": processed_video},
                    "metadata": envelope_public_metadata,
                }
            return processed_video
        if isinstance(audio, torch.Tensor):
            audio = audio.detach().cpu()
        result = {
            "video": processed_video,
            "audio": audio,
            "fps": _resolve_output_fps(sampling_params),
        }
        if audio_sample_rate is not None:
            result["audio_sample_rate"] = int(audio_sample_rate)
        if pending_action is not None:
            return {
                "payload": {
                    "video": result["video"],
                    "audio": result["audio"],
                    "actions": pending_action,
                },
                "metadata": {
                    **pending_action_metadata,
                    "video": {"fps": result["fps"]},
                    "audio": {"sample_rate": result.get("audio_sample_rate")},
                },
            }
        if envelope_public_metadata:
            return {
                "payload": {
                    "video": result["video"],
                    "audio": result["audio"],
                },
                "metadata": {
                    **envelope_public_metadata,
                    "video": {"fps": result["fps"], **envelope_public_metadata.get("video", {})}
                    if isinstance(envelope_public_metadata.get("video"), dict)
                    else {"fps": result["fps"]},
                    "audio": {"sample_rate": result.get("audio_sample_rate")},
                },
            }
        return result

    return post_process_func


def get_cosmos3_ir_op_priority_func(od_config: OmniDiffusionConfig):
    del od_config

    def ir_op_priority_func(ir_op_priority, vllm_config=None):
        del vllm_config
        from vllm.config.kernel import IrOpPriorityConfig

        priority_kwargs = {field.name: list(getattr(ir_op_priority, field.name)) for field in fields(ir_op_priority)}
        priority_kwargs["rms_norm"] = ["native"]
        priority_kwargs["fused_add_rms_norm"] = ["native"]
        return IrOpPriorityConfig(**priority_kwargs)

    return ir_op_priority_func


# ---------------------------------------------------------------------------
# Pipeline
# ---------------------------------------------------------------------------
class Cosmos3OmniDiffusersPipeline(
    nn.Module,
    CFGParallelMixin,
    SupportImageInput,
    SupportsComponentDiscovery,
    ProgressBarMixin,
    DiffusionPipelineProfilerMixin,
):
    """Cosmos3 text/image/video/sound/action pipeline.

    Architecture: Mixture-of-Transformers with Qwen3-VL backbone.
    - Understanding pathway: causal self-attention on text (runs once, K/V cached)
    - Generation pathway: cross-attention on visual latents and optional
      transfer-control, action, and sound latents (runs each step)

    Supports T2V, I2V, V2V, T2I, transfer, sound-enabled video, and action
    generation from the same class. Mode is selected at runtime:

    * **T2I** when ``prompt["modalities"]`` contains ``"image"``.  Latent
      T-dim is forced to 1, T2I-specific scheduler defaults are applied (50 steps,
      flow_shift=3.0, guidance_interval=[400, 1000]), the duration
      template is suppressed, and post-process emits PIL images.
    * **I2V** when the request supplies a preprocessed image via
      ``multi_modal_data['image']`` (handled by
      :func:`get_cosmos3_pre_process_func`) and the requested output modality
      is not image.
      Frame 0 of the initial latent is set to the VAE-encoded conditioning
      image, frame-0 noise predictions are masked to zero, and the clean
      image latent is re-injected at frame 0 after each scheduler step.
    * **V2V** when the request supplies a preprocessed video via
      ``multi_modal_data['video']`` without an action mode. Explicit latent
      frame indexes are kept clean with ``noisy_frame_mask`` and re-injected
      after each scheduler step.
    * **Transfer** when ``edge``, ``blur``, ``depth``, ``seg``, or ``wsm`` hints
      are supplied. Transfer is video-output only and cannot be combined with
      sound or action generation.
    * **Sound-enabled video** when ``generate_sound`` or ``sound_gen`` is true.
      Sound is generated from sound latents, not from ``multi_modal_data['audio']``;
      T2I, transfer, and action+sound are rejected.
    * **Action generation** when ``action_mode`` is provided. ``policy`` and
      ``forward_dynamics`` require an image or video input; ``inverse_dynamics``
      requires video input. Action predictions are returned in the diffusion
      output payload/metadata envelope.
      RoboLab/OpenPI observations in ``extra_args['robot_obs']`` or
      ``extra_args['observation']`` return an action-only envelope.
    * **T2V** otherwise (default video generation).
    """

    support_image_input: ClassVar[bool] = True
    color_format: ClassVar[str] = "RGB"
    _dit_modules: ClassVar[list[str]] = ["transformer.language_model", "transformer"]
    _encoder_modules: ClassVar[list[str]] = []
    _vae_modules: ClassVar[list[str]] = ["vae"]
    _resident_modules: ClassVar[list[str]] = []
    sampling_dtype: ClassVar[torch.dtype] = torch.float32

    @classmethod
    def reference_video_decode_spec(
        cls,
        *,
        num_frames: int | None = None,
        extra_args: dict[str, Any] | None = None,
    ) -> ReferenceVideoDecodeSpec:
        extra_args = extra_args if isinstance(extra_args, dict) else {}
        if has_transfer_hints(extra_args):
            max_frames = transfer_max_frames_from_extra_args(extra_args)
            if num_frames is not None:
                max_frames = min(max_frames, int(num_frames))
            return ReferenceVideoDecodeSpec(max_frames=max_frames, keep="first")

        action_mode = normalize_action_mode(extra_args.get("action_mode"))
        if action_mode is not None:
            if num_frames is not None:
                return ReferenceVideoDecodeSpec(max_frames=int(num_frames), keep="first")
            action_chunk_size = extra_args.get("action_chunk_size")
            if action_chunk_size is not None:
                try:
                    max_frames = int(action_chunk_size) + 1
                except (TypeError, ValueError):
                    max_frames = None
                if max_frames is not None and max_frames > 0:
                    return ReferenceVideoDecodeSpec(max_frames=max_frames, keep="first")
            return ReferenceVideoDecodeSpec(max_frames=None, keep="first")

        condition_indexes = normalize_condition_frame_indexes_vision(extra_args.get("condition_frame_indexes_vision"))
        max_frames = condition_pixel_frame_count(condition_indexes)
        if num_frames is not None:
            max_frames = min(max_frames, int(num_frames))
        keep = normalize_condition_video_keep(extra_args.get("condition_video_keep"))
        return ReferenceVideoDecodeSpec(max_frames=max_frames, keep=keep)

    def __init__(
        self,
        *,
        od_config: OmniDiffusionConfig,
        prefix: str = "",
    ) -> None:
        super().__init__()
        self.od_config = od_config
        resolved_offload = resolve_offload(od_config)
        if resolved_offload.public is not None and resolved_offload.strategy is OffloadStrategy.MODEL_LEVEL:
            raise ValueError(
                "Cosmos3 model offload uses reasoner/generator topology and does "
                "not support the dit/text_encoder component selector"
            )
        self.device = get_local_device()
        self.dtype = od_config.dtype

        model_path = od_config.model
        local_files_only = os.path.exists(model_path)

        # Validate guidance parallelism from checkpoint metadata before loading
        # components. Distilled checkpoints have only a conditional branch.
        scheduler_config = FlowUniPCMultistepScheduler.load_config(
            model_path,
            subfolder="scheduler",
            local_files_only=local_files_only,
        )
        self.is_distilled_model = scheduler_config.get("_class_name") == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS
        if self.is_distilled_model and od_config.parallel_config.cfg_parallel_size > 1:
            raise ValueError(
                "Distilled Cosmos3 checkpoints run without classifier-free guidance. "
                "Set --cfg-parallel-size 1 and use --ulysses-degree for multi-GPU inference."
            )

        # --- Tokenizer ---
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            subfolder="text_tokenizer",
            local_files_only=local_files_only,
        )

        # --- VAE ---
        self.vae = DistributedAutoencoderKLWan.from_pretrained(
            model_path,
            subfolder="vae",
            torch_dtype=self.dtype,
            local_files_only=local_files_only,
        ).to(self.device)

        if not hasattr(self.vae.config, "scale_factor_temporal"):
            raise ValueError(
                "Cosmos3 Diffusers VAE config must define scale_factor_temporal "
                "so transformer mRoPE temporal positions can be computed correctly."
            )
        self.vae_scale_factor_temporal = int(self.vae.config.scale_factor_temporal)
        self.vae_scale_factor_spatial = getattr(self.vae.config, "scale_factor_spatial", 16)

        sound_gen = resolve_sound_gen(od_config)
        sound_dim = None
        sound_latent_fps = None
        self._sound_tokenizer = None
        if sound_gen:
            self._sound_tokenizer = self._get_sound_tokenizer()
            sound_dim = self._sound_tokenizer.latent_ch
            sound_latent_fps = self._sound_tokenizer.latent_fps

        # --- Transformer (weights loaded later via weights_sources) ---
        transformer_cls = resolve_cosmos3_transformer_cls(od_config.tf_model_config)
        self.transformer = transformer_cls(
            od_config=od_config,
            temporal_compression_factor=self.vae_scale_factor_temporal,
            sound_gen=sound_gen,
            sound_dim=sound_dim,
            sound_latent_fps=sound_latent_fps,
        )
        self.is_edge_model = transformer_cls is Cosmos3EdgeVFMTransformer

        # Opt-in: route the per-CFG-branch UND text K/V through the shared
        # SessionStateManager (RFC #4480) instead of the transformer-instance
        # cache. This establishes request-keyed state isolation; it does not by
        # itself make the mutable pipeline safe for concurrent execution.
        # Default off -> the bespoke transformer-instance path below is unchanged.
        self._use_session_state, mm_max_sessions = resolve_session_state_config(
            enable=od_config.enable_session_state_manager,
        )
        self._memory_manager: SessionStateManager | None = (
            SessionStateManager(max_sessions=mm_max_sessions) if self._use_session_state else None
        )
        # Forward declaration for the future BDE paged path (Stage C/D). Nothing
        # writes this field in Phase 0; a non-None value is rejected by the guards
        # below. None means the dense session adapter or the bespoke path.
        self._bde_kv_state = None
        if getattr(self, "_use_session_state", False):
            logger.info("Cosmos3: session state manager enabled (max_sessions=%d)", mm_max_sessions)

        # --- Scheduler ---
        if self.is_distilled_model:
            fixed_step_config = scheduler_config.get("fixed_step_sampler_config")
            if not isinstance(fixed_step_config, dict) or fixed_step_config.get("sample_type") != "sde":
                raise ValueError("Cosmos3 distilled scheduler requires fixed_step_sampler_config.sample_type=sde.")
            t_list = fixed_step_config.get("t_list")
            if not isinstance(t_list, list) or not t_list:
                raise ValueError("Cosmos3 distilled scheduler requires a non-empty fixed_step_sampler_config.t_list.")
            self.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
                scheduler_config,
                stochastic_sampling=True,
            )
            self._scheduler_init_t_list = list(t_list)
        else:
            # Preserve compatible solver settings from the checkpoint, but keep
            # the base shift neutral. The concrete request shift is applied when
            # FlowUniPC builds its timesteps.
            self.scheduler = FlowUniPCMultistepScheduler.from_config(
                scheduler_config,
                shift=1.0,
                use_dynamic_shifting=False,
                prediction_type="flow_prediction",
            )
        self._cpu_scheduler_state()

        # --- Video processor for post-decode ---
        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)

        # --- Weight sources for DiffusersPipelineLoader ---
        self.weights_sources = [
            DiffusersPipelineLoader.ComponentSource(
                model_or_path=model_path,
                subfolder="transformer",
                revision=None,
                prefix="transformer.",
                fall_back_to_pt=True,
            ),
        ]

        # An engine-level override becomes the video default; otherwise use
        # the matching regular/Edge default. Per-request values still take
        # precedence.
        default_video_flow_shift = (
            COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT if self.is_edge_model else COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT
        )
        self._engine_init_flow_shift = float(
            od_config.flow_shift if od_config.flow_shift is not None else default_video_flow_shift
        )
        self._current_flow_shift = self._engine_init_flow_shift

        self._guidance_scale = None
        self._num_timesteps = None
        self._current_step_index = None
        self._current_sigma = None
        self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None
        self._robolab_transforms: dict[bool, Any] = {}

        # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on
        # this pipeline. Tells the sequential-CFG loop to keep paired
        # cond/uncond forwards so cache-dit's has_separate_cfg step accounting
        # stays in sync.
        self._cache_dit_requires_paired_cfg = False

        self.setup_diffusion_pipeline_profiler(
            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
        )

    def enable_omni_model_cpu_offload(
        self,
        *,
        device: torch.device,
        pin_memory: bool = True,
        use_hsdp: bool = False,
        offload_components: frozenset[str] | None = None,
    ) -> None:
        """Enable Cosmos3 component-level model offload.

        Cosmos3 has a nested reasoner/generator transformer instead of separate
        text-encoder and DiT pipeline components, so the transformer owns the
        mutual-exclusion swaps.  The VAE stays resident on GPU like the generic
        model-level offloader.
        """
        if offload_components is not None:
            raise ValueError(
                "Cosmos3 model offload uses reasoner/generator topology and does not support "
                "the dit/text_encoder offload_components selector"
            )
        self.vae.to(device, non_blocking=True)
        if isinstance(self._sound_tokenizer, nn.Module):
            self._sound_tokenizer.to(device)
        self.transformer.enable_model_cpu_offload(
            device=device,
            pin_memory=pin_memory,
            use_hsdp=use_hsdp,
        )

    def disable_omni_model_cpu_offload(self) -> None:
        self.transformer.disable_model_cpu_offload()

    # -- Weight loading --------------------------------------------------------

    @staticmethod
    def _remap_ckpt_key(key: str) -> str | None:
        """Remap a Diffusers transformer key to the model parameter namespace.

        Checkpoint keys arrive with a synthetic ``transformer.`` prefix from
        ``weights_sources``.  The source checkpoint itself uses the prefixless
        Diffusers transformer namespace: top-level projections plus Qwen3-VL
        backbone keys.  UND and GEN components share each layer in the source
        and are split into separate module lists here.  Some sources wrap the
        transformer namespace under ``model.``; that wrapper is structural and
        is stripped before applying the Cosmos3 leaf-name remap.

        Returns the remapped name under ``transformer.``, or None to skip.
        """
        k = key
        # Strip the weights_sources prefix
        if k.startswith("transformer."):
            k = k[len("transformer.") :]
        if k.startswith("model."):
            k = k[len("model.") :]

        # Top-level generation components.
        if k.startswith(
            (
                "proj_in.",
                "proj_out.",
                "time_embedder.",
                "audio_proj_in.",
                "audio_proj_out.",
                "action_proj_in.",
                "action_proj_out.",
            )
        ):
            return f"transformer.{k}"
        if k in ("audio_modality_embed", "audio_modality_embed.weight"):
            return "transformer.audio_modality_embed"
        if k in ("action_modality_embed", "action_modality_embed.weight"):
            return "transformer.action_modality_embed"
        if k.startswith("action_pos_embed."):
            return None

        # Skip lm_head
        if k.startswith("lm_head."):
            return None

        # embed_tokens / norm -> language_model.*
        if k.startswith("embed_tokens."):
            return f"transformer.language_model.{k}"
        if k.startswith("norm."):
            return f"transformer.language_model.{k}"

        # norm_moe_gen -> top level
        if k.startswith("norm_moe_gen."):
            return f"transformer.{k}"

        if not k.startswith("layers."):
            return None

        parts = k.split(".", 2)  # ['layers', '{i}', '{rest}']
        if len(parts) != 3:
            return None
        layer_idx = parts[1]
        rest = parts[2]

        und_lp = f"transformer.language_model.layers.{layer_idx}"
        gen_lp = f"transformer.gen_layers.{layer_idx}"

        _LAYER_MAP = {
            # UND attention
            "self_attn.to_q.": f"{und_lp}.self_attn.to_q.",
            "self_attn.to_k.": f"{und_lp}.self_attn.to_k.",
            "self_attn.to_v.": f"{und_lp}.self_attn.to_v.",
            "self_attn.to_out.": f"{und_lp}.self_attn.to_out.",
            "self_attn.norm_q.": f"{und_lp}.self_attn.norm_q.",
            "self_attn.norm_k.": f"{und_lp}.self_attn.norm_k.",
            "self_attn.k_norm_und_for_gen.": f"{und_lp}.self_attn.k_norm_und_for_gen.",
            # GEN attention
            "self_attn.add_q_proj.": f"{gen_lp}.cross_attention.to_q.",
            "self_attn.add_k_proj.": f"{gen_lp}.cross_attention.to_k.",
            "self_attn.add_v_proj.": f"{gen_lp}.cross_attention.to_v.",
            "self_attn.to_add_out.": f"{gen_lp}.cross_attention.to_out.",
            "self_attn.norm_added_q.": f"{gen_lp}.cross_attention.norm_q.",
            "self_attn.norm_added_k.": f"{gen_lp}.cross_attention.norm_k.",
            # Norms
            "input_layernorm.": f"{und_lp}.input_layernorm.",
            "post_attention_layernorm.": f"{und_lp}.post_attention_layernorm.",
            "input_layernorm_moe_gen.": f"{gen_lp}.input_layernorm.",
            "post_attention_layernorm_moe_gen.": f"{gen_lp}.post_attention_layernorm.",
            # UND MLP
            "mlp.gate_proj.": f"{und_lp}.mlp.gate_proj.",
            "mlp.up_proj.": f"{und_lp}.mlp.up_proj.",
            "mlp.down_proj.": f"{und_lp}.mlp.down_proj.",
            # GEN MLP
            "mlp_moe_gen.gate_proj.": f"{gen_lp}.mlp.gate_proj.",
            "mlp_moe_gen.up_proj.": f"{gen_lp}.mlp.up_proj.",
            "mlp_moe_gen.down_proj.": f"{gen_lp}.mlp.down_proj.",
        }

        for pattern, replacement in _LAYER_MAP.items():
            if rest.startswith(pattern):
                suffix = rest[len(pattern) :]
                return replacement + suffix

        return None

    # Checkpoint adapters use this hook before model-specific weight loading.
    remap_checkpoint_key = _remap_ckpt_key

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        """Stream-remap checkpoint weights and load via AutoWeightsLoader.

        Handles quantization, TP-aware weight_loader, and buffer loading.
        Returns the set of loaded parameter names for strict validation.
        """
        state = self.state_dict()
        allowed = set(state.keys())
        tp_aware = {n for n, p in self.named_parameters() if hasattr(p, "weight_loader")}

        def _remapped_weights() -> Iterable[tuple[str, torch.Tensor]]:
            total = kept = 0
            for name, tensor in weights:
                total += 1
                if name in allowed or name in tp_aware:
                    kept += 1
                    yield name, tensor
                    continue
                remapped = self._remap_ckpt_key(name)
                if remapped is not None and (remapped in allowed or remapped in tp_aware):
                    kept += 1
                    yield remapped, tensor
            if _is_rank_zero():
                logger.info(
                    "Cosmos3 weight remap: kept %d/%d tensors",
                    kept,
                    total,
                )

        loader = AutoWeightsLoader(self)
        loaded = loader.load_weights(_remapped_weights())
        self.transformer.post_load_weights()
        self.transformer.eval()
        self.transformer.validate_loaded_weights(loaded)
        if getattr(self.transformer, "sound_gen", False):
            sound_markers = ("audio_proj_in.", "audio_proj_out.", "audio_modality_embed")
            missing = [marker.rstrip(".") for marker in sound_markers if not any(marker in name for name in loaded)]
            if missing:
                raise ValueError(
                    "Cosmos3 transformer config enables sound generation, but "
                    f"the checkpoint is missing sound weights for {missing}. "
                    "Use a sound-capable transformer checkpoint."
                )
        if getattr(self.transformer, "action_gen", False):
            action_markers = ("action_proj_in.", "action_proj_out.", "action_modality_embed")
            missing = [marker.rstrip(".") for marker in action_markers if not any(marker in name for name in loaded)]
            if missing:
                raise ValueError(
                    "Cosmos3 transformer config enables action generation, but "
                    f"the checkpoint is missing action weights for {missing}. "
                    "Use an action-capable transformer checkpoint."
                )
        return loaded

    def predict_noise(self, **kwargs) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """Override CFGParallelMixin.predict_noise for Cosmos3.

        The transformer returns the raw prediction: video-only as a tensor,
        or a tuple in video, action, sound order for multimodal generation.
        """

        def _to_model_dtype(value: Any) -> Any:
            if isinstance(value, torch.Tensor):
                return value.to(self.dtype)
            if isinstance(value, list):
                return [_to_model_dtype(item) for item in value]
            if isinstance(value, tuple):
                return tuple(_to_model_dtype(item) for item in value)
            return value

        def _to_sampling_dtype(
            prediction: torch.Tensor | tuple[torch.Tensor, ...],
        ) -> torch.Tensor | tuple[torch.Tensor, ...]:
            if isinstance(prediction, tuple):
                return tuple(item.to(self.sampling_dtype) for item in prediction)
            return prediction.to(self.sampling_dtype)

        cache_key = kwargs.pop("_cosmos3_cache_key", None)
        context_name = str(kwargs.pop("_cache_context", "cond"))
        for key in (
            "hidden_states",
            "action_latents",
            "sound_latents",
            "control_latents",
            "noisy_frame_mask",
            "action_noisy_mask",
        ):
            if key in kwargs and kwargs[key] is not None:
                kwargs[key] = _to_model_dtype(kwargs[key])

        context_factory = getattr(self, "_cache_context_factory", None)
        context = context_factory(context_name) if callable(context_factory) else nullcontext()
        with context:
            if cache_key is None:
                prediction = self.transformer(**kwargs)
            else:
                branch_caches = self._cosmos3_branch_caches
                if branch_caches is None:
                    prediction = self.transformer(**kwargs)
                else:
                    cache_key = str(cache_key)
                    self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get(
                        cache_key, (None, None)
                    )
                    prediction = self.transformer(**kwargs)
                    branch_caches[cache_key] = (self.transformer.cached_kv, self.transformer.cached_freqs_gen)
        return _to_sampling_dtype(prediction)

    def combine_multi_branch_cfg_noise(
        self,
        predictions: list[torch.Tensor | tuple[torch.Tensor, ...]],
        true_cfg_scale: float | dict[str, float],
        cfg_normalize: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        if not isinstance(true_cfg_scale, dict) or true_cfg_scale.get("mode") != "cosmos3_transfer":
            return super().combine_multi_branch_cfg_noise(predictions, true_cfg_scale, cfg_normalize)
        mode = str(true_cfg_scale.get("branch_mode", "control_and_text"))
        guidance_scale = float(true_cfg_scale.get("guidance_scale", 1.0))
        control_guidance = float(true_cfg_scale.get("control_guidance", 1.0))

        if mode == "control_only":
            if len(predictions) != 2:
                raise ValueError(f"Cosmos3 transfer control-only CFG expects 2 branches, got {len(predictions)}.")
            cond_full, cond_no_control = predictions
            if isinstance(cond_full, tuple) or isinstance(cond_no_control, tuple):
                raise ValueError("Cosmos3 transfer control-only CFG expects video-only tensor predictions.")
            cfg_reference = cond_full
            combined = cond_no_control + control_guidance * (cond_full - cond_no_control)
        elif mode == "text_only":
            if len(predictions) != 2:
                raise ValueError(f"Cosmos3 transfer text-only CFG expects 2 branches, got {len(predictions)}.")
            cond_full, uncond_full = predictions
            if isinstance(cond_full, tuple) or isinstance(uncond_full, tuple):
                raise ValueError("Cosmos3 transfer text-only CFG expects video-only tensor predictions.")
            cfg_reference = cond_full
            combined = uncond_full + guidance_scale * (cond_full - uncond_full)
        elif mode == "control_and_text":
            if len(predictions) != 3:
                raise ValueError(f"Cosmos3 transfer control+text CFG expects 3 branches, got {len(predictions)}.")
            cond_full, cond_no_control, uncond_full = predictions
            if isinstance(cond_full, tuple) or isinstance(cond_no_control, tuple) or isinstance(uncond_full, tuple):
                raise ValueError("Cosmos3 transfer control+text CFG expects video-only tensor predictions.")
            cfg_reference = cond_full
            control_cond = cond_no_control + control_guidance * (cond_full - cond_no_control)
            combined = uncond_full + guidance_scale * (control_cond - uncond_full)
        else:
            raise ValueError(f"Unknown Cosmos3 transfer CFG branch_mode={mode!r}.")

        if cfg_normalize:
            combined = self.cfg_normalize_function(cfg_reference, combined)
        return combined

    @staticmethod
    def _cfg_parallel_active() -> bool:
        try:
            return get_classifier_free_guidance_world_size() > 1
        except Exception:
            return False

    def _cache_requires_paired_cfg(self) -> bool:
        """Whether the sequential-CFG denoising loop must keep paired forwards.

        cache-dit wraps the GEN pathway with ``has_separate_cfg=True`` and
        distinguishes the conditional vs unconditional passes purely by the
        parity of its transformer-forward counter.  The T2I ``guidance_interval``
        optimization that skips the uncond pass outside the interval would
        desync that accounting (cond passes get mislabeled as uncond and the
        per-generation step counter drifts).  ``enable_cache_for_cosmos3`` sets
        the marker below when it enables cache-dit on this pipeline; the loop
        then keeps both passes and neutralizes CFG via scale=1.0 instead.

        Returns False when cache-dit is not active, preserving the skip speedup.
        """
        return self._cache_dit_requires_paired_cfg

    @staticmethod
    def _get_sp_param(sp: OmniDiffusionSamplingParams, key: str, default: Any = None) -> Any:
        """Read a runtime control from sampling params.

        Order of precedence:
            1. ``sp.extra_args[key]`` - preferred path; the OpenAI image/video
               endpoints surface custom controls here (see e.g.
               ``serving_video.py`` writing ``extra_args['flow_shift']``).
            2. direct attribute on ``sp`` - backward compat for callers that
               set attributes directly.
            3. ``default``.

        Skipping this helper would cause API-driven overrides like
        ``request.flow_shift`` (forwarded as ``extra_args['flow_shift']``) to
        be silently ignored.
        """
        extra = sp.extra_args or {}
        if extra.get(key) is not None:
            return extra[key]
        val = getattr(sp, key, None)
        if val is not None:
            return val
        return default

    def _get_robolab_transform(self, *, format_prompt_as_json: bool = False):
        transforms = getattr(self, "_robolab_transforms", None)
        if transforms is None:
            transforms = {}
            self._robolab_transforms = transforms
        if format_prompt_as_json not in transforms:
            action_dim = int(getattr(self.transformer, "action_dim", 64))
            transforms[format_prompt_as_json] = lazy_action_transform_pipeline(
                action_dim,
                format_prompt_as_json=format_prompt_as_json,
            )
        return transforms[format_prompt_as_json]

    def _build_robolab_policy_inputs(
        self,
        sp: OmniDiffusionSamplingParams,
        prompt_data: Any | None = None,
        request_id: str | None = None,
    ) -> RoboLabPolicyInputs | None:
        extra = sp.extra_args if isinstance(sp.extra_args, dict) else {}
        obs = extra.get("robot_obs")
        if obs is None:
            obs = extra.get("observation")
        if obs is None:
            return None
        if not isinstance(obs, dict):
            raise TypeError(f"Cosmos3 RoboLab observation must be a dict, got {type(obs)!r}.")

        prompt = obs.get("prompt")
        if not isinstance(prompt, str):
            raise ValueError("RoboLab observation must contain string key 'prompt'.")

        def extra_param(key: str, default: Any) -> Any:
            value = extra.get(key)
            return default if value is None else value

        def extra_param_alias(primary_key: str, alias_key: str, default: Any) -> Any:
            value = extra.get(primary_key)
            if value is not None:
                return value
            value = extra.get(alias_key)
            return default if value is None else value

        action_space = normalize_robolab_action_space(extra_param("action_space", ROBOLAB_DEFAULT_ACTION_SPACE))
        action_chunk_size = int(extra_param("action_chunk_size", ROBOLAB_DEFAULT_ACTION_CHUNK_SIZE))
        raw_action_dim_default = (
            ROBOLAB_DEFAULT_RAW_ACTION_DIM if action_space == "joint_pos" else ROBOLAB_MIDTRAIN_RAW_ACTION_DIM
        )
        raw_action_dim = int(extra_param("raw_action_dim", raw_action_dim_default))
        image_h = int(extra_param("image_height", ROBOLAB_DEFAULT_IMAGE_HEIGHT))
        image_w = int(extra_param("image_width", ROBOLAB_DEFAULT_IMAGE_WIDTH))
        history_length = int(extra_param("history_length", 1))
        use_state = self._truthy(extra_param("use_state", True))
        resolution = str(extra_param("resolution", ROBOLAB_DEFAULT_RESOLUTION))
        fps = float(extra_param("conditioning_fps", ROBOLAB_DEFAULT_CONDITIONING_FPS))
        domain_name = str(extra_param("domain_name", ROBOLAB_DEFAULT_DOMAIN_NAME))
        domain_id = get_robolab_domain_id(domain_name)
        format_prompt_as_json = self._truthy(extra_param("format_prompt_as_json", False))

        if use_state and history_length < 1:
            raise ValueError("RoboLab history_length must be >= 1 when use_state is true.")
        if action_chunk_size <= 0:
            raise ValueError(f"RoboLab action_chunk_size must be positive, got {action_chunk_size}.")
        if raw_action_dim <= 0:
            raise ValueError(f"RoboLab raw_action_dim must be positive, got {raw_action_dim}.")

        try:
            image = extract_robolab_image(obs)
        except ValueError as exc:
            image = extract_robolab_prompt_image(prompt_data)
            if image is None:
                raise exc
        if image.shape[:2] != (image_h, image_w):
            image = resize_rgb_uint8(image, (image_h, image_w))

        t_frames = action_chunk_size + 1
        video = torch.zeros((3, t_frames, image_h, image_w), dtype=torch.uint8)
        video[:, 0] = torch.from_numpy(image.copy()).permute(2, 0, 1)

        use_state_rows = 1 if use_state else 0
        action = torch.zeros((action_chunk_size + use_state_rows, raw_action_dim), dtype=torch.float32)
        history_action = None
        num_history_rows = history_length - use_state_rows
        gripper_position = 1.0 - ensure_gripper_array(obs["observation/gripper_position"])

        if action_space == "joint_pos":
            joint_position = ensure_2d_float_array(obs["observation/joint_position"], "observation/joint_position", 7)
            if use_state:
                action[0] = torch.from_numpy(np.concatenate((joint_position[-1], gripper_position[-1])))
            if num_history_rows > 0:
                if len(joint_position) < num_history_rows + 1:
                    raise ValueError("Not enough joint_position rows for requested history_length.")
                history_np = np.concatenate(
                    (joint_position[-num_history_rows - 1 : -1], gripper_position[-num_history_rows - 1 : -1]),
                    axis=-1,
                )
                history_action = torch.from_numpy(history_np).float()
        else:
            eef_pos = ensure_2d_float_array(obs["observation/eef_pos"], "observation/eef_pos", 3)
            eef_quat = ensure_2d_float_array(obs["observation/eef_quat"], "observation/eef_quat", 4)
            if use_state:
                rot6d = convert_midtrain_rotation(eef_quat[-1], "quat_xyzw", "rot6d")
                action[0] = torch.from_numpy(np.concatenate((eef_pos[-1], rot6d, gripper_position[-1])))
            if num_history_rows > 0:
                if len(eef_pos) < num_history_rows + 1 or len(eef_quat) < num_history_rows + 1:
                    raise ValueError("Not enough eef_pos/eef_quat rows for requested history_length.")
                poses_abs = build_abs_pose_from_components(eef_pos, eef_quat, "quat_xyzw")
                poses_rel = pose_abs_to_rel(poses_abs, rotation_format="rot6d", pose_convention="backward_framewise")
                history_np = np.concatenate(
                    [poses_rel[-num_history_rows:], gripper_position[-num_history_rows:]],
                    axis=-1,
                )
                history_action = torch.from_numpy(history_np).float()

        sample: dict[str, Any] = {
            "ai_caption": prompt,
            "video": video,
            "action": action,
            # Cosmos Framework consumes this as an integer conditioning bucket.
            "conditioning_fps": torch.tensor(fps, dtype=torch.long),
            # Cosmos Framework 1.2 renamed the internal policy transform mode to
            # ``wam``. The public vLLM action_mode remains ``policy``.
            "mode": "wam",
            "domain_id": torch.tensor(domain_id, dtype=torch.long),
            "viewpoint": "concat_view",
            "additional_view_description": ROBOLAB_CONCAT_VIEW_DESCRIPTION,
        }
        if history_action is not None:
            sample["history_action"] = history_action

        sample = self._get_robolab_transform(format_prompt_as_json=format_prompt_as_json)(sample, resolution)
        if isinstance(sample.get("ai_caption"), dict):
            sample["ai_caption"] = json.dumps(sample["ai_caption"])
        sequence_plan = sample["sequence_plan"]
        video_tensor = sample["video"].float() / 127.5 - 1.0
        raw_action_dim_tensor = sample.get("raw_action_dim")
        if isinstance(raw_action_dim_tensor, torch.Tensor):
            transformed_raw_action_dim = int(raw_action_dim_tensor.item())
        else:
            transformed_raw_action_dim = raw_action_dim

        return RoboLabPolicyInputs(
            prompt=sample["ai_caption"],
            video_tensor=video_tensor.unsqueeze(0),
            action_tensor=sample["action"].float(),
            action_condition_indexes=list(getattr(sequence_plan, "condition_frame_indexes_action", []) or []),
            action_start_frame_offset=int(getattr(sequence_plan, "action_start_frame_offset", 1)),
            raw_action_dim=transformed_raw_action_dim,
            domain_id=domain_id,
            fps=fps,
            height=int(video_tensor.shape[-2]),
            width=int(video_tensor.shape[-1]),
            image_size=sample.get("image_size"),
            num_frames=int(video_tensor.shape[1]),
            num_inference_steps=int(
                extra_param_alias("num_inference_steps", "num_steps", ROBOLAB_DEFAULT_NUM_INFERENCE_STEPS)
            ),
            guidance_scale=float(extra_param_alias("guidance_scale", "guidance", ROBOLAB_DEFAULT_GUIDANCE_SCALE)),
            flow_shift=float(extra_param_alias("flow_shift", "shift", ROBOLAB_DEFAULT_FLOW_SHIFT)),
            seed=next_robolab_seed(extra, obs, request_id),
            history_length=history_length,
            action_space=action_space,
            observation=obs,
        )

    @staticmethod
    def _build_action_condition_mask_from_indexes(
        indexes: list[int],
        action_length: int,
        *,
        device: torch.device,
        dtype: torch.dtype,
    ) -> torch.Tensor:
        mask = torch.zeros(1, action_length, 1, device=device, dtype=dtype)
        for idx in indexes:
            if idx < 0 or idx >= action_length:
                raise ValueError(f"Action condition index {idx} is out of range for action length {action_length}.")
            mask[:, idx, :] = 1.0
        return mask

    def _forward_robolab_policy(
        self,
        sp: OmniDiffusionSamplingParams,
        inputs: RoboLabPolicyInputs,
        pipeline_start: float,
        session_id: str | None = None,
    ) -> DiffusionOutput:
        if self.is_distilled_model:
            raise self._distilled_unsupported_error("RoboLab/action policy requests are unsupported.")
        if not getattr(self.transformer, "action_gen", False):
            raise ValueError(
                "Cosmos3 RoboLab policy serving was requested, but the transformer "
                "was initialized without action modules. Check that the checkpoint "
                "config enables action_gen and includes action weights."
            )

        action_mode = ACTION_MODE_POLICY
        height = inputs.height
        width = inputs.width
        num_frames = inputs.num_frames
        action_chunk_size = int(inputs.action_tensor.shape[0])
        num_inference_steps = inputs.num_inference_steps
        guidance_scale = float(inputs.guidance_scale)
        flow_shift_target = float(inputs.flow_shift)
        domain_id = int(inputs.domain_id)
        frame_rate = self._get_sp_param(sp, "resolved_frame_rate") or self._get_sp_param(sp, "frame_rate") or inputs.fps
        max_sequence_length = (
            self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH)
            or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH
        )
        use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False))

        self._guidance_scale = guidance_scale
        self._num_timesteps = num_inference_steps

        generator = sp.generator
        if generator is None:
            generator = torch.Generator(device=self.device).manual_seed(int(inputs.seed))

        cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts(
            inputs.prompt,
            "",
            num_frames,
            frame_rate,
            height,
            width,
            max_sequence_length,
            sp,
            use_system_prompt,
            is_t2i=False,
        )

        action_video_tensor = inputs.video_tensor
        if action_video_tensor.ndim == 4:
            action_video_tensor = action_video_tensor.unsqueeze(0)
        if action_video_tensor.ndim != 5:
            raise ValueError(
                "Cosmos3 RoboLab action video tensor must have shape [1, 3, T, H, W] "
                f"or [3, T, H, W], got {tuple(action_video_tensor.shape)}."
            )
        if action_video_tensor.shape[2] < num_frames:
            pad = action_video_tensor[:, :, -1:].repeat(1, 1, num_frames - action_video_tensor.shape[2], 1, 1)
            action_video_tensor = torch.cat([action_video_tensor, pad], dim=2)
        elif action_video_tensor.shape[2] > num_frames:
            action_video_tensor = action_video_tensor[:, :, :num_frames]

        action_latents, action_velocity_mask, action_condition_latents, raw_action_dim = self._prepare_action_latents(
            mode=action_mode,
            action_chunk_size=action_chunk_size,
            raw_action_dim=int(inputs.raw_action_dim),
            generator=generator,
            sp=sp,
            clean_action=inputs.action_tensor,
            condition_indexes=inputs.action_condition_indexes,
        )
        action_offset = int(inputs.action_start_frame_offset)

        latents, velocity_mask, condition_latents = self._prepare_latents_action_video(
            action_video_tensor,
            action_mode,
            height,
            width,
            num_frames,
            generator,
            image_size=inputs.image_size,
        )
        image_latent = condition_latents[:, :, 0:1]

        video_shape = (latents.shape[2], latents.shape[3], latents.shape[4])
        shared_kwargs = dict(
            video_shape=video_shape,
            fps=frame_rate,
            noisy_frame_mask=velocity_mask,
            action_domain_ids=torch.tensor([domain_id], dtype=torch.long, device=self.device),
            action_noisy_mask=action_velocity_mask,
            action_start_frame_offset=action_offset,
            action_fps=float(self._get_sp_param(sp, "action_fps", frame_rate) or frame_rate),
        )

        scheduler = build_robolab_unipc_scheduler(num_inference_steps, flow_shift_target, self.device)
        _, action_latents = self.diffuse(
            latents=latents,
            timesteps=scheduler.timesteps,
            cond_ids=cond_ids,
            cond_mask=cond_mask,
            uncond_ids=uncond_ids,
            uncond_mask=uncond_mask,
            guidance_scale=guidance_scale,
            shared_kwargs=shared_kwargs,
            action_latents=action_latents,
            action_velocity_mask=action_velocity_mask,
            action_condition_latents=action_condition_latents,
            sound_latents=None,
            velocity_mask=velocity_mask,
            image_latent=image_latent,
            condition_latents=condition_latents,
            guidance_interval=None,
            raw_action_dim=raw_action_dim,
            scheduler=scheduler,
            session_id=session_id,
            generator=generator,
        )

        if _is_rank_zero():
            logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)

        action = action_latents[:, :, :raw_action_dim].detach().cpu()
        action_output: dict[str, Any] = {
            "payload": {
                "actions": action,
            },
            "metadata": {
                "actions": {
                    "raw_action_dim": raw_action_dim,
                    "action_mode": action_mode,
                    "domain_id": domain_id,
                },
                "common": {
                    "action_only_output": True,
                },
                "internal": {
                    "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs),
                },
            },
        }
        return DiffusionOutput(
            output=action_output,
            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
        )

    @staticmethod
    def _truthy(value) -> bool:
        if isinstance(value, str):
            return value.strip().lower() in {"1", "true", "yes", "on"}
        return bool(value)

    @classmethod
    def _get_prompt_param(cls, prompt_data, key: str, default=None):
        if not isinstance(prompt_data, dict):
            return default
        if prompt_data.get(key) is not None:
            return prompt_data[key]
        additional = prompt_data.get("additional_information")
        if isinstance(additional, dict) and additional.get(key) is not None:
            return additional[key]
        return default

    @classmethod
    def _is_sound_request(cls, prompt_data, sp) -> bool:
        for key in ("generate_sound", "sound_gen"):
            if cls._truthy(cls._get_prompt_param(prompt_data, key, None)):
                return True
            if cls._truthy(cls._get_sp_param(sp, key, None)):
                return True
        return False

    @classmethod
    def _get_action_mode(cls, prompt_data, sp) -> str | None:
        return normalize_action_mode(
            cls._get_sp_param(sp, "action_mode", cls._get_prompt_param(prompt_data, "action_mode", None))
        )

    def _get_sound_tokenizer(self):
        if self._sound_tokenizer is None:
            from .sound_tokenizer import Cosmos3SoundTokenizer

            self._sound_tokenizer = Cosmos3SoundTokenizer.from_config(self.od_config)
        return self._sound_tokenizer

    @staticmethod
    def _is_t2i_request(req: DiffusionRequestBatch) -> bool:
        """Return whether request-level modalities select image output.

        Only ``"image"`` switches Cosmos3 into T2I. ``"video"`` and omitted
        modalities select video output. ``"text"`` and ``"audio"`` are accepted
        compatibility values for callers that share prompt schemas, but they do
        not select text or audio output in this pipeline. ``"image"`` and
        ``"video"`` cannot be requested together.
        """
        if not req.prompts:
            return False
        first_prompt = req.prompts[0]
        modalities = first_prompt.get("modalities", []) if isinstance(first_prompt, dict) else []
        if modalities is None:
            modalities = []
        if isinstance(modalities, str):
            modalities = [modalities]
        if "image" in modalities and "video" in modalities:
            raise ValueError("Cosmos3 prompt modalities cannot request both image and video output.")

        accepted_modalities = ["image", "video", "text", "audio"]
        if any(x not in accepted_modalities for x in modalities):
            raise ValueError(f"Incorrect modality value in {modalities}, expected one of {accepted_modalities}.")
        return "image" in modalities

    def _set_timesteps(self, num_inference_steps: int, device: str | torch.device, shift: float) -> None:
        if self.is_distilled_model:
            self.scheduler.set_timesteps(sigmas=self._scheduler_init_t_list, device=device)
        else:
            self.scheduler.set_timesteps(
                num_inference_steps,
                device=device,
                shift=shift,
            )

    def _set_flow_shift(self, target_shift: float) -> None:
        """Select the shift applied by FlowUniPC when timesteps are built."""
        self._current_flow_shift = float(target_shift)

    @staticmethod
    def _resolve_seed(sp: OmniDiffusionSamplingParams, generator: torch.Generator | None, default: int = 42) -> int:
        if sp.seed is not None:
            return int(sp.seed)
        if isinstance(generator, torch.Generator):
            return int(generator.initial_seed())
        return int(default)

    def _resolve_guidance_scale(self, sp: OmniDiffusionSamplingParams, default: float) -> float:
        if self.is_distilled_model:
            if sp.guidance_scale_provided and float(sp.guidance_scale) != 1.0:
                logger.warning_once(
                    "Distilled Cosmos3 checkpoints run without classifier-free guidance. "
                    "The requested guidance_scale is overridden to 1.0; negative_prompt does not affect generation."
                )
            return 1.0
        if sp.guidance_scale_provided:
            return float(sp.guidance_scale)
        return float(default)

    def _cpu_scheduler_state(self) -> None:
        # We need to move scheduler tensors to CPU, as unipc from diffusers assumes they are on CPU.
        # However, after the creation they are on GPU due to "with target_device:" in diffusers_loader.py
        for name, value in vars(self.scheduler).items():
            if isinstance(value, torch.Tensor) and value.device.type != "cpu":
                setattr(self.scheduler, name, value.cpu())

    @property
    def guidance_scale(self):
        return self._guidance_scale

    @property
    def do_classifier_free_guidance(self):
        return self._guidance_scale is not None and self._guidance_scale > 1.0

    @property
    def num_timesteps(self):
        return self._num_timesteps

    @property
    def current_step_index(self):
        return self._current_step_index

    @property
    def current_sigma(self):
        return self._current_sigma

    def _set_denoise_step_metadata(
        self,
        step_index: int,
        timesteps: torch.Tensor,
        scheduler: Any,
    ) -> None:
        self._current_step_index = step_index
        self._num_timesteps = len(timesteps)
        sigmas = getattr(scheduler, "sigmas", None)
        self._current_sigma = sigmas[step_index] if sigmas is not None and step_index < len(sigmas) else None

    def _clear_denoise_step_metadata(self) -> None:
        self._current_step_index = None
        self._current_sigma = None

    def _set_mixed_precision_step(self, step_index: int, num_steps: int) -> None:
        setter = getattr(self.transformer, "set_mixed_precision_step", None)
        if setter is not None:
            setter(step_index, num_steps)

    def _reset_mixed_precision(self) -> None:
        resetter = getattr(self.transformer, "reset_mixed_precision", None)
        if resetter is not None:
            resetter()

    @staticmethod
    def _distilled_unsupported_error(detail: str) -> ValueError:
        return ValueError(
            f"Cosmos3 distilled checkpoints support only text-to-image and image-to-video generation; {detail}"
        )

    def _validate_distilled_generation_mode(
        self,
        *,
        is_t2i: bool,
        image_tensor: torch.Tensor | None,
        action_enabled: bool,
        transfer_config: Cosmos3TransferConfig | None,
        is_v2v: bool,
        sound_enabled: bool,
    ) -> None:
        if not self.is_distilled_model:
            return
        if action_enabled:
            raise self._distilled_unsupported_error("action requests are unsupported.")
        if transfer_config is not None:
            raise self._distilled_unsupported_error("transfer requests are unsupported.")
        if is_v2v:
            raise self._distilled_unsupported_error("video-to-video requests are unsupported.")
        if sound_enabled:
            raise self._distilled_unsupported_error("sound generation is unsupported.")
        if not is_t2i and image_tensor is None:
            raise self._distilled_unsupported_error("text-to-video requests are unsupported.")
        if is_t2i and image_tensor is not None:
            raise self._distilled_unsupported_error("image-conditioned image generation is unsupported.")

    def _validate_edge_generation_mode(
        self,
        *,
        transfer_config: Cosmos3TransferConfig | None,
        is_v2v: bool,
        sound_enabled: bool,
    ) -> None:
        if not self.is_edge_model:
            return
        if sound_enabled:
            raise ValueError("Cosmos3 Edge checkpoints do not support sound generation.")
        if transfer_config is not None:
            raise ValueError("Cosmos3 Edge checkpoints do not support transfer inference.")
        if is_v2v:
            raise ValueError("Cosmos3 Edge checkpoints do not support video-to-video generation.")

    # -- Prompt formatting -----------------------------------------------------

    @staticmethod
    def _apply_metadata_templates(
        prompt: str,
        num_frames: int,
        frame_rate: float,
        height: int,
        width: int,
        duration_template: str | None = COSMOS3_DURATION_TEMPLATE,
        resolution_template: str | None = COSMOS3_RESOLUTION_TEMPLATE,
        force_duration_template: bool = False,
    ) -> str:
        """
        Append duration and resolution metadata to a prompt.
        """
        prompt = prompt.strip()
        if duration_template is None and resolution_template is None:
            return prompt

        parts: list[str] = []
        head = prompt.rstrip(".").strip()
        if head:
            parts.append(head)
        if duration_template is not None and (num_frames > 1 or force_duration_template):
            duration = num_frames / frame_rate
            parts.append(duration_template.format(duration=duration, fps=frame_rate).rstrip("."))
        if resolution_template is not None:
            parts.append(resolution_template.format(height=height, width=width).rstrip("."))
        if not parts:
            return ""
        return ". ".join(parts) + "."

    # -- Tokenization --------------------------------------------------------

    def _tokenize_prompt(
        self,
        text: str,
        max_sequence_length: int,
        use_system_prompt: bool = False,
        system_prompt: str | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Tokenize a prompt using the Qwen2 chat template.

        Returns (input_ids, attention_mask) as [1, S] tensors on device.
        """
        conversations = []
        if use_system_prompt:
            conversations.append(
                {
                    "role": "system",
                    "content": system_prompt or COSMOS3_SYSTEM_PROMPT,
                }
            )
        conversations.append({"role": "user", "content": text})

        token_ids = self._normalize_token_ids(
            self.tokenizer.apply_chat_template(conversations, tokenize=True, add_generation_prompt=True)
        )
        original_token_count = len(token_ids)
        if original_token_count > max_sequence_length and _is_rank_zero():
            logger.warning(
                "Cosmos3 prompt token_ids shortened to max_sequence_length: "
                "original_token_count=%d, max_sequence_length=%d, removed_token_count=%d",
                original_token_count,
                max_sequence_length,
                original_token_count - max_sequence_length,
            )
        token_ids = token_ids[:max_sequence_length]
        token_ids.append(self.tokenizer.eos_token_id)  # 151645
        token_ids.append(self.tokenizer.convert_tokens_to_ids("<|vision_start|>"))  # 151652
        seq_len = len(token_ids)

        # No right-padding: the prompt is tokenized to its natural length.
        # The UND pathway uses causal self-attention with no padding mask and
        # the GEN cross-attention K/V is trimmed to the real text length, so
        # padding to a fixed length only added dead compute and never changed
        # the output.  ``max_sequence_length`` is kept purely as a truncation
        # cap (above).  The mask is therefore all ones.
        attention_mask = [1] * seq_len

        input_ids = torch.tensor([token_ids], dtype=torch.long, device=self.device)
        attention_mask = torch.tensor([attention_mask], dtype=torch.long, device=self.device)
        return input_ids, attention_mask

    @staticmethod
    def _normalize_token_ids(tokenized_output: object) -> list[int]:
        """Normalize tokenizer outputs into a flat ``list[int]``.

        Different Transformers/tokenizers versions can return ``list[int]``,
        a mapping/BatchEncoding with ``input_ids``, tensors, or
        ``tokenizers.Encoding`` objects from ``apply_chat_template``.
        """
        token_ids = tokenized_output
        while True:
            if isinstance(token_ids, dict) and "input_ids" in token_ids:
                token_ids = token_ids["input_ids"]
            elif hasattr(token_ids, "input_ids"):
                token_ids = token_ids.input_ids
            elif hasattr(token_ids, "ids"):
                token_ids = token_ids.ids
            elif hasattr(token_ids, "tolist"):
                token_ids = token_ids.tolist()
            elif isinstance(token_ids, tuple):
                token_ids = list(token_ids)
            elif isinstance(token_ids, list) and len(token_ids) == 1:
                first = token_ids[0]
                if isinstance(first, list | tuple):
                    token_ids = list(first)
                elif hasattr(first, "ids") or hasattr(first, "input_ids"):
                    token_ids = first
                elif hasattr(first, "tolist"):
                    first_list = first.tolist()
                    if isinstance(first_list, list | tuple):
                        token_ids = list(first_list)
                    else:
                        break
                else:
                    break
            else:
                break

        if not isinstance(token_ids, list):
            raise TypeError(
                "Cosmos3 tokenizer must return token IDs as a list-like value; "
                f"got {type(token_ids).__name__}: {token_ids!r}"
            )

        normalized_ids = []
        for idx, token_id in enumerate(token_ids):
            if hasattr(token_id, "item"):
                token_id = token_id.item()
            try:
                normalized_ids.append(int(token_id))
            except (TypeError, ValueError) as exc:
                raise TypeError(
                    "Cosmos3 tokenizer returned a non-integer token at "
                    f"index {idx}: {type(token_id).__name__}: {token_id!r}"
                ) from exc
        return normalized_ids

    # -- Latent preparation --------------------------------------------------

    def _prepare_latents(
        self,
        height: int,
        width: int,
        num_frames: int,
        generator: torch.Generator,
    ) -> torch.Tensor:
        num_channels_latents = self.transformer.latent_channel_size
        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
        shape = (
            1,
            num_channels_latents,
            num_latent_frames,
            height // self.vae_scale_factor_spatial,
            width // self.vae_scale_factor_spatial,
        )
        return randn_tensor(shape, generator=generator, device=self.device, dtype=self.sampling_dtype)

    def _prepare_sound_latents(
        self,
        target_audio_samples: int,
        generator: torch.Generator,
        *,
        sp_video_shape: tuple[int, int, int] | None = None,
        sp_num_vision_items: int = 1,
    ) -> tuple[torch.Tensor, int]:
        sound_tokenizer = self._get_sound_tokenizer()
        hop_size = int(
            getattr(sound_tokenizer, "hop_size", None) or getattr(sound_tokenizer, "temporal_compression_factor")
        )
        latent_frames = max(1, math.ceil(max(1, int(target_audio_samples)) / hop_size))
        if sp_video_shape is not None:
            latent_frames = self.transformer.sound_latent_frames_for_sequence_parallel(
                video_shape=sp_video_shape,
                sound_frames=latent_frames,
                num_vision_items=sp_num_vision_items,
            )
        sound_dim = int(getattr(sound_tokenizer, "latent_ch", 64))
        transformer_sound_dim = int(getattr(self.transformer, "sound_dim", sound_dim))
        if sound_dim != transformer_sound_dim:
            raise ValueError(
                "Cosmos3 sound tokenizer latent channels do not match transformer "
                f"sound_dim: tokenizer={sound_dim}, transformer={transformer_sound_dim}."
            )
        latents = randn_tensor(
            (1, sound_dim, latent_frames),
            generator=generator,
            device=self.device,
            dtype=self.sampling_dtype,
        )
        return latents, latent_frames

    def _resolve_sound_target_samples(
        self,
        sp,
        num_frames: int,
        frame_rate: float,
    ) -> tuple[int, float, int]:
        sound_tokenizer = self._get_sound_tokenizer()
        duration = self._get_sp_param(sp, "sound_duration", None)
        if duration is None:
            duration = self._get_sp_param(sp, "audio_duration", None)
        if duration is None:
            duration = num_frames / frame_rate
        duration = max(float(duration), 1.0 / max(float(frame_rate), 1.0))
        sample_rate = int(getattr(sound_tokenizer, "sample_rate", 48000))
        return max(1, int(round(duration * sample_rate))), duration, sample_rate

    # -- VAE decode ----------------------------------------------------------

    def _get_latents_mean_std(self, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]:
        cached = getattr(self, "_latents_mean_std", None)
        if cached is not None:
            latents_mean, latents_std = cached
            if latents_mean.device == device and latents_mean.dtype == dtype:
                return latents_mean, latents_std

        latents_mean = torch.as_tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
        latents_std = torch.as_tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(1, -1, 1, 1, 1)
        self._latents_mean_std = (latents_mean, latents_std)
        return latents_mean, latents_std

    def _to_vae_device(self, tensor: torch.Tensor, *, pin_cpu: bool = False) -> torch.Tensor:
        if tensor.device == self.device and tensor.dtype == self.vae.dtype:
            return tensor

        non_blocking = False
        if tensor.device.type == "cpu" and self.device.type == "cuda":
            if pin_cpu and not tensor.is_pinned():
                tensor = tensor.pin_memory()
            non_blocking = tensor.is_pinned()

        return tensor.to(device=self.device, dtype=self.vae.dtype, non_blocking=non_blocking)

    def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
        latents = latents.to(self.vae.dtype)

        if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"):
            latents_mean, latents_std = self._get_latents_mean_std(latents.device, latents.dtype)
            latents = (latents * latents_std) + latents_mean
        else:
            scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
            latents = latents / scaling_factor

        video = self.vae.decode(latents, return_dict=False)[0]
        return video

    def _decode_sound_latents(
        self,
        sound_latents: torch.Tensor,
        target_audio_samples: int,
    ) -> torch.Tensor:
        sound_tokenizer = self._get_sound_tokenizer()
        audio = sound_tokenizer.decode(sound_latents.to(self.dtype))
        if audio.shape[-1] > target_audio_samples:
            audio = audio[..., :target_audio_samples]
        elif audio.shape[-1] < target_audio_samples:
            audio = torch.nn.functional.pad(audio, (0, target_audio_samples - audio.shape[-1]))
        return audio.detach().cpu()

    # -- Prompt formatting + tokenization (shared by T2V and I2V) ------------

    def _format_and_tokenize_prompts(
        self,
        prompt: str,
        negative_prompt: str,
        num_frames: int,
        frame_rate: float,
        height: int,
        width: int,
        max_sequence_length: int,
        sp: OmniDiffusionSamplingParams,
        use_system_prompt: bool = False,
        is_t2i: bool = False,
        system_prompt: str | None = None,
        prompt_suffix: str | None = None,
        use_duration_template: bool | None = None,
        use_resolution_template: bool | None = None,
        negative_metadata_mode: str | None = None,
        aspect_ratio_override: str | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Format prompts with metadata templates and tokenize.

        Returns (cond_ids, cond_mask, uncond_ids, uncond_mask).

        For T2I (``is_t2i=True``) the duration template is suppressed (no FPS
        or duration concept for a single image) and the image-flavored
        resolution template is used. ``prompt_suffix`` is appended only to the
        positive prompt, after metadata formatting. ``negative_metadata_mode``
        controls whether the negative branch receives no metadata, the same
        metadata as the positive branch, or inverse metadata.
        """
        # Route cosmos3-specific controls through ``_get_sp_param`` so they
        # are picked up from ``extra_args`` (OpenAI endpoint path) as well
        # as from direct attributes.
        if use_duration_template is None:
            use_duration_template = bool(self._get_sp_param(sp, "use_duration_template", False))
        if use_resolution_template is None:
            use_resolution_template = bool(self._get_sp_param(sp, "use_resolution_template", False))
        if negative_metadata_mode is None:
            negative_metadata_mode = str(self._get_sp_param(sp, "negative_metadata_mode", "inverse"))
        negative_metadata_mode = negative_metadata_mode.strip().lower()
        if negative_metadata_mode not in {"none", "same", "inverse"}:
            raise ValueError(
                "Cosmos3 negative_metadata_mode must be one of 'none', 'same', or 'inverse'; "
                f"got {negative_metadata_mode!r}."
            )

        use_duration_template = use_duration_template and not is_t2i
        dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None
        if use_resolution_template:
            res_tmpl = COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_RESOLUTION_TEMPLATE
        else:
            res_tmpl = None
        requested_aspect_ratio = self._get_sp_param(sp, "aspect_ratio", None)
        prompt_aspect_ratio = _json_object_aspect_ratio(prompt)
        metadata_aspect_ratio = requested_aspect_ratio if requested_aspect_ratio is not None else prompt_aspect_ratio
        canvas_aspect_ratio = aspect_ratio_override or _aspect_ratio_for_dimensions(width, height)
        normalized_metadata_ratio = _normalize_aspect_ratio(metadata_aspect_ratio)
        has_aspect_ratio_conflict = (
            aspect_ratio_override is None
            and metadata_aspect_ratio is not None
            and normalized_metadata_ratio != canvas_aspect_ratio
        )
        if has_aspect_ratio_conflict:
            aspect_ratio_source = (
                "requested aspect_ratio" if requested_aspect_ratio is not None else "JSON prompt aspect_ratio"
            )
            if _is_rank_zero():
                logger.warning(
                    "Cosmos3 %s=%r conflicts with the generated %dx%d canvas (WxH); using %s for prompt metadata.",
                    aspect_ratio_source,
                    metadata_aspect_ratio,
                    width,
                    height,
                    canvas_aspect_ratio,
                )
            metadata_aspect_ratio = canvas_aspect_ratio
        elif aspect_ratio_override is not None:
            metadata_aspect_ratio = aspect_ratio_override

        json_prompt = _format_json_object_prompt(
            prompt,
            num_frames=num_frames,
            frame_rate=frame_rate,
            height=height,
            width=width,
            aspect_ratio=metadata_aspect_ratio,
        )
        if json_prompt is not None:
            prompt = json_prompt
        else:
            prompt = self._apply_metadata_templates(
                prompt,
                num_frames,
                frame_rate,
                height,
                width,
                duration_template=dur_tmpl,
                resolution_template=res_tmpl,
            )
        if prompt_suffix:
            prompt = f"{prompt.rstrip()} {prompt_suffix.lstrip()}".strip()
        if _is_rank_zero():
            logger.debug("Final prompt: '%s'", prompt)

        if negative_metadata_mode == "none":
            negative_dur_tmpl = None
            negative_res_tmpl = None
        elif negative_metadata_mode == "same":
            negative_dur_tmpl = dur_tmpl
            negative_res_tmpl = res_tmpl
        else:
            negative_dur_tmpl = COSMOS3_INVERSE_DURATION_TEMPLATE if dur_tmpl else None
            if res_tmpl is None:
                negative_res_tmpl = None
            elif is_t2i:
                negative_res_tmpl = COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE
            else:
                negative_res_tmpl = COSMOS3_INVERSE_RESOLUTION_TEMPLATE
        negative_prompt = self._apply_metadata_templates(
            negative_prompt,
            num_frames,
            frame_rate,
            height,
            width,
            duration_template=negative_dur_tmpl,
            resolution_template=negative_res_tmpl,
            force_duration_template=negative_metadata_mode == "inverse",
        )

        if system_prompt is None:
            default_sys_prompt = COSMOS3_T2I_SYSTEM_PROMPT if is_t2i else COSMOS3_SYSTEM_PROMPT
            system_prompt = self._get_sp_param(sp, "system_prompt", default_sys_prompt) or default_sys_prompt
        cond_ids, cond_mask = self._tokenize_prompt(
            prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt
        )
        uncond_ids, uncond_mask = self._tokenize_prompt(
            negative_prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt
        )
        return cond_ids, cond_mask, uncond_ids, uncond_mask

    # -- I2V latent preparation ---------------------------------------------

    def _normalize_vae_latent(self, latent: torch.Tensor) -> torch.Tensor:
        if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"):
            latents_mean, latents_std = self._get_latents_mean_std(latent.device, latent.dtype)
            latent = (latent - latents_mean) / latents_std
        else:
            scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
            latent = latent * scaling_factor

        return latent

    def _encode_conditioning_image_latent(self, image_tensor: torch.Tensor) -> torch.Tensor:
        """VAE-encode the first I2V conditioning frame.

        I2V only consumes the first conditioning latent frame.  Encoding the
        input image as a one-frame video keeps that latent while avoiding VAE
        work for repeated frames that are later replaced by noise.
        """
        image_tensor = self._to_vae_device(image_tensor, pin_cpu=True)
        video = image_tensor.unsqueeze(2)
        latent = self.vae.encode(video).latent_dist.mode()
        latent = self._normalize_vae_latent(latent)
        return latent[:, :, 0:1, :, :].to(self.sampling_dtype)

    def _latent_hw_from_image_size(self, image_size: Any | None) -> tuple[int, int] | None:
        if image_size is None:
            return None
        if isinstance(image_size, torch.Tensor):
            frame_size = image_size.detach().cpu().flatten()
        else:
            frame_size = torch.as_tensor(image_size).flatten()
        if frame_size.numel() < 4:
            return None
        orig_h = int(frame_size[2].item())
        orig_w = int(frame_size[3].item())
        spatial_factor = int(self.vae_scale_factor_spatial)
        return max(orig_h // spatial_factor, 1), max(orig_w // spatial_factor, 1)

    def _crop_latent_to_image_size(self, latent: torch.Tensor, image_size: Any | None) -> torch.Tensor:
        latent_hw = self._latent_hw_from_image_size(image_size)
        if latent_hw is None:
            return latent
        h_latent, w_latent = latent_hw
        return latent[:, :, :, :h_latent, :w_latent].contiguous()

    def _encode_video_tensor(self, video_tensor: torch.Tensor, image_size: Any | None = None) -> torch.Tensor:
        """VAE-encode a preprocessed pixel video [1, 3, T, H, W]."""
        if video_tensor.ndim == 4:
            video_tensor = video_tensor.unsqueeze(0)
        if video_tensor.ndim != 5:
            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
        if video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")

        video = self._to_vae_device(video_tensor)
        latent = self.vae.encode(video).latent_dist.mode()
        latent = self._normalize_vae_latent(latent)
        latent = self._crop_latent_to_image_size(latent, image_size)
        return latent.to(self.sampling_dtype)

    def _prepare_latents_i2v(
        self,
        image_tensor: torch.Tensor,
        height: int,
        width: int,
        num_frames: int,
        generator: torch.Generator,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Prepare initial latents with frame 0 conditioned on the input image.

        Returns:
            latents: [1, C, T_lat, H_lat, W_lat] with frame 0 = image, rest = noise
            velocity_mask: [1, 1, T_lat, 1, 1] with frame 0 = 0, rest = 1
            image_latent: [1, C, 1, H_lat, W_lat] clean frame 0 for re-injection
        """
        C = self.transformer.latent_channel_size
        T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
        H_lat = height // self.vae_scale_factor_spatial
        W_lat = width // self.vae_scale_factor_spatial

        noise = randn_tensor(
            (1, C, T_lat, H_lat, W_lat),
            generator=generator,
            device=self.device,
            dtype=self.sampling_dtype,
        )

        image_latent = self._encode_conditioning_image_latent(image_tensor)
        latents = noise
        latents[:, :, 0:1, :, :] = image_latent

        velocity_mask = torch.ones(1, 1, T_lat, 1, 1, device=self.device, dtype=self.sampling_dtype)
        velocity_mask[:, :, 0, :, :] = 0.0
        return latents, velocity_mask, image_latent

    def _prepare_latents_v2v(
        self,
        video_tensor: torch.Tensor,
        height: int,
        width: int,
        num_frames: int,
        generator: torch.Generator,
        condition_frame_indexes_vision: Iterable[int] | int | str | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Prepare V2V latents with explicit clean conditioned latent frames."""
        del height, width
        if video_tensor.ndim == 4:
            video_tensor = video_tensor.unsqueeze(0)
        if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
        if video_tensor.shape[2] < 1:
            raise ValueError("Cosmos3 V2V video tensor must contain at least one frame.")

        C = self.transformer.latent_channel_size
        T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
        H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial
        W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial
        indexes = normalize_condition_frame_indexes_vision(condition_frame_indexes_vision)
        out_of_range = [index for index in indexes if index >= T_lat]
        if out_of_range:
            raise ValueError(
                "Cosmos3 condition_frame_indexes_vision contains indexes outside the latent video: "
                f"indexes={indexes}, latent_frames={T_lat}."
            )

        noise = randn_tensor(
            (1, C, T_lat, H_lat, W_lat),
            generator=generator,
            device=self.device,
            dtype=self.sampling_dtype,
        )
        condition_pixel_frames = condition_pixel_frame_count(indexes, self.vae_scale_factor_temporal)
        condition_video = video_tensor[:, :, :condition_pixel_frames]
        if condition_video.shape[2] < condition_pixel_frames:
            pad = condition_video[:, :, -1:].repeat(1, 1, condition_pixel_frames - condition_video.shape[2], 1, 1)
            condition_video = torch.cat([condition_video, pad], dim=2)

        cond_prefix_latent = self._encode_video_tensor(condition_video)
        expected_prefix = (1, C, max(indexes) + 1, H_lat, W_lat)
        if (
            cond_prefix_latent.shape[0] != expected_prefix[0]
            or cond_prefix_latent.shape[1] != expected_prefix[1]
            or cond_prefix_latent.shape[2] < expected_prefix[2]
            or cond_prefix_latent.shape[3:] != expected_prefix[3:]
        ):
            raise ValueError(
                "Cosmos3 V2V condition latent shape mismatch: "
                f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}."
            )

        condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.sampling_dtype)
        condition_latents = torch.zeros_like(noise)
        for index in indexes:
            condition_mask[:, :, index, :, :] = 1.0
            condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1]
        latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise
        velocity_mask = 1.0 - condition_mask
        return latents, velocity_mask, condition_latents

    def _prepare_latents_action_video(
        self,
        video_tensor: torch.Tensor,
        mode: str,
        height: int,
        width: int,
        num_frames: int,
        generator: torch.Generator,
        image_size: Any | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Prepare video latents for action modes with mode-specific conditioning.

        Policy and forward-dynamics modes condition only latent frame zero.
        The Wan VAE is temporally causal, so encoding only the first pixel
        frame preserves that latent while avoiding unused future-frame work.
        Inverse dynamics conditions every latent and keeps the full encode.
        """
        del height, width
        if video_tensor.ndim == 4:
            video_tensor = video_tensor.unsqueeze(0)
        if video_tensor.ndim != 5 or video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
        if video_tensor.shape[2] < 1:
            raise ValueError("Cosmos3 action video tensor must contain at least one frame.")

        C = self.transformer.latent_channel_size
        T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
        latent_hw = self._latent_hw_from_image_size(image_size)
        if latent_hw is None:
            H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial
            W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial
        else:
            H_lat, W_lat = latent_hw

        noise = randn_tensor(
            (1, C, T_lat, H_lat, W_lat),
            generator=generator,
            device=self.device,
            dtype=self.sampling_dtype,
        )
        condition_indexes = vision_condition_indexes(mode, num_frames, self.vae_scale_factor_temporal)
        condition_video = video_tensor[:, :, :1] if condition_indexes == [0] else video_tensor
        cond_prefix_latent = self._encode_video_tensor(condition_video, image_size=image_size)
        expected_prefix = (1, C, max(condition_indexes) + 1, H_lat, W_lat)
        if (
            cond_prefix_latent.shape[0] != expected_prefix[0]
            or cond_prefix_latent.shape[1] != expected_prefix[1]
            or cond_prefix_latent.shape[2] < expected_prefix[2]
            or cond_prefix_latent.shape[3:] != expected_prefix[3:]
        ):
            raise ValueError(
                "Cosmos3 action video latent shape mismatch: "
                f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}."
            )

        condition_latents = torch.zeros_like(noise)
        for index in condition_indexes:
            condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1]
        condition_mask = build_vision_condition_mask(
            mode,
            num_frames,
            self.vae_scale_factor_temporal,
            device=self.device,
            dtype=self.sampling_dtype,
        )
        latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise
        velocity_mask = 1.0 - condition_mask
        return latents, velocity_mask, condition_latents

    def _prepare_action_latents(
        self,
        *,
        mode: str,
        action_chunk_size: int,
        raw_action_dim: int | None,
        generator: torch.Generator,
        sp,
        clean_action: torch.Tensor | None = None,
        condition_indexes: list[int] | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
        action_dim = int(getattr(self.transformer, "action_dim", 64))
        if clean_action is not None:
            action = clean_action.detach().to(dtype=torch.float32)
            if action.ndim == 3 and action.shape[0] == 1:
                action = action.squeeze(0)
            if action.ndim != 2:
                raise ValueError(f"Cosmos3 clean action must have shape [T, D], got {tuple(action.shape)}.")
            if action.shape[0] < action_chunk_size:
                pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1)
                action = torch.cat([action, pad], dim=0)
            elif action.shape[0] > action_chunk_size:
                action = action[:action_chunk_size]
            if raw_action_dim is None:
                raw_action_dim = int(action.shape[-1])
            clean_action = pad_action_to_dim(action, action_dim)
        elif mode == ACTION_MODE_FORWARD_DYNAMICS:
            action = load_action_tensor(self._get_sp_param(sp, "action", None))
            if action.shape[0] < action_chunk_size:
                pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1)
                action = torch.cat([action, pad], dim=0)
            elif action.shape[0] > action_chunk_size:
                action = action[:action_chunk_size]
            if raw_action_dim is None:
                raw_action_dim = int(action.shape[-1])
            clean_action = pad_action_to_dim(action, action_dim)
        else:
            if raw_action_dim is None:
                raise ValueError(
                    "Cosmos3 action_mode='policy' and 'inverse_dynamics' require extra_args['raw_action_dim']."
                )
            clean_action = torch.zeros(action_chunk_size, action_dim, dtype=torch.float32)

        raw_action_dim = int(raw_action_dim)
        if raw_action_dim <= 0 or raw_action_dim > action_dim:
            raise ValueError(f"Cosmos3 raw_action_dim must be in [1, {action_dim}], got {raw_action_dim}.")

        clean_action = clean_action.to(device=self.device, dtype=self.sampling_dtype).unsqueeze(0)
        if condition_indexes is None:
            condition_mask = build_action_condition_mask(
                mode,
                action_chunk_size,
                device=self.device,
                dtype=self.sampling_dtype,
            )
        else:
            condition_mask = self._build_action_condition_mask_from_indexes(
                condition_indexes,
                action_chunk_size,
                device=self.device,
                dtype=self.sampling_dtype,
            )
        noise = randn_tensor(
            (1, action_chunk_size, action_dim),
            generator=generator,
            device=self.device,
            dtype=self.sampling_dtype,
        )
        noise[:, :, raw_action_dim:] = 0
        clean_action[:, :, raw_action_dim:] = 0
        action_latents = condition_mask * clean_action + (1.0 - condition_mask) * noise
        action_velocity_mask = 1.0 - condition_mask
        return action_latents, action_velocity_mask, clean_action, raw_action_dim

    # -- Session-memory bypass for UND text K/V (RFC #4480) -----------------
    # Mirrors DreamZero's `_kv_*` pattern: route through the BDE pool when set
    # (Stage C/D), else the session adapter (Stage B), else the bespoke
    # transformer-instance cache. All gated so the default path is unchanged.

    def _new_cosmos3_state(self, session_id: str | None) -> Cosmos3StateAdapter | None:
        if getattr(self, "_memory_manager", None) is None:
            return None
        return Cosmos3StateAdapter(session_id, self._memory_manager)

    def _kv_load_und(self, state: Cosmos3StateAdapter | None, is_negative: bool) -> bool:
        """Load this branch's UND K/V onto the transformer before forward.

        Returns True if the session/BDE path handled the load (caller skips the
        bespoke ``cond_cache`` assignment), False for the bespoke path.
        """
        if getattr(self, "_bde_kv_state", None) is not None:
            raise NotImplementedError("BDE pool path for Cosmos3 UND K/V is Stage C/D")
        if state is not None:
            state.load_into_transformer(self.transformer, is_negative)
            return True
        return False

    def _kv_capture_und(self, state: Cosmos3StateAdapter | None, is_negative: bool) -> None:
        """After forward, store the freshly computed UND K/V for this branch (encode-once)."""
        if getattr(self, "_bde_kv_state", None) is not None:
            raise NotImplementedError("BDE pool path for Cosmos3 UND K/V is Stage C/D")
        if state is not None:
            state.capture_from_transformer(self.transformer, is_negative)

    def _kv_reset_und(self, state: Cosmos3StateAdapter | None) -> None:
        """Reset UND K/V at the start of a generation (replaces ``transformer.reset_cache()``).

        Always clears the transformer-instance cache so no stale UND K/V leaks into
        paths that read it directly (cfg-parallel / no-CFG / bespoke fallback); also
        clears the session objects when session state is active.
        """
        if getattr(self, "_bde_kv_state", None) is not None:
            raise NotImplementedError("BDE pool path for Cosmos3 UND K/V is Stage C/D")
        self.transformer.reset_cache()
        if state is not None:
            state.reset()

    # -- Denoising loop (shared by T2V and I2V) -----------------------------

    def diffuse(
        self,
        latents: torch.Tensor,
        timesteps: torch.Tensor,
        cond_ids: torch.Tensor,
        cond_mask: torch.Tensor,
        uncond_ids: torch.Tensor,
        uncond_mask: torch.Tensor,
        guidance_scale: float,
        shared_kwargs: dict,
        *,
        action_latents: torch.Tensor | None = None,
        action_velocity_mask: torch.Tensor | None = None,
        action_condition_latents: torch.Tensor | None = None,
        sound_latents: torch.Tensor | None = None,
        velocity_mask: torch.Tensor | None = None,
        image_latent: torch.Tensor | None = None,
        condition_latents: torch.Tensor | None = None,
        guidance_interval: tuple[float, float] | None = None,
        raw_action_dim: int | None = None,
        scheduler: Any | None = None,
        session_id: str | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
        """Denoising loop with 3-mode CFG support (parallel, sequential, none).

        Cosmos3's UND pathway is text-dependent, so CFG needs separate K/V
        caches for conditional and unconditional text.

        Two modes:
          1. CFG parallel (multi-GPU): each rank handles one condition via
             predict_noise_maybe_with_cfg; caching is rank-local.
          2. Sequential CFG (single-GPU or cfg_size=1): two separate
             forward passes with explicit cache swapping.  We cannot
             batch B=2 because different text lengths would cause the
             shorter branch to attend to padding in cross-attention.

        I2V conditioning (when both arguments are supplied):
          * ``velocity_mask`` zeros frame-0 noise predictions before stepping.
          * ``image_latent`` is re-injected into frame 0 after each scheduler
            step, since UniPC's predictor-corrector update rescales the
            sample (sigma-dependent), so even zero velocity does not preserve
            frame 0.

        ``guidance_interval`` (T2I) restricts CFG to
        timesteps inside the closed interval ``[lo, hi]``.  The interval is
        compared against the raw scheduler timestep value; works for both
        the [0, 1000] discrete scale and normalized flow-matching scales.
        Outside the interval the cond/uncond delta is zeroed so all ranks
        continue to execute identical control flow (CFG-Parallel safe).
        """
        do_cfg = guidance_scale > 1.0
        cfg_parallel = self._cfg_parallel_active() and do_cfg
        step_scheduler = scheduler if scheduler is not None else self.scheduler
        # Session-keyed UND K/V (RFC #4480); None => bespoke transformer-instance cache.
        kv_state = self._new_cosmos3_state(session_id)
        self._kv_reset_und(kv_state)

        def _cfg_active_at(t: torch.Tensor) -> bool:
            if guidance_interval is None:
                return True
            t_scalar = float(t.item()) if torch.is_tensor(t) else float(t)
            lo, hi = guidance_interval
            return lo <= t_scalar <= hi

        # Joint scheduler step over multiple modalities. Safe for flow-matching schedulers
        # because the update is linear per element; revisit this if Cosmos3 adopts a
        # scheduler with cross-element dependencies (e.g. per-modality timestep).
        def _pack_joint(
            video_tensor: torch.Tensor,
            action_tensor: torch.Tensor | None = None,
            sound_tensor: torch.Tensor | None = None,
        ):
            batch = video_tensor.shape[0]
            tensors = [video_tensor]
            if action_tensor is not None:
                tensors.append(action_tensor)
            if sound_tensor is not None:
                tensors.append(sound_tensor)
            flats = [tensor.reshape(batch, -1) for tensor in tensors]
            return torch.cat(flats, dim=1), [tensor.shape for tensor in tensors], [flat.shape[1] for flat in flats]

        def _unpack_joint(
            packed: torch.Tensor,
            shapes: list[torch.Size],
            numels: list[int],
        ) -> tuple[torch.Tensor, ...]:
            outputs = []
            offset = 0
            for shape, numel in zip(shapes, numels, strict=True):
                outputs.append(packed[:, offset : offset + numel].reshape(shape))
                offset += numel
            return tuple(outputs)

        def _split_noise_pred(
            noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
        ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
            has_action = action_latents is not None
            has_sound = sound_latents is not None
            if not has_action and not has_sound:
                if isinstance(noise_pred, tuple):
                    raise ValueError("Cosmos3 video-only diffusion received tuple predictions.")
                return noise_pred, None, None
            if not isinstance(noise_pred, tuple):
                raise ValueError("Cosmos3 multimodal diffusion expects transformer predictions as a tuple.")
            expected = 1 + int(has_action) + int(has_sound)
            if len(noise_pred) != expected:
                raise ValueError(
                    f"Cosmos3 multimodal diffusion expected {expected} predictions, got {len(noise_pred)}."
                )
            video_pred = noise_pred[0]
            idx = 1
            action_pred = noise_pred[idx] if has_action else None
            if has_action:
                idx += 1
            sound_pred = noise_pred[idx] if has_sound else None
            return video_pred, action_pred, sound_pred

        def _step(
            noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
            t: torch.Tensor,
            latents: torch.Tensor,
            action_latents: torch.Tensor | None,
            sound_latents: torch.Tensor | None,
        ) -> torch.Tensor | tuple[torch.Tensor, ...]:
            video_pred, action_pred, sound_pred = _split_noise_pred(noise_pred)
            if velocity_mask is not None:
                if image_latent is not None and condition_latents is None:
                    video_pred[:, :, 0:1, :, :] = 0
                else:
                    video_pred = video_pred * velocity_mask
            if action_pred is not None and action_velocity_mask is not None:
                action_pred = action_pred * action_velocity_mask
                if raw_action_dim is not None and 0 < raw_action_dim < action_pred.shape[-1]:
                    action_pred[..., raw_action_dim:] = 0
            if action_latents is None and sound_latents is None:
                latents = step_scheduler.step(
                    video_pred,
                    t,
                    latents,
                    generator=generator,
                    return_dict=False,
                )[0]
            else:
                packed_noise, shapes, numels = _pack_joint(video_pred, action_pred, sound_pred)
                packed_latents, _, _ = _pack_joint(latents, action_latents, sound_latents)
                packed_next = step_scheduler.step(
                    packed_noise,
                    t,
                    packed_latents,
                    generator=generator,
                    return_dict=False,
                )[0]
                unpacked = _unpack_joint(packed_next, shapes, numels)
                latents = unpacked[0]
                idx = 1
                if action_latents is not None:
                    action_latents = unpacked[idx]
                    idx += 1
                if sound_latents is not None:
                    sound_latents = unpacked[idx]
            if condition_latents is not None and velocity_mask is not None:
                latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
            elif image_latent is not None:
                latents[:, :, 0:1, :, :] = image_latent
            if action_latents is not None and action_condition_latents is not None and action_velocity_mask is not None:
                action_latents = (
                    action_velocity_mask * action_latents + (1.0 - action_velocity_mask) * action_condition_latents
                )
            outputs = [latents]
            if action_latents is not None:
                outputs.append(action_latents)
            if sound_latents is not None:
                outputs.append(sound_latents)
            return outputs[0] if len(outputs) == 1 else tuple(outputs)

        def _assign_step_out(step_out: torch.Tensor | tuple[torch.Tensor, ...]) -> None:
            nonlocal latents, action_latents, sound_latents
            if action_latents is None and sound_latents is None:
                assert isinstance(step_out, torch.Tensor)
                latents = step_out
                return
            if not isinstance(step_out, tuple):
                raise ValueError("Cosmos3 multimodal diffusion step returned a non-tuple result.")
            latents = step_out[0]
            idx = 1
            if action_latents is not None:
                action_latents = step_out[idx]
                idx += 1
            if sound_latents is not None:
                sound_latents = step_out[idx]

        try:
            if cfg_parallel:
                # Each CFG-parallel rank runs exactly one branch (rank 0 -> cond,
                # else uncond), so session keying loads/stores only this rank's branch.
                cfg_rank_is_negative = get_classifier_free_guidance_rank() != 0
                for step_index, t in enumerate(self.progress_bar(timesteps)):
                    self._set_denoise_step_metadata(step_index, timesteps, step_scheduler)
                    self._set_mixed_precision_step(step_index, len(timesteps))
                    timestep = t.unsqueeze(0)
                    # Outside the interval, scale=1 makes the combined output equal
                    # the cond branch. Every rank remains on the same iteration and
                    # collective schedule.
                    step_scale = guidance_scale if _cfg_active_at(t) else 1.0
                    self._kv_load_und(kv_state, is_negative=cfg_rank_is_negative)
                    noise_pred = self.predict_noise_maybe_with_cfg(
                        do_true_cfg=True,
                        true_cfg_scale=step_scale,
                        positive_kwargs=dict(
                            _cache_context="cond",
                            hidden_states=latents,
                            timestep=timestep,
                            text_ids=cond_ids,
                            text_mask=cond_mask,
                            action_latents=action_latents,
                            sound_latents=sound_latents,
                            **shared_kwargs,
                        ),
                        negative_kwargs=dict(
                            _cache_context="uncond",
                            hidden_states=latents,
                            timestep=timestep,
                            text_ids=uncond_ids,
                            text_mask=uncond_mask,
                            action_latents=action_latents,
                            sound_latents=sound_latents,
                            **shared_kwargs,
                        ),
                        cfg_normalize=False,
                    )
                    if kv_state is not None:
                        self._kv_capture_und(kv_state, is_negative=cfg_rank_is_negative)
                    _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents))

            elif do_cfg:
                cond_cache: tuple = (None, None)
                uncond_cache: tuple = (None, None)
                keep_uncond_for_cache = self._cache_requires_paired_cfg()

                for step_index, t in enumerate(self.progress_bar(timesteps)):
                    self._set_denoise_step_metadata(step_index, timesteps, step_scheduler)
                    self._set_mixed_precision_step(step_index, len(timesteps))
                    timestep = t.unsqueeze(0)
                    cfg_active = _cfg_active_at(t)

                    if not self._kv_load_und(kv_state, is_negative=False):
                        self.transformer.cached_kv, self.transformer.cached_freqs_gen = cond_cache
                    noise_cond = self.predict_noise(
                        _cache_context="cond",
                        hidden_states=latents,
                        timestep=timestep,
                        text_ids=cond_ids,
                        text_mask=cond_mask,
                        action_latents=action_latents,
                        sound_latents=sound_latents,
                        **shared_kwargs,
                    )
                    if kv_state is not None:
                        self._kv_capture_und(kv_state, is_negative=False)
                    elif cond_cache[0] is None:
                        cond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen)

                    if cfg_active or keep_uncond_for_cache:
                        if not self._kv_load_und(kv_state, is_negative=True):
                            self.transformer.cached_kv, self.transformer.cached_freqs_gen = uncond_cache
                        noise_uncond = self.predict_noise(
                            _cache_context="uncond",
                            hidden_states=latents,
                            timestep=timestep,
                            text_ids=uncond_ids,
                            text_mask=uncond_mask,
                            action_latents=action_latents,
                            sound_latents=sound_latents,
                            **shared_kwargs,
                        )
                        if kv_state is not None:
                            self._kv_capture_und(kv_state, is_negative=True)
                        elif uncond_cache[0] is None:
                            uncond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen)
                        # Outside the interval, scale=1.0 makes the combined result
                        # equal to noise_cond; the uncond pass is computed only to
                        # preserve cache-dit's cond/uncond parity.
                        step_scale = guidance_scale if cfg_active else 1.0
                        noise_pred = self.combine_cfg_noise(
                            noise_cond,
                            noise_uncond,
                            step_scale,
                            cfg_normalize=False,
                        )
                    else:
                        noise_pred = noise_cond

                    _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents))

            else:
                # No CFG: a single cond branch per step. Bespoke (state None) keeps
                # using the transformer-instance cache exactly as before.
                for step_index, t in enumerate(self.progress_bar(timesteps)):
                    self._set_denoise_step_metadata(step_index, timesteps, step_scheduler)
                    self._set_mixed_precision_step(step_index, len(timesteps))
                    timestep = t.unsqueeze(0)
                    self._kv_load_und(kv_state, is_negative=False)
                    noise_pred = self.predict_noise(
                        _cache_context="cond",
                        hidden_states=latents,
                        timestep=timestep,
                        text_ids=cond_ids,
                        text_mask=cond_mask,
                        action_latents=action_latents,
                        sound_latents=sound_latents,
                        **shared_kwargs,
                    )
                    if kv_state is not None:
                        self._kv_capture_und(kv_state, is_negative=False)
                    _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents))
        finally:
            self._clear_denoise_step_metadata()
            self._reset_mixed_precision()
            # Cosmos3 currently receives a unique request_id rather than a
            # reusable rollout session id. Retaining its state would only pin
            # K/V buffers on device after this generation finishes.
            if kv_state is not None:
                self._memory_manager.drop_session(session_id)

        outputs = [latents]
        if action_latents is not None:
            outputs.append(action_latents)
        if sound_latents is not None:
            outputs.append(sound_latents)
        return outputs[0] if len(outputs) == 1 else tuple(outputs)

    @staticmethod
    def _get_transfer_num_chunks(
        total_frames: int,
        frames_per_chunk: int,
        conditional_frames: int,
    ) -> tuple[int, int]:
        if frames_per_chunk <= 0:
            raise ValueError("Cosmos3 transfer frames_per_chunk must be positive.")
        if total_frames <= frames_per_chunk:
            return 1, frames_per_chunk
        stride = frames_per_chunk - conditional_frames
        if stride <= 0:
            raise ValueError("Cosmos3 transfer num_conditional_frames must be smaller than num_video_frames_per_chunk.")
        remaining = total_frames - frames_per_chunk
        extra_chunks = remaining // stride + (1 if remaining % stride else 0)
        return 1 + extra_chunks, stride

    def _prepare_transfer_latents(
        self,
        target_video: torch.Tensor,
        current_conditional_frames: int,
        generator: torch.Generator,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        condition_latents = self._encode_video_tensor(target_video)
        noise = randn_tensor(
            condition_latents.shape,
            generator=generator,
            device=self.device,
            dtype=self.sampling_dtype,
        )
        condition_mask = torch.zeros(
            1,
            1,
            condition_latents.shape[2],
            1,
            1,
            device=self.device,
            dtype=self.sampling_dtype,
        )
        if current_conditional_frames > 0:
            latent_frames = (current_conditional_frames - 1) // self.vae_scale_factor_temporal + 1
            condition_mask[:, :, :latent_frames] = 1.0
        latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise
        velocity_mask = 1.0 - condition_mask
        return latents, velocity_mask, condition_mask * condition_latents

    def _transfer_bucket_size(
        self,
        sp: OmniDiffusionSamplingParams,
        source_hw: tuple[int, int] | None,
    ) -> tuple[int, int, str]:
        resolution = self._get_sp_param(sp, "resolution", self._get_sp_param(sp, "image_size", 720))
        source_h, source_w = source_hw or (COSMOS3_T2V_DEFAULT_HEIGHT, COSMOS3_T2V_DEFAULT_WIDTH)
        target_w, target_h = find_closest_target_size(int(source_h), int(source_w), resolution)
        aspect_ratio = next(
            ratio for ratio, size in VIDEO_RES_SIZE_INFO[str(resolution)].items() if size == (target_w, target_h)
        )
        return int(target_h), int(target_w), aspect_ratio

    def _warn_transfer_bucket_conflicts(
        self,
        sp: OmniDiffusionSamplingParams,
        prompt: str,
        *,
        source_hw: tuple[int, int] | None,
        height: int,
        width: int,
        aspect_ratio: str,
    ) -> None:
        """Explain transfer's bucket selection when it overrides request metadata."""
        if not _is_rank_zero():
            return

        requested_size_data = self._get_sp_param(sp, COSMOS3_TRANSFER_REQUESTED_SIZE_KEY, None)
        if isinstance(requested_size_data, Mapping):
            requested_width = requested_size_data.get("width")
            requested_height = requested_size_data.get("height")
        else:
            requested_width = getattr(sp, "width", None)
            requested_height = getattr(sp, "height", None)
        if requested_width is not None and requested_height is not None:
            try:
                requested_size = (int(requested_width), int(requested_height))
            except (TypeError, ValueError):
                requested_size = None
            if requested_size is not None and requested_size != (width, height):
                source_description = (
                    f"control input {source_hw[1]}x{source_hw[0]} (WxH)"
                    if source_hw is not None
                    else "the default transfer geometry"
                )
                resolution = self._get_sp_param(sp, "resolution", self._get_sp_param(sp, "image_size", 720))
                logger.warning(
                    "Cosmos3 transfer ignores requested size=%dx%d (WxH); %s selected the %s resolution=%s "
                    "bucket, %dx%d (WxH). Transfer supports canonical Cosmos buckets, not arbitrary output sizes.",
                    *requested_size,
                    source_description,
                    aspect_ratio,
                    resolution,
                    width,
                    height,
                )

        request_aspect_ratio = self._get_sp_param(sp, "aspect_ratio", None)
        aspect_ratio_source = "requested aspect_ratio"
        if request_aspect_ratio is None:
            request_aspect_ratio = _json_object_aspect_ratio(prompt)
            aspect_ratio_source = "JSON prompt aspect_ratio"
        normalized_requested_ratio = _normalize_aspect_ratio(request_aspect_ratio)
        if normalized_requested_ratio is not None and normalized_requested_ratio != aspect_ratio:
            logger.warning(
                "Cosmos3 transfer %s=%r conflicts with the control-selected %s bucket, %dx%d (WxH); "
                "using %s for prompt metadata.",
                aspect_ratio_source,
                request_aspect_ratio,
                aspect_ratio,
                width,
                height,
                aspect_ratio,
            )

    @staticmethod
    def _first_transfer_control_hw(transfer_config: Cosmos3TransferConfig) -> tuple[int, int] | None:
        for hint in transfer_config.ordered_hints:
            if hint.control is not None:
                detected = media_hw(hint.control)
                if detected is not None:
                    return detected
            if hint.control_path is not None:
                detected = media_hw(hint.control_path)
                if detected is not None:
                    return detected
        return None

    def diffuse_transfer(
        self,
        latents: torch.Tensor,
        timesteps: torch.Tensor,
        cond_ids: torch.Tensor,
        cond_mask: torch.Tensor,
        uncond_ids: torch.Tensor,
        uncond_mask: torch.Tensor,
        guidance_scale: float,
        control_guidance: float,
        control_guidance_interval: tuple[float, float] | None,
        control_latents: list[torch.Tensor],
        shared_kwargs: dict[str, Any],
        *,
        control_weights: list[float] | None = None,
        velocity_mask: torch.Tensor,
        condition_latents: torch.Tensor,
        guidance_interval: tuple[float, float] | None = None,
        generator: torch.Generator | None = None,
    ) -> torch.Tensor:
        if getattr(self, "_use_session_state", False):
            raise NotImplementedError(
                "Cosmos3 transfer does not yet support SessionStateManager: "
                "its control/text CFG branches require richer cache keys. "
                "Disable the session state manager for transfer requests."
            )

        def _active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool:
            if interval is None:
                return True
            t_scalar = float(t.item()) if torch.is_tensor(t) else float(t)
            lo, hi = interval
            return lo <= t_scalar <= hi

        self.transformer.reset_cache()
        self._cosmos3_branch_caches = {}
        try:
            for step_index, t in enumerate(self.progress_bar(timesteps)):
                self._set_denoise_step_metadata(step_index, timesteps, self.scheduler)
                self._set_mixed_precision_step(step_index, len(timesteps))
                timestep = t.unsqueeze(0)
                step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0
                step_control = control_guidance if _active_at(t, control_guidance_interval) else 1.0
                needs_text_cfg = step_guidance > 1.0
                needs_control_cfg = step_control != 1.0

                cond_full_kwargs = dict(
                    _cache_context="cond",
                    hidden_states=latents,
                    timestep=timestep,
                    text_ids=cond_ids,
                    text_mask=cond_mask,
                    control_latents=control_latents,
                    control_weights=control_weights,
                    _cosmos3_cache_key="transfer_cond_full",
                    **shared_kwargs,
                )
                if needs_control_cfg and needs_text_cfg:
                    branches_kwargs = [
                        cond_full_kwargs,
                        dict(
                            _cache_context="cond_no_control",
                            hidden_states=latents,
                            timestep=timestep,
                            text_ids=cond_ids,
                            text_mask=cond_mask,
                            control_latents=None,
                            _cosmos3_cache_key="transfer_cond_no_control",
                            **shared_kwargs,
                        ),
                        dict(
                            _cache_context="uncond",
                            hidden_states=latents,
                            timestep=timestep,
                            text_ids=uncond_ids,
                            text_mask=uncond_mask,
                            control_latents=control_latents,
                            control_weights=control_weights,
                            _cosmos3_cache_key="transfer_uncond_full",
                            **shared_kwargs,
                        ),
                    ]
                    noise_pred = self.predict_noise_with_multi_branch_cfg(
                        do_true_cfg=True,
                        true_cfg_scale={
                            "mode": "cosmos3_transfer",
                            "branch_mode": "control_and_text",
                            "guidance_scale": step_guidance,
                            "control_guidance": step_control,
                        },
                        branches_kwargs=branches_kwargs,
                        cfg_normalize=False,
                    )
                elif needs_control_cfg:
                    branches_kwargs = [
                        cond_full_kwargs,
                        dict(
                            _cache_context="cond_no_control",
                            hidden_states=latents,
                            timestep=timestep,
                            text_ids=cond_ids,
                            text_mask=cond_mask,
                            control_latents=None,
                            _cosmos3_cache_key="transfer_cond_no_control",
                            **shared_kwargs,
                        ),
                    ]
                    noise_pred = self.predict_noise_with_multi_branch_cfg(
                        do_true_cfg=True,
                        true_cfg_scale={
                            "mode": "cosmos3_transfer",
                            "branch_mode": "control_only",
                            "control_guidance": step_control,
                        },
                        branches_kwargs=branches_kwargs,
                        cfg_normalize=False,
                    )
                elif needs_text_cfg:
                    branches_kwargs = [
                        cond_full_kwargs,
                        dict(
                            _cache_context="uncond",
                            hidden_states=latents,
                            timestep=timestep,
                            text_ids=uncond_ids,
                            text_mask=uncond_mask,
                            control_latents=control_latents,
                            control_weights=control_weights,
                            _cosmos3_cache_key="transfer_uncond_full",
                            **shared_kwargs,
                        ),
                    ]
                    noise_pred = self.predict_noise_with_multi_branch_cfg(
                        do_true_cfg=True,
                        true_cfg_scale={
                            "mode": "cosmos3_transfer",
                            "branch_mode": "text_only",
                            "guidance_scale": step_guidance,
                        },
                        branches_kwargs=branches_kwargs,
                        cfg_normalize=False,
                    )
                else:
                    noise_pred = self.predict_noise(**cond_full_kwargs)
                if isinstance(noise_pred, tuple):
                    raise ValueError("Cosmos3 transfer diffusion expects video-only tensor predictions.")
                noise_pred = noise_pred * velocity_mask
                latents = self.scheduler.step(
                    noise_pred,
                    t,
                    latents,
                    generator=generator,
                    return_dict=False,
                )[0]
                latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
        finally:
            self._clear_denoise_step_metadata()
            self._reset_mixed_precision()
            self._cosmos3_branch_caches = None
            self.transformer.reset_cache()
        return latents

    def _transfer_vae_executor(self) -> Any | None:
        """Return the active distributed VAE executor, if any."""
        executor = getattr(self.vae, "distributed_executor", None)
        is_distributed_enabled = getattr(self.vae, "is_distributed_enabled", None)
        if executor is None or not callable(is_distributed_enabled) or not is_distributed_enabled():
            return None
        return executor

    def _sync_transfer_overlap(
        self,
        output_video: torch.Tensor,
        *,
        overlap_frames: int,
        reference_video: torch.Tensor,
        vae_executor: Any | None,
    ) -> torch.Tensor | None:
        """Broadcast only the decoded frames needed to condition the next chunk."""
        if overlap_frames <= 0:
            return None

        is_output_rank = vae_executor is None or vae_executor.rank == 0
        if is_output_rank:
            overlap = output_video[:, :, -overlap_frames:].contiguous()
        else:
            overlap = torch.empty(
                (
                    reference_video.shape[0],
                    reference_video.shape[1],
                    overlap_frames,
                    reference_video.shape[3],
                    reference_video.shape[4],
                ),
                device=self.device,
                dtype=self.vae.dtype,
            )

        if vae_executor is not None:
            overlap = vae_executor.broadcast_tensor(overlap)
        return overlap

    def _forward_transfer(
        self,
        *,
        prompt: str,
        negative_prompt: str,
        sp: OmniDiffusionSamplingParams,
        transfer_config: Cosmos3TransferConfig,
        transfer_video_tensor: torch.Tensor | None,
        transfer_input_fps: float | None,
    ) -> DiffusionOutput:
        if self.is_distilled_model:
            raise self._distilled_unsupported_error("transfer requests are unsupported.")
        input_frames = None
        if transfer_video_tensor is not None:
            input_frames = normalized_video_to_uint8_cthw(transfer_video_tensor)
            source_hw = (int(input_frames.shape[-2]), int(input_frames.shape[-1]))
        else:
            source_hw = self._first_transfer_control_hw(transfer_config)
        height, width, transfer_aspect_ratio = self._transfer_bucket_size(sp, source_hw)
        self._warn_transfer_bucket_conflicts(
            sp,
            prompt,
            source_hw=source_hw,
            height=height,
            width=width,
            aspect_ratio=transfer_aspect_ratio,
        )

        if input_frames is not None:
            if tuple(input_frames.shape[-2:]) != (height, width):
                input_frames = resize_center_crop_uint8_cthw(input_frames, height, width)
            input_frames = input_frames[:, : transfer_config.max_frames]

        per_hint_frames: dict[str, torch.Tensor] = {}
        for hint in transfer_config.ordered_hints:
            frames = load_or_compute_control_frames(
                hint,
                height=height,
                width=width,
                max_frames=transfer_config.max_frames,
                input_frames=input_frames,
            )
            if frames.shape[1] < 1:
                raise ValueError(f"Cosmos3 transfer hint '{hint.key}' produced no frames.")
            per_hint_frames[hint.key] = frames
        if not per_hint_frames:
            raise ValueError("Cosmos3 transfer requires at least one control hint.")

        total_frames = next(iter(per_hint_frames.values())).shape[1]
        if transfer_config.num_frames is not None:
            total_frames = min(total_frames, int(transfer_config.num_frames))
        total_frames = max(1, total_frames)
        per_hint_frames = {key: pad_temporal_frames(frames, total_frames) for key, frames in per_hint_frames.items()}
        if input_frames is not None:
            input_frames = pad_temporal_frames(input_frames, total_frames)

        temporal_compression = self.vae_scale_factor_temporal
        chunk_frames = 1 if total_frames == 1 else transfer_config.num_video_frames_per_chunk
        chunk_frames = math.ceil((chunk_frames - 1) / temporal_compression) * temporal_compression + 1
        num_chunks, stride = self._get_transfer_num_chunks(
            total_frames,
            chunk_frames,
            transfer_config.num_conditional_frames,
        )
        padded_frames = max(total_frames, chunk_frames)
        per_hint_frames = {key: pad_temporal_frames(frames, padded_frames) for key, frames in per_hint_frames.items()}
        if input_frames is not None:
            input_frames = pad_temporal_frames(input_frames, padded_frames)

        configured_frame_rate = positive_float(transfer_config.fps)
        input_frame_rate = positive_float(transfer_input_fps)
        sampling_frame_rate = (
            positive_float(self._get_sp_param(sp, "resolved_frame_rate"))
            or positive_float(self._get_sp_param(sp, "frame_rate"))
            or positive_float(self._get_sp_param(sp, "fps"))
        )
        is_wsm_only = len(transfer_config.hints) == 1 and "wsm" in transfer_config.hints
        if is_wsm_only:
            frame_rate = configured_frame_rate or input_frame_rate or sampling_frame_rate or 24.0
        else:
            frame_rate = input_frame_rate or configured_frame_rate or sampling_frame_rate or 24.0
        num_inference_steps = sp.num_inference_steps or COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS
        guidance_scale = (
            float(transfer_config.guidance_scale)
            if transfer_config.guidance_scale is not None
            else self._resolve_guidance_scale(sp, COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE)
        )
        flow_shift_target = float(
            transfer_config.flow_shift
            if transfer_config.flow_shift is not None
            else self._get_sp_param(sp, "flow_shift", COSMOS3_V2V_DEFAULT_FLOW_SHIFT)
        )
        max_sequence_length = (
            self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH)
            or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH
        )
        transfer_prompt_suffix = None
        if transfer_config.emphasize_control_in_prompt:
            hint_names = ", ".join(hint.key for hint in transfer_config.ordered_hints)
            transfer_prompt_suffix = COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names=hint_names)

        self._guidance_scale = guidance_scale
        self._num_timesteps = num_inference_steps
        self._set_flow_shift(flow_shift_target)

        generator = sp.generator
        seed = self._resolve_seed(sp, generator)
        if generator is None:
            generator = torch.Generator(device=self.device).manual_seed(seed)

        for ignored_key in ("system_prompt", "use_system_prompt"):
            if self._get_sp_param(sp, ignored_key, None) is not None:
                logger.warning_once(
                    "Cosmos3 transfer ignores the '%s' request parameter and always "
                    "tokenizes both CFG branches with the transfer-specific system prompt.",
                    ignored_key,
                )
        cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts(
            prompt,
            negative_prompt,
            chunk_frames,
            frame_rate,
            height,
            width,
            max_sequence_length,
            sp,
            use_system_prompt=True,
            is_t2i=False,
            system_prompt=COSMOS3_TRANSFER_SYSTEM_PROMPT,
            prompt_suffix=transfer_prompt_suffix,
            use_duration_template=bool(self._get_sp_param(sp, "use_duration_template", True)),
            use_resolution_template=bool(self._get_sp_param(sp, "use_resolution_template", True)),
            negative_metadata_mode=str(self._get_sp_param(sp, "negative_metadata_mode", "same")),
            aspect_ratio_override=transfer_aspect_ratio,
        )

        output_chunks: list[torch.Tensor] = []
        control_chunks_per_hint: dict[str, list[torch.Tensor]] = {key: [] for key in per_hint_frames}
        previous_output: torch.Tensor | None = None
        vae_executor = self._transfer_vae_executor()
        is_output_rank = vae_executor is None or vae_executor.rank == 0

        for chunk_id in range(num_chunks):
            start_frame = chunk_id * stride
            end_frame = min(start_frame + chunk_frames, total_frames)
            control_norms = {
                key: uint8_cthw_to_normalized_5d(
                    pad_temporal_frames(frames[:, start_frame:end_frame], chunk_frames),
                    dtype=self.dtype,
                )
                for key, frames in per_hint_frames.items()
            }
            target_norm = torch.zeros_like(next(iter(control_norms.values())))
            current_conditional_frames = 0

            if chunk_id == 0 and transfer_config.num_first_chunk_conditional_frames > 0:
                if input_frames is None:
                    raise ValueError("Cosmos3 transfer num_first_chunk_conditional_frames > 0 requires a video input.")
                current_conditional_frames = min(
                    transfer_config.num_first_chunk_conditional_frames,
                    input_frames.shape[1],
                    chunk_frames,
                )
                if current_conditional_frames > 0:
                    input_cond = uint8_cthw_to_normalized_5d(
                        input_frames[:, :current_conditional_frames],
                        dtype=self.dtype,
                    )
                    target_norm[:, :, :current_conditional_frames] = input_cond
                    if current_conditional_frames < chunk_frames:
                        fill = target_norm[:, :, current_conditional_frames - 1 : current_conditional_frames]
                        target_norm[:, :, current_conditional_frames:] = fill.expand(
                            -1,
                            -1,
                            chunk_frames - current_conditional_frames,
                            -1,
                            -1,
                        )
            elif chunk_id > 0 and previous_output is not None:
                current_conditional_frames = min(
                    transfer_config.num_conditional_frames,
                    previous_output.shape[2],
                    chunk_frames,
                )
                if current_conditional_frames > 0:
                    target_norm[:, :, :current_conditional_frames] = previous_output[
                        :, :, -current_conditional_frames:
                    ].to(target_norm)
                    if current_conditional_frames < chunk_frames:
                        fill = target_norm[:, :, current_conditional_frames - 1 : current_conditional_frames]
                        target_norm[:, :, current_conditional_frames:] = fill.expand(
                            -1,
                            -1,
                            chunk_frames - current_conditional_frames,
                            -1,
                            -1,
                        )

            control_latents = [self._encode_video_tensor(video) for video in control_norms.values()]
            latents, velocity_mask, condition_latents = self._prepare_transfer_latents(
                target_norm,
                current_conditional_frames,
                generator,
            )
            video_shape = (latents.shape[2], latents.shape[3], latents.shape[4])
            shared_kwargs = dict(
                video_shape=video_shape,
                fps=frame_rate,
                noisy_frame_mask=velocity_mask,
                transfer_share_vision_temporal_positions=transfer_config.share_vision_temporal_positions,
            )

            self._set_timesteps(
                num_inference_steps,
                device=self.device,
                shift=self._current_flow_shift,
            )
            latents = self.diffuse_transfer(
                latents=latents,
                timesteps=self.scheduler.timesteps,
                cond_ids=cond_ids,
                cond_mask=cond_mask,
                uncond_ids=uncond_ids,
                uncond_mask=uncond_mask,
                guidance_scale=guidance_scale,
                control_guidance=transfer_config.control_guidance,
                control_guidance_interval=transfer_config.control_guidance_interval,
                control_latents=control_latents,
                control_weights=transfer_config.normalized_control_weights,
                shared_kwargs=shared_kwargs,
                velocity_mask=velocity_mask,
                condition_latents=condition_latents,
                generator=generator,
            )
            output_video = self._decode_latents(latents).clamp(-1, 1)
            if chunk_id + 1 < num_chunks:
                previous_output = self._sync_transfer_overlap(
                    output_video,
                    overlap_frames=min(transfer_config.num_conditional_frames, chunk_frames),
                    reference_video=target_norm,
                    vae_executor=vae_executor,
                )

            if is_output_rank:
                if chunk_id == 0:
                    output_chunks.append(output_video)
                    for key, control in control_norms.items():
                        control_chunks_per_hint[key].append(control)
                else:
                    output_chunks.append(output_video[:, :, current_conditional_frames:])
                    for key, control in control_norms.items():
                        control_chunks_per_hint[key].append(control[:, :, current_conditional_frames:])

        if not is_output_rank:
            return DiffusionOutput(
                output={
                    "payload": {"video": output_video},
                    "metadata": {"video": {"fps": frame_rate}},
                },
                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
            )

        full_output = torch.cat(output_chunks, dim=2)[:, :, :total_frames]
        full_controls = {
            key: torch.cat(chunks, dim=2)[:, :, :total_frames] for key, chunks in control_chunks_per_hint.items()
        }

        if transfer_config.show_control_condition:
            all_controls = torch.cat([full_controls[key] for key in per_hint_frames], dim=-1)
            all_controls = all_controls.to(full_output)
            full_output = torch.cat([all_controls, full_output], dim=-1)
        if transfer_config.show_input and input_frames is not None:
            normalized_input = uint8_cthw_to_normalized_5d(input_frames[:, :total_frames], dtype=torch.float32)
            full_output = torch.cat([normalized_input.to(full_output), full_output], dim=-1)

        return DiffusionOutput(
            output={
                "payload": {"video": full_output},
                "metadata": {
                    "transfer": {
                        "controls": full_controls,
                        "hints": list(per_hint_frames),
                    },
                    "video": {"fps": frame_rate},
                },
            },
            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
        )

    # -- Forward (main generation entry point) -------------------------------

    def forward(
        self,
        req: DiffusionRequestBatch,
    ) -> DiffusionOutput:
        pipeline_start = time.time()

        # --- Parse request ---
        prompt_data = req.prompts[0] if req.prompts else ""
        if len(req.prompts) > 1:
            raise ValueError("Cosmos3OmniDiffusersPipeline currently supports a single prompt per request.")

        sp = req.sampling_params
        robolab_inputs = self._build_robolab_policy_inputs(sp, prompt_data, getattr(req, "request_id", None))
        if robolab_inputs is not None:
            return self._forward_robolab_policy(
                sp,
                robolab_inputs,
                pipeline_start,
                session_id=getattr(req, "request_id", None),
            )

        if isinstance(prompt_data, str):
            prompt = prompt_data
            negative_prompt = None
            image_tensor = None
            video_tensor = None
            transfer_video_tensor = None
            transfer_input_fps = None
        else:
            prompt = prompt_data.get("prompt", "")
            negative_prompt = prompt_data.get("negative_prompt")
            additional_info = prompt_data.get("additional_information", {}) or {}
            image_tensor = additional_info.get("preprocessed_image")
            video_tensor = additional_info.get("preprocessed_video")
            transfer_video_tensor = additional_info.get("preprocessed_transfer_video")
            transfer_input_fps = positive_float(additional_info.get("transfer_input_fps"))

        is_t2i = self._is_t2i_request(req)
        sound_enabled = self._is_sound_request(prompt_data, sp)
        action_mode = self._get_action_mode(prompt_data, sp)
        action_enabled = action_mode is not None
        transfer_config = resolve_transfer_config(sp, prompt_data)
        action_video_tensor = video_tensor if action_enabled else None
        is_v2v = video_tensor is not None and not is_t2i and not action_enabled
        self._validate_distilled_generation_mode(
            is_t2i=is_t2i,
            image_tensor=image_tensor,
            action_enabled=action_enabled,
            transfer_config=transfer_config,
            is_v2v=is_v2v,
            sound_enabled=sound_enabled,
        )
        self._validate_edge_generation_mode(
            transfer_config=transfer_config,
            is_v2v=is_v2v,
            sound_enabled=sound_enabled,
        )
        if transfer_config is not None:
            if is_t2i:
                raise ValueError("Cosmos3 transfer inference is supported only for video outputs.")
            if action_enabled:
                raise ValueError("Cosmos3 transfer inference cannot be combined with action generation.")
            if sound_enabled:
                raise ValueError("Cosmos3 transfer inference cannot be combined with sound generation.")
        if action_enabled and is_t2i:
            raise ValueError("Cosmos3 action generation is supported only for video outputs.")
        if action_enabled and sound_enabled:
            raise ValueError("Cosmos3 action+sound joint generation is not supported in this phase.")
        if action_enabled and not getattr(self.transformer, "action_gen", False):
            raise ValueError(
                "Cosmos3 action generation was requested, but the transformer was "
                "initialized without action modules. Check that the checkpoint config "
                "enables action_gen and includes action weights."
            )
        if sound_enabled and is_t2i:
            raise ValueError(
                "Cosmos3 sound generation is supported only for video outputs in "
                "this phase; text-to-image with sound is unsupported."
            )
        if sound_enabled and not getattr(self.transformer, "sound_gen", False):
            raise ValueError(
                "Cosmos3 sound generation was requested, but the transformer was "
                "initialized without sound modules. Check that the checkpoint config "
                "enables sound_gen or defines sound_dim and includes sound weights."
            )
        if negative_prompt is None:
            negative_prompt = ""
        if transfer_config is not None:
            if image_tensor is not None:
                raise ValueError("Cosmos3 transfer inference accepts video inputs or control_path values, not images.")
            if transfer_video_tensor is None and video_tensor is not None:
                transfer_video_tensor = video_tensor
            return self._forward_transfer(
                prompt=prompt,
                negative_prompt=negative_prompt,
                sp=sp,
                transfer_config=transfer_config,
                transfer_video_tensor=transfer_video_tensor,
                transfer_input_fps=transfer_input_fps,
            )
        if image_tensor is not None and video_tensor is not None and not action_enabled:
            raise ValueError("Cosmos3 non-action generation accepts either image or video input, not both.")
        if video_tensor is not None and is_t2i:
            raise ValueError("Cosmos3 video-to-video generation is supported only for video outputs.")

        # T2I and T2V share the same model and forward path; their defaults are:
        #   T2I: regular 1024x1024, Edge 640x640; 50 steps, shift=3.0,
        #        guidance_interval=[400, 1000]
        #   T2V: 189 frames and 35 steps; regular guidance=6, shift=10;
        #        Edge guidance=5, shift=3;
        #        no guidance interval
        if is_t2i:
            height = sp.height or (
                COSMOS3_EDGE_T2I_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2I_DEFAULT_HEIGHT
            )
            width = sp.width or (COSMOS3_EDGE_T2I_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2I_DEFAULT_WIDTH)
            num_frames = 1
            num_inference_steps = sp.num_inference_steps or COSMOS3_T2I_DEFAULT_NUM_INFERENCE_STEPS
            guidance_scale = self._resolve_guidance_scale(sp, COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE)
            default_flow_shift = COSMOS3_T2I_DEFAULT_FLOW_SHIFT
            default_guidance_interval: tuple[float, float] | None = COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL
            batch_size = max(1, int(sp.num_outputs_per_prompt or 1))
        else:
            height = sp.height or (
                COSMOS3_EDGE_T2V_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2V_DEFAULT_HEIGHT
            )
            width = sp.width or (COSMOS3_EDGE_T2V_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2V_DEFAULT_WIDTH)
            default_guidance_scale = (
                COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE if self.is_edge_model else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE
            )
            num_frames = sp.num_frames or COSMOS3_T2V_DEFAULT_NUM_FRAMES
            num_inference_steps = sp.num_inference_steps or COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS
            guidance_scale = self._resolve_guidance_scale(sp, default_guidance_scale)
            default_flow_shift = self._engine_init_flow_shift
            default_guidance_interval = None
            batch_size = 1  # Existing video pipeline assumes B=1.

        if action_enabled:
            action_chunk_param = self._get_sp_param(sp, "action_chunk_size", None)
            if action_chunk_param is not None:
                action_chunk_size = int(action_chunk_param)
                if sp.num_frames is None:
                    num_frames = action_chunk_size + 1
            elif sp.num_frames is None:
                action_chunk_size = 16
                num_frames = action_chunk_size + 1
            else:
                action_chunk_size = int(num_frames) - 1
            if action_chunk_size <= 0:
                raise ValueError(f"Cosmos3 action_chunk_size must be positive, got {action_chunk_size}.")
            if num_frames not in (action_chunk_size, action_chunk_size + 1):
                raise ValueError(
                    "Cosmos3 action requests require num_frames to equal action_chunk_size "
                    f"or action_chunk_size + 1; got num_frames={num_frames}, action_chunk_size={action_chunk_size}."
                )
            num_inference_steps = sp.num_inference_steps or 30
            guidance_scale = self._resolve_guidance_scale(sp, 1.0)
            default_flow_shift = 5.0

        if not is_t2i and not action_enabled:
            requested_num_frames = int(num_frames)
            num_frames = _ceil_video_num_frames(
                requested_num_frames,
                self.vae_scale_factor_temporal,
            )
            if num_frames != requested_num_frames and _is_rank_zero():
                logger.info(
                    "Rounded Cosmos3 num_frames from %d to %d for temporal compression factor %d.",
                    requested_num_frames,
                    num_frames,
                    self.vae_scale_factor_temporal,
                )

        domain_id = None
        if action_enabled:
            domain_id = resolve_domain_id(
                domain_id=self._get_sp_param(sp, "domain_id", None),
                domain_name=self._get_sp_param(sp, "domain_name", None),
                require_explicit=True,
            )

        # Runtime controls: prefer ``extra_args`` (OpenAI endpoints write
        # there) over direct attrs.
        flow_shift_target = float(self._get_sp_param(sp, "flow_shift", default_flow_shift))
        guidance_interval = self._get_sp_param(sp, "guidance_interval", default_guidance_interval)

        frame_rate = self._get_sp_param(sp, "resolved_frame_rate") or self._get_sp_param(sp, "frame_rate") or 24.0
        max_sequence_length = (
            self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH)
            or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH
        )
        use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", is_v2v))

        if action_enabled and action_video_tensor is None:
            extra_action_video = self._get_sp_param(sp, "action_video", None)
            if isinstance(extra_action_video, torch.Tensor):
                action_video_tensor = extra_action_video
        if action_enabled and isinstance(action_video_tensor, torch.Tensor):
            if action_video_tensor.ndim == 4:
                action_video_tensor = action_video_tensor.unsqueeze(0)
            if action_video_tensor.ndim != 5:
                raise ValueError(
                    "Cosmos3 extra_args['action_video'] must have shape [1, 3, T, H, W] "
                    f"or [3, T, H, W], got {tuple(action_video_tensor.shape)}."
                )
            if sp.height is None:
                height = int(action_video_tensor.shape[-2])
            if sp.width is None:
                width = int(action_video_tensor.shape[-1])

        self._guidance_scale = guidance_scale
        self._num_timesteps = num_inference_steps

        # Always resolve to a concrete target shift for this request.
        self._set_flow_shift(flow_shift_target)

        generator = sp.generator
        seed = self._resolve_seed(sp, generator)
        if generator is None:
            generator = torch.Generator(device=self.device).manual_seed(seed)

        # --- Format prompts & tokenize (B=1; reused across loop iterations
        # for T2I num_outputs_per_prompt > 1) ---
        cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts(
            prompt,
            negative_prompt,
            num_frames,
            frame_rate,
            height,
            width,
            max_sequence_length,
            sp,
            use_system_prompt,
            is_t2i=is_t2i,
        )

        # --- Prepare latents (T2I, T2V, or I2V) ---
        # T2I shares _prepare_latents with T2V; the math collapses cleanly
        # at num_frames=1 ((1-1)//4 + 1 = 1 latent frame).  For T2I with
        # ``num_outputs_per_prompt > 1`` we loop the diffusion below;
        # batching B=N together would require expanding text K/V (UND
        # pathway is text-only and cached) and is left as a future
        # optimization.
        action_latents = None
        action_velocity_mask = None
        action_condition_latents = None
        raw_action_dim = None
        action_offset = 1
        if action_enabled:
            if action_video_tensor is not None and action_video_tensor.ndim == 4:
                action_video_tensor = action_video_tensor.unsqueeze(0)
            if action_video_tensor is not None and action_video_tensor.ndim != 5:
                raise ValueError(
                    "Cosmos3 action video tensor must have shape [1, 3, T, H, W] "
                    f"or [3, T, H, W], got {tuple(action_video_tensor.shape)}."
                )
            if action_video_tensor is not None and action_video_tensor.shape[2] < num_frames:
                pad = action_video_tensor[:, :, -1:].repeat(1, 1, num_frames - action_video_tensor.shape[2], 1, 1)
                action_video_tensor = torch.cat([action_video_tensor, pad], dim=2)
            elif action_video_tensor is not None and action_video_tensor.shape[2] > num_frames:
                action_video_tensor = action_video_tensor[:, :, :num_frames]

            if action_mode == ACTION_MODE_INVERSE_DYNAMICS and action_video_tensor is None:
                raise ValueError("Cosmos3 inverse_dynamics action mode requires multi_modal_data['video'].")
            if action_mode in {ACTION_MODE_POLICY, ACTION_MODE_FORWARD_DYNAMICS} and image_tensor is None:
                if action_video_tensor is None:
                    raise ValueError(
                        f"Cosmos3 action_mode={action_mode!r} requires multi_modal_data['image'] "
                        "or multi_modal_data['video']."
                    )
                image_tensor = action_video_tensor[:, :, 0]

            raw_action_dim_param = self._get_sp_param(sp, "raw_action_dim", None)
            raw_action_dim = int(raw_action_dim_param) if raw_action_dim_param is not None else None
            clean_action = None
            action_condition_indexes = None
            action_prepared = self._prepare_action_latents(
                mode=action_mode,
                action_chunk_size=action_chunk_size,
                raw_action_dim=raw_action_dim,
                generator=generator,
                sp=sp,
                clean_action=clean_action,
                condition_indexes=action_condition_indexes,
            )
            action_latents, action_velocity_mask, action_condition_latents, raw_action_dim = action_prepared
            action_offset = action_start_frame_offset(action_mode, action_chunk_size, num_frames)

        if action_enabled and action_video_tensor is not None:
            latents, velocity_mask, condition_latents = self._prepare_latents_action_video(
                action_video_tensor,
                action_mode,
                height,
                width,
                num_frames,
                generator,
            )
            image_latent = condition_latents[:, :, 0:1]
        elif is_v2v:
            condition_frame_indexes_vision = normalize_condition_frame_indexes_vision(
                self._get_sp_param(
                    sp,
                    "condition_frame_indexes_vision",
                    self._get_prompt_param(prompt_data, "condition_frame_indexes_vision", None),
                )
            )
            latents, velocity_mask, condition_latents = self._prepare_latents_v2v(
                video_tensor,
                height,
                width,
                num_frames,
                generator,
                condition_frame_indexes_vision,
            )
            image_latent = None
        elif image_tensor is not None and not is_t2i:
            latents, velocity_mask, image_latent = self._prepare_latents_i2v(
                image_tensor,
                height,
                width,
                num_frames,
                generator,
            )
            condition_latents = None
        else:
            latents = self._prepare_latents(height, width, num_frames, generator)
            velocity_mask = None
            image_latent = None
            condition_latents = None

        T_latent = latents.shape[2]
        H_latent = latents.shape[3]
        W_latent = latents.shape[4]
        video_shape = (T_latent, H_latent, W_latent)

        sound_latents = None
        target_audio_samples = None
        sound_sample_rate = None
        if sound_enabled:
            target_audio_samples, _, sound_sample_rate = self._resolve_sound_target_samples(sp, num_frames, frame_rate)
            sound_latents, _ = self._prepare_sound_latents(
                target_audio_samples,
                generator,
                sp_video_shape=video_shape,
            )

        # --- Denoising loop ---
        shared_kwargs = dict(video_shape=video_shape, fps=frame_rate)
        if velocity_mask is not None:
            shared_kwargs["noisy_frame_mask"] = velocity_mask
        if action_enabled:
            shared_kwargs.update(
                action_domain_ids=torch.tensor([domain_id], dtype=torch.long, device=self.device),
                action_noisy_mask=action_velocity_mask,
                action_start_frame_offset=action_offset,
                action_fps=float(self._get_sp_param(sp, "action_fps", frame_rate) or frame_rate),
            )

        def _run_diffusion(start_latents):
            self._set_timesteps(
                num_inference_steps,
                device=self.device,
                shift=self._current_flow_shift,
            )
            scheduler = self.scheduler
            return self.diffuse(
                latents=start_latents,
                timesteps=scheduler.timesteps,
                cond_ids=cond_ids,
                cond_mask=cond_mask,
                uncond_ids=uncond_ids,
                uncond_mask=uncond_mask,
                guidance_scale=guidance_scale,
                shared_kwargs=shared_kwargs,
                action_latents=action_latents,
                action_velocity_mask=action_velocity_mask,
                action_condition_latents=action_condition_latents,
                sound_latents=sound_latents,
                velocity_mask=velocity_mask,
                image_latent=image_latent,
                condition_latents=condition_latents,
                guidance_interval=guidance_interval,
                raw_action_dim=raw_action_dim,
                scheduler=scheduler,
                session_id=getattr(req, "request_id", None),
                generator=generator,
            )

        if is_t2i and batch_size > 1:
            # Generate N independent images by re-running the full diffusion
            # loop with different noise seeds.  The first sample reuses
            # ``latents`` already drawn from ``generator``; subsequent
            # samples draw fresh noise from the same generator (state
            # advances per call), giving distinct outputs from a single
            # user-provided seed.  Batched B=N would be more efficient but
            # requires expanding cached UND text K/V to match.
            samples = [_run_diffusion(latents)]
            for _ in range(batch_size - 1):
                next_latents = self._prepare_latents(height, width, num_frames, generator)
                samples.append(_run_diffusion(next_latents))
            latents = torch.cat(samples, dim=0)
        else:
            diffusion_output = _run_diffusion(latents)
            if action_enabled and sound_enabled:
                latents, action_latents, sound_latents = diffusion_output
            elif action_enabled:
                latents, action_latents = diffusion_output
            elif sound_enabled:
                latents, sound_latents = diffusion_output
            else:
                latents = diffusion_output

        # --- Decode ---
        if _is_rank_zero():
            logger.info("Decoding video...")
        decode_start = time.time()
        video = self._decode_latents(latents)
        if _is_rank_zero():
            logger.info("Video decoded in %.2fs", time.time() - decode_start)
            if not sound_enabled:
                logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)

        if sound_enabled:
            if sound_latents is None or target_audio_samples is None or sound_sample_rate is None:
                raise ValueError("Cosmos3 sound generation finished without sound latents.")
            if _is_rank_zero():
                logger.info("Decoding sound...")
            sound_decode_start = time.time()
            audio = self._decode_sound_latents(sound_latents, target_audio_samples)
            if _is_rank_zero():
                logger.info("Sound tokenizer decoded in %.2fs", time.time() - sound_decode_start)
                logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start)
            return DiffusionOutput(
                output={"video": video, "audio": audio, "audio_sample_rate": sound_sample_rate},
                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
            )

        if action_enabled:
            if action_latents is None or raw_action_dim is None or domain_id is None:
                raise ValueError("Cosmos3 action generation finished without action latents.")
            action = action_latents[:, :, :raw_action_dim].detach().cpu()
            return DiffusionOutput(
                output={
                    "payload": {
                        "video": video,
                        "actions": action,
                    },
                    "metadata": {
                        "actions": {
                            "raw_action_dim": raw_action_dim,
                            "action_mode": action_mode,
                            "domain_id": domain_id,
                        },
                    },
                },
                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
            )

        return DiffusionOutput(
            output={"image": video} if is_t2i else {"video": video},
            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
        )