cosmos-framework/packages/diffusers-cosmos3/diffusers_cosmos3/pipeline.py

# Copyright 2025 The NVIDIA Team and The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import math
from pathlib import Path
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
import torchvision.transforms.functional as TF
from diffusers.models.autoencoders.autoencoder_kl_wan import AutoencoderKLWan
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.schedulers import UniPCMultistepScheduler
from diffusers.utils.torch_utils import randn_tensor
from einops import rearrange
from tqdm import tqdm
from transformers import AutoTokenizer

from diffusers_cosmos3.sequence_packing import (
    GenerationDataClean,
    SequencePlan,
    build_packed_sequence,
    build_sequence_plans_from_data_batch,
    get_all_seq,
    pack_input_sequence,
)
from diffusers_cosmos3.transformer import (
    Cosmos3OmniTransformer,
)

_SYSTEM_PROMPT_IMAGE = "You are a helpful assistant who will generate images from a give prompt."
_SYSTEM_PROMPT_VIDEO = "You are a helpful assistant who will generate videos from a give prompt."


def save_img_or_video(sample, save_fp_wo_ext, fps=24, quality=10, ffmpeg_params=None, **kwargs):
    # TODO: remove this function and use diffusers-style vidoe processor
    # However, it may cause numerical differences in the saved video, so we keep it for now for exact reproducibility of saved videos.
    import imageio
    from PIL import Image as PILImage

    assert sample.ndim == 4, "Only support 4D tensor"

    if torch.is_floating_point(sample):
        sample = sample.clamp(0, 1)
    else:
        assert sample.dtype == torch.uint8, "Only support uint8 tensor"
        sample = sample.float().div(255)

    if sample.shape[1] == 1:
        save_obj = PILImage.fromarray(
            rearrange((sample.cpu().float().numpy() * 255), "c 1 h w -> h w c").astype(np.uint8),
            mode="RGB",
        )
        save_obj.save(f"{save_fp_wo_ext}.jpg", format="JPEG", quality=85 if quality is None else quality)
    else:
        frames = rearrange((sample.cpu().float().numpy() * 255), "c t h w -> t h w c").astype(np.uint8)
        h, w = frames.shape[1], frames.shape[2]
        out_ffmpeg_params = ffmpeg_params if ffmpeg_params is not None else ["-s", f"{w}x{h}"]
        imageio.mimsave(
            f"{save_fp_wo_ext}.mp4",
            frames,
            fps=fps,
            quality=quality,
            macro_block_size=1,
            ffmpeg_params=out_ffmpeg_params,
            output_params=["-f", "mp4"],
        )


class DiffusersWan22VAE:
    """
    Drop-in replacement for Wan2pt2VAEInterface, backed by AutoencoderKLWan.

    Bridges the following interface differences:

    1. encode – AutoencoderKLWan returns AutoencoderKLOutput(latent_dist=
       DiagonalGaussianDistribution); we extract .mode() and apply the same
       (μ - mean) * inv_std normalization that WanVAE does internally.

    2. decode – AutoencoderKLWan expects un-normalized z and returns
       DecoderOutput(sample=…); we invert the normalization before calling
       decode and unwrap the result to a plain tensor.

    3. spatial/temporal_compression_factor properties – AutoencoderKLWan
       stores these as config.scale_factor_spatial / scale_factor_temporal
       and exposes spatial_compression_ratio (not *_factor).

    Note: AutoencoderKLWan._decode() clamps the output to [-1, 1];
    Wan2pt2VAEInterface does not.  The pipeline applies .clamp(0, 1) after
    decode so this difference does not affect saved videos.

    Numerical equivalence requirements (needed for bitwise-identical output
    vs Wan2pt2VAEInterface):

    - No torch.amp.autocast: Wan2pt2VAEInterface constructs WanVAE with
      is_amp=False, so the encoder/decoder run as pure bfloat16 with no
      autocast context.  Wrapping calls in autocast changes how ops such as
      F.normalize accumulate internally and breaks the match.

    - mean / inv_std must be initialised directly in `dtype` (bfloat16).
      WanVAE.__init__ does:
          self.std    = torch.tensor(std, dtype=bfloat16)
          self.scale  = [self.mean, 1.0 / self.std]  # division in bfloat16
      Computing 1/std in float32 and then casting to bfloat16 can yield
      different bit patterns, so we must perform the division in bfloat16
      from the start.
    """

    def __init__(self, vae: AutoencoderKLWan, dtype: torch.dtype = torch.bfloat16):
        self.vae = vae
        self.dtype = dtype
        # Initialise in `dtype` so 1/std is computed in bfloat16, matching WanVAE.
        mean = torch.tensor(vae.config.latents_mean, dtype=dtype)
        std = torch.tensor(vae.config.latents_std, dtype=dtype)
        self._mean = mean  # [z_dim]
        self._inv_std = 1.0 / std  # [z_dim]

    @torch.no_grad()
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """[B,3,T,H,W] -> [B,z_dim,T//4,H//16,W//16]  (normalized μ, matching Wan2pt2VAEInterface)"""
        in_dtype = x.dtype
        device = x.device
        mean = self._mean.to(device=device, dtype=self.dtype)
        inv_std = self._inv_std.to(device=device, dtype=self.dtype)
        # No autocast — mirrors WanVAE(is_amp=False), pure bfloat16 forward pass.
        raw_mu = self.vae.encode(x.to(self.dtype)).latent_dist.mode()
        normalized = (raw_mu - mean.view(1, -1, 1, 1, 1)) * inv_std.view(1, -1, 1, 1, 1)
        return normalized.to(in_dtype)

    @torch.no_grad()
    def decode(self, z: torch.Tensor) -> torch.Tensor:
        """[B,z_dim,T_lat,H_lat,W_lat] -> [B,3,T,H,W]"""
        in_dtype = z.dtype
        device = z.device
        mean = self._mean.to(device=device, dtype=self.dtype)
        inv_std = self._inv_std.to(device=device, dtype=self.dtype)
        z_raw = z.to(self.dtype) / inv_std.view(1, -1, 1, 1, 1) + mean.view(1, -1, 1, 1, 1)
        # No autocast — mirrors WanVAE(is_amp=False), pure bfloat16 forward pass.
        out = self.vae.decode(z_raw).sample
        return out.to(in_dtype)

    @property
    def spatial_compression_factor(self) -> int:
        return self.vae.config.scale_factor_spatial

    @property
    def temporal_compression_factor(self) -> int:
        return self.vae.config.scale_factor_temporal


class Cosmos3OmniDiffusersPipeline(DiffusionPipeline):
    _optional_components = ["vision_encoder"]
    model_cpu_offload_seq = "transformer"

    def __init__(
        self,
        transformer: Cosmos3OmniTransformer,
        text_tokenizer: AutoTokenizer,
        vae: AutoencoderKLWan,
        scheduler: UniPCMultistepScheduler,
        vision_encoder=None,
    ):
        super().__init__()
        self.register_modules(
            transformer=transformer,
            text_tokenizer=text_tokenizer,
            vae=vae,
            scheduler=scheduler,
            vision_encoder=vision_encoder,
        )
        # Plain attribute (not registered): registering the wrapper would cause save_pretrained to call
        # wrapper.save_pretrained(), which fails since DiffusersWan22VAE has no such method.
        self.vision_tokenizer = DiffusersWan22VAE(vae)

        self.llm_special_tokens = {
            "start_of_generation": text_tokenizer.convert_tokens_to_ids("<|vision_start|>"),
            "end_of_generation": text_tokenizer.convert_tokens_to_ids("<|vision_end|>"),
            "eos_token_id": text_tokenizer.eos_token_id,
        }

    def tokenize_caption(
        self,
        caption: str,
        is_video: bool = False,
        use_system_prompt: bool = False,
    ) -> list[int]:
        """Tokenize a text caption into token IDs using the Qwen2 chat template.
        Returns:
            List of token IDs representing the full chat-formatted caption.
        """
        conversations = []
        # Optionally prepend a system prompt that tells the model whether it is generating
        # an image or a video. This changes the conditioning context for the LLM.
        if use_system_prompt:
            _system_prompt = _SYSTEM_PROMPT_VIDEO if is_video else _SYSTEM_PROMPT_IMAGE
            conversations.append({"role": "system", "content": _system_prompt})
        conversations.append({"role": "user", "content": caption})

        tokenizer_output = self.text_tokenizer.apply_chat_template(
            conversations,
            tokenize=True,
            add_generation_prompt=True,
            add_vision_id=False,
        )
        return tokenizer_output

    def apply_timestep_embeds_to_noisy_tokens(
        self,
        packed_tokens: torch.Tensor,
        packed_timestep_embeds: torch.Tensor,
        noisy_frame_indexes: List[torch.Tensor],
        token_shapes: list[tuple[int, ...]],
    ) -> torch.Tensor:
        start_noisy_index = 0
        flattened_noisy_frame_indexes = []
        for noisy_indexes_i, token_shape_i in zip(noisy_frame_indexes, token_shapes):
            assert noisy_indexes_i.numel() <= token_shape_i[0]
            spatial_numel_i = math.prod(token_shape_i[1:])
            spatial_indexes_i = torch.arange(spatial_numel_i, device=packed_tokens.device)
            noisy_indexes_i = (noisy_indexes_i * spatial_numel_i).unsqueeze(-1).expand(-1, spatial_numel_i)
            noisy_indexes_i = noisy_indexes_i.clone() + spatial_indexes_i + start_noisy_index
            flattened_noisy_frame_indexes.append(noisy_indexes_i.flatten())
            start_noisy_index += math.prod(token_shape_i)
        flattened_noisy_frame_indexes = torch.cat(flattened_noisy_frame_indexes, dim=0)
        assert packed_tokens.dim() == 2
        assert packed_timestep_embeds.dim() == 2
        assert packed_timestep_embeds.shape[1] == packed_tokens.shape[1]
        assert packed_timestep_embeds.shape[0] <= packed_tokens.shape[0]
        assert flattened_noisy_frame_indexes.dim() == 1
        assert flattened_noisy_frame_indexes.shape[0] == packed_timestep_embeds.shape[0]
        flattened_noisy_frame_indexes = flattened_noisy_frame_indexes.unsqueeze(-1).expand(
            -1,
            packed_tokens.shape[1],
        )
        return packed_tokens.scatter_add(
            dim=0,
            index=flattened_noisy_frame_indexes,
            src=packed_timestep_embeds,
        )

    def patchify_and_pack_latents(
        self,
        latent_patch_size: int,
        latent_channel: int,
        tokens_vision: torch.Tensor,
        token_shapes_vision: List[Tuple[int, int, int]],
    ) -> tuple[torch.Tensor, List[Tuple[int, int, int]]]:
        p = latent_patch_size
        packed_latent = []
        original_latent_shapes = []
        for latent, (t, h, w) in zip(tokens_vision, token_shapes_vision):
            latent = latent.squeeze(0)  # [C,T,H,W]
            _, t_actual, h_actual, w_actual = latent.shape
            original_latent_shapes.append((t_actual, h_actual, w_actual))
            h_padded = ((h_actual + p - 1) // p) * p
            w_padded = ((w_actual + p - 1) // p) * p
            if h_padded != h_actual or w_padded != w_actual:
                padded = torch.zeros(
                    (latent_channel, t_actual, h_padded, w_padded),
                    device=latent.device,
                    dtype=latent.dtype,
                )
                padded[:, :, :h_actual, :w_actual] = latent
                latent = padded
            h_patches = h_padded // p
            w_patches = w_padded // p
            latent = latent.reshape(latent_channel, t_actual, h_patches, p, w_patches, p)
            latent = torch.einsum("cthpwq->thwpqc", latent).reshape(-1, p * p * latent_channel)
            packed_latent.append(latent)
        return torch.cat(packed_latent, dim=0), original_latent_shapes

    def unpatchify_and_unpack_latents(
        self,
        latent_patch_size: int,
        latent_channel: int,
        packed_mse_preds: torch.Tensor,
        token_shapes_vision: List[Tuple[int, int, int]],
        noisy_frame_indexes_vision: list[torch.Tensor],
        original_latent_shapes: List[Tuple[int, int, int]] | None = None,
    ) -> list[torch.Tensor]:
        p = latent_patch_size
        unpatchified_latents = []
        start_idx = 0
        for i, (t_c, h_c, w_c) in enumerate(token_shapes_vision):
            if original_latent_shapes is not None:
                t_orig, h_orig, w_orig = original_latent_shapes[i]
                h_padded = ((h_orig + p - 1) // p) * p
                w_padded = ((w_orig + p - 1) // p) * p
                h_patches = h_padded // p
                w_patches = w_padded // p
            else:
                t_orig, h_orig, w_orig = t_c, h_c * p, w_c * p
                h_patches, w_patches = h_c, w_c
            noisy_frame_indexes = noisy_frame_indexes_vision[i]
            t_n = len(noisy_frame_indexes)
            output_tensor = torch.zeros(
                (latent_channel, t_c, h_orig, w_orig),
                device=packed_mse_preds.device,
                dtype=packed_mse_preds.dtype,
            )
            num_patches = t_n * h_patches * w_patches
            if num_patches > 0:
                end_idx = start_idx + num_patches
                latent_patches = packed_mse_preds[start_idx:end_idx]
                latent_patches = latent_patches.reshape(t_n, h_patches, w_patches, p, p, latent_channel)
                latent = torch.einsum("thwpqc->cthpwq", latent_patches)
                latent = latent.reshape(latent_channel, t_n, h_patches * p, w_patches * p)
                latent = latent[:, :, :h_orig, :w_orig]
                output_tensor[:, noisy_frame_indexes] = latent
                start_idx = end_idx
            unpatchified_latents.append(output_tensor.unsqueeze(0))
        return unpatchified_latents

    def decode_vision(
        self,
        patch_latent_dim: int,
        latent_patch_size: int,
        latent_channel: int,
        packed_seq,
        last_hidden_state: torch.Tensor,
        original_latent_shapes: List[Tuple[int, int, int]] | None = None,
    ) -> list[torch.Tensor]:
        """Decode vision predictions from last_hidden_state. Returns preds_vision list."""
        vision = packed_seq.vision
        has_noisy_vision = (
            vision is not None
            and vision.tokens is not None
            and isinstance(vision.mse_loss_indexes, torch.Tensor)
            and vision.mse_loss_indexes.numel() > 0
        )
        if not has_noisy_vision:
            preds_vision = torch.zeros(
                [1, patch_latent_dim], device=last_hidden_state.device, dtype=last_hidden_state.dtype
            )
            preds_vision = self.transformer.proj_in(preds_vision)
            preds_vision = self.transformer.proj_out(preds_vision)
            if vision is not None and vision.tokens is not None:
                preds_vision_list = [torch.zeros_like(tok) for tok in vision.tokens]
                preds_vision_list[0] = preds_vision_list[0] + 0.0 * preds_vision.sum()
            else:
                preds_vision_list = [preds_vision]
        else:
            assert vision is not None
            assert isinstance(vision.mse_loss_indexes, torch.Tensor)
            assert vision.noisy_frame_indexes is not None
            preds_vision = self.transformer.proj_out(last_hidden_state[vision.mse_loss_indexes])
            preds_vision_list = self.unpatchify_and_unpack_latents(
                latent_patch_size,
                latent_channel,
                preds_vision,
                token_shapes_vision=vision.token_shapes,
                noisy_frame_indexes_vision=vision.noisy_frame_indexes,
                original_latent_shapes=original_latent_shapes,
            )
        return preds_vision_list

    def normalize_video_databatch_inplace(
        self,
        input_video_key: str,
        data_batch: dict,
        input_key: str | None = None,
        device: str = "cuda",
        dtype: torch.dtype = torch.bfloat16,
    ) -> None:
        input_key = input_video_key if input_key is None else input_key
        if input_key in data_batch:
            if data_batch.get("is_preprocessed", False) is True:
                for i in range(len(data_batch[input_key])):
                    assert torch.is_floating_point(data_batch[input_key][i])
                    assert torch.all((data_batch[input_key][i] >= -1.0001) & (data_batch[input_key][i] <= 1.0001))
            else:
                for i in range(len(data_batch[input_key])):
                    item = data_batch[input_key][i]
                    if isinstance(item, torch.Tensor):
                        item = [item]
                    assert item[0].dtype == torch.uint8
                    data_batch[input_key][i] = torch.stack(item).to(device=device, dtype=dtype) / 127.5 - 1.0
                data_batch["is_preprocessed"] = True

    def augment_image_dim_inplace(
        self,
        input_image_key: str,
        data_batch: dict,
        input_key: str | None = None,
        device: str = "cuda",
        dtype: torch.dtype = torch.bfloat16,
    ) -> None:
        input_key = input_image_key if input_key is None else input_key
        if input_key in data_batch:
            if data_batch.get("is_preprocessed", False) is True:
                for i in range(len(data_batch[input_key])):
                    assert data_batch[input_key][i].shape[2] == 1
                return
            else:
                new_image_tensor_list = []
                for i in range(len(data_batch[input_key])):
                    for img_tensor in data_batch[input_key][i]:
                        img_tensor = rearrange(img_tensor, "c h w -> 1 c 1 h w").contiguous()
                        if img_tensor.dtype == torch.uint8:
                            img_tensor = img_tensor.to(device=device, dtype=dtype) / 127.5 - 1.0
                        new_image_tensor_list.append(img_tensor)
                data_batch[input_key] = new_image_tensor_list
                data_batch["is_preprocessed"] = True

    def remove_padding_from_latent(
        self,
        spatial_compression_factor: int,
        x0_tokens_vision: list[torch.Tensor],
        frame_size: list[torch.Tensor],
    ) -> list[torch.Tensor]:
        cropped_latents = []
        for i in range(len(x0_tokens_vision)):
            fs = frame_size[i]
            if fs.dim() == 2:
                fs = fs[0]
            orig_h = int(fs[2].item())
            orig_w = int(fs[3].item())
            orig_h_latent = orig_h // spatial_compression_factor
            orig_w_latent = orig_w // spatial_compression_factor
            cropped_latents.append(x0_tokens_vision[i][:, :, :, :orig_h_latent, :orig_w_latent].contiguous())
        return cropped_latents

    def get_data_and_condition(
        self,
        input_image_key: str,
        input_video_key: str,
        data_batch: dict,
        device: str = "cuda",
        dtype: torch.dtype = torch.bfloat16,
    ) -> GenerationDataClean:
        assert (input_image_key in data_batch) != (input_video_key in data_batch)
        is_img = input_image_key in data_batch
        sample_vision_list = data_batch[input_image_key if is_img else input_video_key]

        if "num_vision_items_per_sample" not in data_batch:
            has_multiple_vision_per_sample = any(
                isinstance(v, (list, tuple)) and len(v) > 1 for v in sample_vision_list
            )
            num_vision_items_per_sample: list[int] | None = (
                [len(v) for v in sample_vision_list] if has_multiple_vision_per_sample else None
            )
            data_batch["num_vision_items_per_sample"] = num_vision_items_per_sample
            if has_multiple_vision_per_sample:
                media_key = input_video_key if not is_img else input_image_key
                data_batch[media_key] = [item.unsqueeze(0) for sublist in sample_vision_list for item in sublist]
                if data_batch[media_key][0].dtype == torch.float32 and not is_img:
                    data_batch["is_preprocessed"] = True
        else:
            num_vision_items_per_sample = data_batch["num_vision_items_per_sample"]

        batch_size = (
            len(sample_vision_list) if num_vision_items_per_sample is None else len(num_vision_items_per_sample)
        )

        self.normalize_video_databatch_inplace(input_video_key, data_batch, device=device, dtype=dtype)
        self.augment_image_dim_inplace(input_image_key, data_batch, device=device, dtype=dtype)
        raw_state_vision = data_batch[input_image_key if is_img else input_video_key]
        x0_tokens_vision = [
            self.vision_tokenizer.encode(raw_state_vision_i).contiguous().float()
            for raw_state_vision_i in raw_state_vision
        ]

        frame_size = data_batch.get("image_size", None)
        if frame_size is not None:
            x0_tokens_vision = self.remove_padding_from_latent(
                self.vision_tokenizer.spatial_compression_factor, x0_tokens_vision, frame_size
            )

        fps_raw = data_batch.get("conditioning_fps", None)
        if isinstance(fps_raw, list):
            fps_raw = torch.stack(fps_raw).flatten()
        fps_vision = fps_raw.to(device=device, dtype=dtype) if fps_raw is not None else None

        return GenerationDataClean(
            batch_size=batch_size,
            is_image_batch=is_img,
            raw_state_vision=raw_state_vision,
            x0_tokens_vision=x0_tokens_vision,
            fps_vision=fps_vision,
            num_vision_items_per_sample=num_vision_items_per_sample,
        )

    def get_inference_text_tokens(
        self, use_system_prompt: bool, input_caption_key: str, data_batch: dict, has_negative_prompt: bool
    ) -> tuple[list[list[int]], list[list[int]]]:
        cond_tokens = [
            self.tokenize_caption(c, is_video=False, use_system_prompt=use_system_prompt)
            for c in data_batch[input_caption_key]
        ]
        if has_negative_prompt:
            neg_key = "neg_" + input_caption_key
            assert neg_key in data_batch, f"Negative prompt ({neg_key}) not found"
            uncond_captions = data_batch[neg_key]
        else:
            uncond_captions = [""] * len(cond_tokens)
        uncond_tokens = [
            self.tokenize_caption(c, is_video=False, use_system_prompt=use_system_prompt) for c in uncond_captions
        ]
        return cond_tokens, uncond_tokens

    def derive_include_end_of_generation_token(self, joint_attn_implementation: str) -> bool:
        assert joint_attn_implementation in ("flex", "two_way", "three_way")
        return joint_attn_implementation == "flex"

    def prepare_inference_data(
        self,
        use_system_prompt: bool,
        prompt: Union[str, List[str]],
        negative_prompt: Optional[Union[str, List[str]]] = None,
        image=None,
        num_frames: int = 189,
        height: int = 720,
        width: int = 1280,
        fps: float = 24.0,
        condition_frame_indexes: Optional[List[int]] = None,
        noises: Optional[List[torch.Tensor]] = None,
        generator: Optional[torch.Generator] = None,
        input_caption_key: str = "ai_caption",
        input_video_key: str = "video",
        input_image_key: str = "images",
        device: str = "cuda",
        dtype: torch.dtype = torch.bfloat16,
    ) -> tuple[
        list[SequencePlan],
        GenerationDataClean,
        list[list[int]],
        list[list[int]],
        torch.Tensor,
    ]:
        # Build data_batch
        prompts = [prompt] if isinstance(prompt, str) else list(prompt)
        batch_size = len(prompts)
        is_image = num_frames == 1

        conditioning_frames = None
        if image is not None:
            conditioning_frames = self._load_image_as_tensor(image, height, width)

        image_size = [
            torch.tensor([[height, width, height, width]], dtype=torch.float32, device=device)
            for _ in range(batch_size)
        ]

        if is_image:
            img_tensor = (
                conditioning_frames.unsqueeze(0).to(device=device, dtype=dtype)
                if conditioning_frames is not None
                else torch.zeros(1, 3, 1, height, width, dtype=dtype, device=device)
            )
            seq_plans = [
                SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=[])
                for _ in range(batch_size)
            ]
            data_batch = {
                input_image_key: [img_tensor] * batch_size,
                "image_size": image_size,
                "is_preprocessed": True,
                "fps": torch.full((batch_size,), float(fps), device=device),
                "conditioning_fps": torch.full((batch_size,), float(fps), device=device),
                "num_frames": torch.full((batch_size,), num_frames, device=device),
                "sequence_plan": seq_plans,
                input_caption_key: prompts,
            }
        else:
            cond_indexes = (
                condition_frame_indexes
                if condition_frame_indexes is not None
                else ([0] if conditioning_frames is not None else [])
            )
            if conditioning_frames is not None:
                video_data = torch.zeros(1, 3, num_frames, height, width, dtype=dtype)
                t_fill = min(conditioning_frames.shape[1], num_frames)
                video_data[0, :, :t_fill] = conditioning_frames[:, :t_fill].to(dtype=dtype)
                if t_fill < num_frames:
                    video_data[0, :, t_fill:] = video_data[0, :, t_fill - 1 : t_fill].expand(
                        -1, num_frames - t_fill, -1, -1
                    )
                video_tensor = video_data.to(device=device)
            else:
                video_tensor = torch.zeros(1, 3, num_frames, height, width, dtype=dtype, device=device)
            seq_plans = [
                SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=list(cond_indexes))
                for _ in range(batch_size)
            ]
            data_batch = {
                input_video_key: [video_tensor] * batch_size,
                "image_size": image_size,
                "is_preprocessed": True,
                "fps": torch.full((batch_size,), float(fps), device=device),
                "conditioning_fps": torch.full((batch_size,), float(fps), device=device),
                "num_frames": torch.full((batch_size,), num_frames, device=device),
                "sequence_plan": seq_plans,
                input_caption_key: prompts,
            }

        has_negative_prompt = negative_prompt is not None
        if has_negative_prompt:
            neg_prompts = [negative_prompt] if isinstance(negative_prompt, str) else list(negative_prompt)
            data_batch["neg_" + input_caption_key] = neg_prompts

        sequence_plans = build_sequence_plans_from_data_batch(
            data_batch=data_batch,
            input_video_key=input_video_key,
            input_image_key=input_image_key,
        )
        gen_data_clean = self.get_data_and_condition(
            input_image_key, input_video_key, data_batch, device=device, dtype=dtype
        )
        num_items_per_sample = gen_data_clean.num_vision_items_per_sample
        cond_text_tokens, uncond_text_tokens = self.get_inference_text_tokens(
            use_system_prompt, input_caption_key, data_batch, has_negative_prompt
        )

        mask_timesteps = torch.zeros((gen_data_clean.batch_size,), dtype=torch.float32)
        packed_seq = pack_input_sequence(
            sequence_plans=sequence_plans,
            input_text_indexes=cond_text_tokens,
            gen_data_clean=gen_data_clean,
            input_timesteps=mask_timesteps,
            special_tokens=self.llm_special_tokens,
            latent_patch_size=self.transformer.config.latent_patch_size,
            include_end_of_generation_token=self.derive_include_end_of_generation_token(
                self.transformer.config.joint_attn_implementation
            ),
            position_embedding_type=self.transformer.config.position_embedding_type,
            unified_3d_mrope_reset_spatial_ids=self.transformer.config.unified_3d_mrope_reset_spatial_ids,
            unified_3d_mrope_temporal_modality_margin=self.transformer.config.unified_3d_mrope_temporal_modality_margin,
            enable_fps_modulation=self.transformer.config.enable_fps_modulation,
            base_fps=float(self.transformer.config.base_fps),
            temporal_compression_factor=self.vision_tokenizer.temporal_compression_factor,
            video_temporal_causal=self.transformer.config.video_temporal_causal,
            action_dim=self.transformer.config.max_action_dim,
        )

        assert packed_seq.vision is not None
        assert packed_seq.vision.condition_mask is not None
        assert isinstance(packed_seq.vision.condition_mask, list)
        assert gen_data_clean.x0_tokens_vision is not None

        noise_vision_list: list[torch.Tensor] = []
        for i, (x0_token, cond_mask) in enumerate(
            zip(gen_data_clean.x0_tokens_vision, packed_seq.vision.condition_mask, strict=True)
        ):
            if noises is not None:
                pure_noise = noises[i].to(device=device, dtype=dtype)
            else:
                pure_noise = randn_tensor(tuple(x0_token.shape), generator=generator, device=device, dtype=dtype)
            noise_vision_list.append(
                cond_mask * x0_token.to(device=device, dtype=dtype) + (1.0 - cond_mask) * pure_noise
            )

        initial_noise = torch.cat([t.reshape(-1) for t in noise_vision_list])

        return sequence_plans, gen_data_clean, cond_text_tokens, uncond_text_tokens, initial_noise

    def encode_text(
        self,
        hidden_size: int,
        packed_seq,
    ) -> tuple[torch.Tensor, torch.dtype]:
        """Embed text tokens. Returns (hidden_states [N_total, H], target_dtype)."""
        packed_text_embedding = self.transformer.embed_tokens(packed_seq.text_ids)  # [N_text,H]
        hidden_states = packed_text_embedding.new_zeros(size=(packed_seq.sequence_length, hidden_size))
        hidden_states[packed_seq.text_indexes] = packed_text_embedding
        return hidden_states, packed_text_embedding.dtype

    def encode_vision(
        self,
        timestep_scale: float,
        latent_patch_size: int,
        latent_channel: int,
        packed_seq,
        hidden_states: torch.Tensor,
        target_dtype: torch.dtype,
        fps: Optional[torch.Tensor] = None,
    ) -> List[Tuple[int, int, int]] | None:
        """Project vision tokens into hidden_states in-place. Returns original_latent_shapes."""
        if packed_seq.vision is None or packed_seq.vision.tokens is None:
            return None
        vision = packed_seq.vision
        assert vision.tokens is not None
        assert vision.token_shapes is not None
        assert isinstance(vision.sequence_indexes, torch.Tensor)
        assert isinstance(vision.timesteps, torch.Tensor)
        assert isinstance(vision.mse_loss_indexes, torch.Tensor)

        packed_tokens_vision, original_latent_shapes = self.patchify_and_pack_latents(
            latent_patch_size, latent_channel, vision.tokens, vision.token_shapes
        )
        packed_tokens_vision = self.transformer.proj_in(packed_tokens_vision)

        if vision.mse_loss_indexes.numel() > 0:
            timesteps_vision = vision.timesteps * timestep_scale
            with torch.autocast("cuda", enabled=True, dtype=torch.float32):
                packed_timestep_embeds_vision = self.transformer.time_embedder(timesteps_vision)
            packed_timestep_embeds_vision = packed_timestep_embeds_vision.to(target_dtype)
            packed_tokens_vision = self.apply_timestep_embeds_to_noisy_tokens(
                packed_tokens=packed_tokens_vision,
                packed_timestep_embeds=packed_timestep_embeds_vision,
                noisy_frame_indexes=vision.noisy_frame_indexes,
                token_shapes=vision.token_shapes,
            )

        hidden_states[vision.sequence_indexes] = packed_tokens_vision
        return original_latent_shapes

    @torch.no_grad()
    def run_single(
        self,
        packed_seq,
        noise_x_vision: list[torch.Tensor],
        hidden_size: int,
        latent_patch_size: int,
        latent_channel: int,
        patch_latent_dim: int,
        timestep_scale: float,
        num_heads: int,
        head_dim: int,
        num_hidden_layers: int,
        use_moe: bool,
        joint_attn_implementation: str,
        fps_vision: Optional[torch.Tensor],
        device: str = "cuda",
        dtype: torch.dtype = torch.bfloat16,
    ) -> list[torch.Tensor]:
        """Inlined forward pass from Cosmos3VFMNetworkSimple.forward()."""
        if packed_seq.vision is not None:
            packed_seq.vision.tokens = [x.to(device=device, dtype=dtype) for x in noise_x_vision]
        packed_seq.to_cuda()

        # 1. Encode text
        hidden_states, target_dtype = self.encode_text(hidden_size, packed_seq)

        # 2. Encode vision
        original_latent_shapes = self.encode_vision(
            timestep_scale,
            latent_patch_size,
            latent_channel,
            packed_seq,
            hidden_states,
            target_dtype,
            fps=fps_vision,
        )

        # 3. Build attention metadata
        assert use_moe
        assert packed_seq.attn_modes is not None
        assert packed_seq.split_lens is not None
        all_gen_indexes = []
        if packed_seq.vision is not None:
            assert packed_seq.vision.token_shapes is not None
            assert isinstance(packed_seq.vision.sequence_indexes, torch.Tensor)
            all_gen_indexes.append(packed_seq.vision.sequence_indexes)
        vision_sequence_indexes = torch.cat(all_gen_indexes, dim=0) if all_gen_indexes else None

        input_pack, attention_meta, _ = build_packed_sequence(
            joint_attn_implementation,
            packed_sequence=hidden_states,
            attn_modes=packed_seq.attn_modes,
            split_lens=packed_seq.split_lens,
            sample_lens=packed_seq.sample_lens,
            packed_und_token_indexes=packed_seq.text_indexes,
            packed_gen_token_indexes=vision_sequence_indexes,
            num_heads=num_heads,
            is_image_batch=packed_seq.is_image_batch,
            head_dim=head_dim,
            num_layers=num_hidden_layers,
            token_shapes=packed_seq.vision.token_shapes,
            natten_parameter_list=None,
            cp_world_size=1,
            video_temporal_causal=False,
            vision_token_shapes=packed_seq.vision.token_shapes if packed_seq.vision else None,
            action_token_shapes=None,
            temporal_compression_factor_vision=self.vision_tokenizer.temporal_compression_factor,
            null_action_supertokens=packed_seq.null_action_supertokens,
            pad_for_cuda_graphs=False,
        )

        # 4. Run transformer
        packed_outputs, _ = self.transformer(
            input_pack,
            attention_mask=attention_meta,
            position_ids=packed_seq.position_ids,
            dual_kv_cache=None,
            frame_idx=None,
            natten_metadata_list=None,
        )
        last_hidden_state = get_all_seq(packed_outputs)

        # 5. Decode vision
        return self.decode_vision(
            patch_latent_dim,
            latent_patch_size,
            latent_channel,
            packed_seq,
            last_hidden_state,
            original_latent_shapes,
        )

    @torch.no_grad()
    def get_cfg_velocity(
        self,
        noise_x: torch.Tensor,
        timestep: torch.Tensor,
        guidance: float,
        gen_data_clean: GenerationDataClean,
        sequence_plans: list[SequencePlan],
        cond_tokens: list[list[int]],
        uncond_tokens: list[list[int]],
        include_eog: bool,
        hidden_size: int,
        latent_patch_size: int,
        latent_channel: int,
        patch_latent_dim: int,
        timestep_scale: float,
        num_heads: int,
        head_dim: int,
        num_hidden_layers: int,
        use_moe: bool,
        joint_attn_implementation: str,
        skip_text_tokens_for_cfg: bool = False,
        normalize_cfg: bool = False,
        device: str = "cuda",
        dtype: torch.dtype = torch.bfloat16,
    ) -> torch.Tensor:
        torch.compiler.cudagraph_mark_step_begin()
        assert timestep.ndim == 2 and timestep.shape == (1, 1)

        num_items = gen_data_clean.num_vision_items_per_sample
        num_vis = num_items[0] if num_items is not None else 1

        noise_x_vision: list[torch.Tensor] = []
        offset = 0
        for j in range(num_vis):
            vision_shape = gen_data_clean.x0_tokens_vision[j].shape
            vision_dim = int(torch.prod(torch.tensor(vision_shape)))
            noise_x_vision.append(noise_x[offset : offset + vision_dim].reshape(vision_shape))
            offset += vision_dim

        gen_data_for_packing = GenerationDataClean(
            batch_size=1,
            is_image_batch=gen_data_clean.is_image_batch,
            raw_state_vision=gen_data_clean.raw_state_vision,
            x0_tokens_vision=noise_x_vision,
            fps_vision=gen_data_clean.fps_vision,
            num_vision_items_per_sample=num_items,
        )

        def _run_cond(text_tokens: list[list[int]], skip_text: bool) -> torch.Tensor:
            packed_seq = pack_input_sequence(
                sequence_plans=sequence_plans,
                input_text_indexes=text_tokens,
                gen_data_clean=gen_data_for_packing,
                input_timesteps=timestep.cpu(),
                special_tokens=self.llm_special_tokens,
                latent_patch_size=self.transformer.config.latent_patch_size,
                include_end_of_generation_token=include_eog,
                skip_text_tokens=skip_text,
                position_embedding_type=self.transformer.config.position_embedding_type,
                unified_3d_mrope_reset_spatial_ids=self.transformer.config.unified_3d_mrope_reset_spatial_ids,
                unified_3d_mrope_temporal_modality_margin=self.transformer.config.unified_3d_mrope_temporal_modality_margin,
                enable_fps_modulation=self.transformer.config.enable_fps_modulation,
                base_fps=float(self.transformer.config.base_fps),
                temporal_compression_factor=self.vision_tokenizer.temporal_compression_factor,
                video_temporal_causal=self.transformer.config.video_temporal_causal,
                action_dim=self.transformer.config.max_action_dim,
            )
            preds = self.run_single(
                packed_seq,
                noise_x_vision,
                hidden_size,
                latent_patch_size,
                latent_channel,
                patch_latent_dim,
                timestep_scale,
                num_heads,
                head_dim,
                num_hidden_layers,
                use_moe,
                joint_attn_implementation,
                fps_vision=gen_data_clean.fps_vision,
                device=device,
                dtype=dtype,
            )

            assert packed_seq.vision is not None
            assert packed_seq.vision.condition_mask is not None
            assert isinstance(packed_seq.vision.condition_mask, list)
            velocity_vision = [
                pred * (1.0 - m).to(dtype=pred.dtype, device=pred.device)
                if (1.0 - m).sum() > 0
                else torch.zeros_like(pred)
                for pred, m in zip(preds, packed_seq.vision.condition_mask)
            ]
            return torch.cat([v.reshape(-1) for v in velocity_vision])

        cond_v = _run_cond(cond_tokens, False)
        uncond_v = _run_cond(uncond_tokens, skip_text_tokens_for_cfg)

        v_pred = uncond_v + guidance * (cond_v - uncond_v)

        if normalize_cfg:
            v_pred = v_pred * (torch.norm(cond_v) / (torch.norm(v_pred) + 1e-8)).clamp(min=0.0, max=1.0)
        return v_pred

    def _load_image_as_tensor(self, image, target_h: int, target_w: int) -> torch.Tensor:
        """Load image from PIL, path, URL, or tensor; returns [3, 1, H, W] in [-1, 1]."""
        from PIL import Image as PILImage

        if isinstance(image, (str, Path)):
            image_str = str(image)
            if image_str.startswith("http://") or image_str.startswith("https://"):
                import io
                import urllib.request

                with urllib.request.urlopen(image_str) as resp:
                    img_bytes = resp.read()
                pil_img = PILImage.open(io.BytesIO(img_bytes)).convert("RGB")
            else:
                with open(image_str, "rb") as f:
                    pil_img = PILImage.open(f).convert("RGB")
            img_t = torch.from_numpy(np.array(pil_img)).permute(2, 0, 1).float()
        elif hasattr(image, "convert"):  # PIL.Image
            img_t = torch.from_numpy(np.array(image.convert("RGB"))).permute(2, 0, 1).float()
        elif isinstance(image, torch.Tensor):
            img_t = image.float()
            if img_t.dim() == 4:
                img_t = img_t.squeeze(0)
            # if already normalized to [-1, 1], skip the /127.5-1 step below
            if img_t.max() <= 1.1:
                img_4d = img_t.unsqueeze(0)
                orig_h, orig_w = img_4d.shape[2], img_4d.shape[3]
                scale = max(target_w / orig_w, target_h / orig_h)
                resize_h = int(math.ceil(scale * orig_h))
                resize_w = int(math.ceil(scale * orig_w))
                img_4d = TF.resize(img_4d, [resize_h, resize_w])
                img_4d = TF.center_crop(img_4d, [target_h, target_w])
                return img_4d.squeeze(0).unsqueeze(1)
        else:
            raise TypeError(f"Unsupported image type: {type(image)}")

        img_4d = img_t.unsqueeze(0)  # [1, 3, H, W]  (uint8-range [0, 255])
        orig_h, orig_w = img_4d.shape[2], img_4d.shape[3]
        scale = max(target_w / orig_w, target_h / orig_h)
        resize_h = int(math.ceil(scale * orig_h))
        resize_w = int(math.ceil(scale * orig_w))
        img_4d = TF.resize(img_4d, [resize_h, resize_w])
        img_4d = TF.center_crop(img_4d, [target_h, target_w])
        img_4d = img_4d / 127.5 - 1.0  # normalize after resize, matching load_conditioning_image
        return img_4d.squeeze(0).unsqueeze(1)  # [3, 1, H, W]

    def _resolve_defaults_and_prompts(
        self,
        prompt: Union[str, List[str]],
        negative_prompt: Optional[Union[str, List[str]]],
        image,
        num_frames: int,
        fps: float,
        height: int,
        width: int,
    ) -> tuple[float, int, float, Union[str, List[str]], Union[str, List[str]]]:
        """Load modality defaults and apply duration/resolution templates to prompts.

        Returns (guidance, num_steps, shift, formatted_prompt, formatted_negative_prompt).
        """
        model_mode = "image2video" if image is not None else "text2video"
        defaults = json.loads((Path(__file__).parent / f"sample_args/{model_mode}.json").read_text())

        guidance = float(defaults["guidance"])
        num_steps = int(defaults["num_steps"])
        shift = float(defaults["shift"])
        print(f"model_mode={model_mode!r}: guidance={guidance}, num_steps={num_steps}, shift={shift}")

        duration_template = defaults.get("duration_template")
        resolution_template = defaults.get("resolution_template")
        negative_prompt_base = defaults.get("negative_prompt", "")
        keep_metadata = defaults.get("negative_prompt_keep_metadata", False)

        def _apply_templates(text: str) -> str:
            if duration_template and num_frames > 1:
                text = text.rstrip(".") + ". " + duration_template.format(duration=num_frames / fps, fps=fps)
            if resolution_template:
                text = text.rstrip(".") + ". " + resolution_template.format(height=height, width=width)
            return text

        if isinstance(prompt, str):
            prompt = _apply_templates(prompt)
        else:
            prompt = [_apply_templates(p) for p in prompt]

        if negative_prompt is None:
            negative_prompt = _apply_templates(negative_prompt_base) if keep_metadata else negative_prompt_base

        return guidance, num_steps, shift, prompt, negative_prompt

    @torch.no_grad()
    def decode_latents(self, vision_list: list[torch.Tensor]) -> list[torch.Tensor]:
        """Decode latents to pixel tensors of shape [C, T, H, W] in [0, 1]."""
        frames = []
        for vision_latent in vision_list:
            vision = self.vision_tokenizer.decode(vision_latent.cuda())  # [1, C, T, H, W]
            frames.append(((1.0 + vision) / 2).clamp(0, 1).squeeze(0))
        return frames

    @torch.no_grad()
    def __call__(
        self,
        prompt: Union[str, List[str]],
        negative_prompt: Optional[Union[str, List[str]]] = None,
        image=None,
        num_frames: int = 189,
        height: int = 720,
        width: int = 1280,
        fps: float = 24.0,
        condition_frame_indexes: Optional[List[int]] = None,
        noises: Optional[List[torch.Tensor]] = None,
        generator: Optional[torch.Generator] = None,
        use_system_prompt: bool = False,
        device: str = "cuda",
        dtype: torch.dtype = torch.bfloat16,
        output_type: str = "video",
    ):
        latent_patch_size = self.transformer.config.latent_patch_size
        latent_channel = self.transformer.config.latent_channel
        patch_latent_dim = self.transformer.config.patch_latent_dim
        timestep_scale = self.transformer.config.timestep_scale
        hidden_size = self.transformer.config.hidden_size
        num_heads = self.transformer.config.num_attention_heads
        head_dim = self.transformer.config.head_dim
        num_hidden_layers = self.transformer.config.num_hidden_layers
        use_moe = self.transformer.config.use_moe
        joint_attn_implementation = self.transformer.config.joint_attn_implementation

        guidance, num_steps, shift, prompt, negative_prompt = self._resolve_defaults_and_prompts(
            prompt, negative_prompt, image, num_frames, fps, height, width
        )

        sequence_plans, gen_data_clean, cond_tokens, uncond_tokens, initial_noise = self.prepare_inference_data(
            use_system_prompt,
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=image,
            num_frames=num_frames,
            height=height,
            width=width,
            fps=fps,
            condition_frame_indexes=condition_frame_indexes,
            noises=noises,
            generator=generator,
            device=device,
            dtype=dtype,
        )

        assert guidance != 1.0, "Guidance weight must be != 1.0"

        device = initial_noise.device
        self.scheduler.set_timesteps(num_steps, device=device)
        timesteps = self.scheduler.timesteps
        # print(f"sigmas:    first={self.scheduler.sigmas[0].item():.4f}  last={self.scheduler.sigmas[-2].item():.4f}")
        # print(f"timesteps: first={timesteps[0].item():.2f}  last={timesteps[-1].item():.2f}")
        # print(f"timestep_scale: {timestep_scale}")
        # breakpoint()
        latent = initial_noise
        include_eog = self.derive_include_end_of_generation_token(joint_attn_implementation)

        # --- Denoising loop ---
        print("Running generate_samples_from_batch …")
        for timestep in tqdm(timesteps, desc="Denoising"):
            velocity_pred = self.get_cfg_velocity(
                latent,
                timestep.reshape(1, 1),
                guidance,
                gen_data_clean,
                sequence_plans,
                cond_tokens,
                uncond_tokens,
                include_eog,
                hidden_size,
                latent_patch_size,
                latent_channel,
                patch_latent_dim,
                timestep_scale,
                num_heads,
                head_dim,
                num_hidden_layers,
                use_moe,
                joint_attn_implementation,
                device=device,
                dtype=dtype,
            )
            latent = self.scheduler.step(
                model_output=velocity_pred,
                timestep=timestep,
                sample=latent.unsqueeze(0),
                return_dict=False,
            )[0].squeeze(0)

        # --- Extract vision results ---
        num_vision_items = gen_data_clean.num_vision_items_per_sample
        n_vis = num_vision_items[0] if num_vision_items is not None else 1
        result_vision: list[torch.Tensor] = []
        offset = 0
        for j in range(n_vis):
            vision_shape = gen_data_clean.x0_tokens_vision[j].shape
            vision_dim = int(torch.prod(torch.tensor(vision_shape)))
            if j == n_vis - 1:
                result_vision.append(latent[offset : offset + vision_dim].reshape(vision_shape))
            offset += vision_dim

        if output_type == "latent":
            return result_vision
        return self.decode_latents(result_vision)