# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import inspect
import json
import logging
import os
from collections.abc import Callable, Iterable
from pathlib import Path
from typing import Any, cast

import numpy as np
import torch
from diffusers.image_processor import VaeImageProcessor
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.utils.torch_utils import randn_tensor
from safetensors.torch import load_file
from torch import nn
from transformers import AutoTokenizer, Qwen3Model, T5TokenizerFast

from vllm_omni.diffusion.config import set_current_diffusion_config
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.models.anima.anima_text_conditioner import (
    ANIMA_TEXT_CONDITIONER_CONFIG,
    AnimaTextConditioner,
)
from vllm_omni.diffusion.models.anima.anima_transformer import ANIMA_TRANSFORMER_CONFIG, AnimaTransformer3DModel
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.offloader.config import OffloadStrategy, resolve_offload_strategy
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
from vllm_omni.diffusion.utils.size_utils import normalize_min_aligned_size
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch

logger = logging.getLogger(__name__)

ANIMA_VAE_SCALE_FACTOR = 8


def retrieve_timesteps(
    scheduler: Any,
    num_inference_steps: int | None = None,
    device: str | torch.device | None = None,
    timesteps: list[int] | None = None,
    sigmas: list[float] | None = None,
    **kwargs: Any,
) -> tuple[torch.Tensor, int]:
    if timesteps is not None and sigmas is not None:
        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
    if timesteps is not None:
        if "timesteps" not in inspect.signature(scheduler.set_timesteps).parameters:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                f" timestep schedules. Please check whether you are using the correct scheduler."
            )
        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
        scheduler_timesteps = scheduler.timesteps
        num_inference_steps = len(scheduler_timesteps)
    elif sigmas is not None:
        if "sigmas" not in inspect.signature(scheduler.set_timesteps).parameters:
            raise ValueError(
                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                f" sigmas schedules. Please check whether you are using the correct scheduler."
            )
        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
        scheduler_timesteps = scheduler.timesteps
        num_inference_steps = len(scheduler_timesteps)
    else:
        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
        scheduler_timesteps = scheduler.timesteps
    if num_inference_steps is None:
        num_inference_steps = len(scheduler_timesteps)
    return cast(torch.Tensor, scheduler_timesteps), num_inference_steps


_ANIMA_ORIGINAL_TRANSFORMER_RENAMES = {
    "t_embedder.1": "time_embed.t_embedder",
    "t_embedding_norm": "time_embed.norm",
    "blocks": "transformer_blocks",
    "adaln_modulation_self_attn.1": "norm1.linear_1",
    "adaln_modulation_self_attn.2": "norm1.linear_2",
    "adaln_modulation_cross_attn.1": "norm2.linear_1",
    "adaln_modulation_cross_attn.2": "norm2.linear_2",
    "adaln_modulation_mlp.1": "norm3.linear_1",
    "adaln_modulation_mlp.2": "norm3.linear_2",
    "self_attn": "attn1",
    "cross_attn": "attn2",
    "q_proj": "to_q",
    "k_proj": "to_k",
    "v_proj": "to_v",
    "output_proj": "to_out.0",
    "q_norm": "norm_q",
    "k_norm": "norm_k",
    "mlp.layer1": "ff.net.0.proj",
    "mlp.layer2": "ff.net.2",
    "x_embedder.proj.1": "patch_embed.proj",
    "final_layer.adaln_modulation.1": "norm_out.linear_1",
    "final_layer.adaln_modulation.2": "norm_out.linear_2",
    "final_layer.linear": "proj_out",
}

_ANIMA_ORIGINAL_TRANSFORMER_DROP_KEYS = (
    "accum_video_sample_counter",
    "accum_image_sample_counter",
    "accum_iteration",
    "accum_train_in_hours",
    "pos_embedder.seq",
    "pos_embedder.dim_spatial_range",
    "pos_embedder.dim_temporal_range",
    "_extra_state",
)


def _anima_vae_scale_factor_from_vae(vae: Any) -> int:
    vae_config = getattr(vae, "config", None)
    spatial_compression_ratio = getattr(vae_config, "spatial_compression_ratio", None)
    if spatial_compression_ratio is not None:
        return int(spatial_compression_ratio)
    temporal_downsample = getattr(vae, "temperal_downsample", None)
    if temporal_downsample is not None:
        return int(2 ** len(temporal_downsample))
    return ANIMA_VAE_SCALE_FACTOR


def get_anima_post_process_func(od_config: OmniDiffusionConfig) -> Callable[[torch.Tensor], Any]:
    image_processor = VaeImageProcessor(vae_scale_factor=ANIMA_VAE_SCALE_FACTOR)

    def post_process_func(images: torch.Tensor) -> Any:
        return image_processor.postprocess(images)

    return post_process_func


class AnimaPipeline(nn.Module, DiffusionPipelineProfilerMixin, ProgressBarMixin):
    """Native Anima model path.

    The Anima transformer checkpoint is distributed as a single safetensors
    file. This loader builds the native pipeline components and runs the denoise
    transformer/text-conditioner directly.
    """

    supports_step_execution = False

    def __init__(self, *, od_config: OmniDiffusionConfig, device: torch.device | str | None = None) -> None:
        super().__init__()
        self.od_config = od_config
        self.parallel_config = od_config.parallel_config
        self.device = device or get_local_device()
        self.weights_sources: list[str] = []

        self._raise_unsupported_features()

        self._interrupt = False
        self._guidance_scale = 1.0
        self._num_timesteps = 0
        self._current_timestep = None
        self._profiler_initialized = False

    def _raise_unsupported_features(self) -> None:
        pc = self.od_config.parallel_config
        if pc.cfg_parallel_size > 1:
            raise NotImplementedError("AnimaPipeline does not yet support CFG parallel. Validate native CFG first.")
        if pc.sequence_parallel_size is not None and pc.sequence_parallel_size > 1:
            raise NotImplementedError("AnimaPipeline does not yet support sequence parallelism. Add an _sp_plan first.")
        if pc.tensor_parallel_size > 1:
            raise NotImplementedError("AnimaPipeline does not yet support tensor parallelism. Port TP linears first.")
        if pc.use_hsdp:
            raise NotImplementedError("AnimaPipeline does not yet support HSDP. Validate sharding after native load.")
        if self.od_config.cache_backend not in ("none", None):
            raise NotImplementedError("AnimaPipeline does not yet support TeaCache or Cache-DiT.")
        if self.od_config.quantization_config is not None:
            raise NotImplementedError("AnimaPipeline does not yet support quantized checkpoints.")
        if resolve_offload_strategy(self.od_config) is not OffloadStrategy.NONE:
            raise NotImplementedError("AnimaPipeline does not yet support CPU offload.")

    @staticmethod
    def _sub_state_dict(state_dict: dict[str, torch.Tensor], prefix: str) -> dict[str, torch.Tensor]:
        return {key.removeprefix(prefix): value for key, value in state_dict.items() if key.startswith(prefix)}

    @staticmethod
    def _convert_original_transformer_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        if "patch_embed.proj.weight" in state_dict:
            return state_dict

        # The raw Anima checkpoint uses the original training names for a
        # Cosmos-style denoiser. Convert those keys to the local
        # AnimaTransformer3DModel module names before strict loading.
        converted_state_dict = {}
        for key, value in state_dict.items():
            new_key = key.removeprefix("net.")
            if any(drop_key in new_key for drop_key in _ANIMA_ORIGINAL_TRANSFORMER_DROP_KEYS):
                continue
            for old_key, new_name in _ANIMA_ORIGINAL_TRANSFORMER_RENAMES.items():
                new_key = new_key.replace(old_key, new_name)
            if new_key in converted_state_dict:
                raise ValueError(f"Duplicate converted Anima transformer key: {new_key}")
            converted_state_dict[new_key] = value
        return converted_state_dict

    @staticmethod
    def _split_single_file_state_dict(
        state_dict: dict[str, torch.Tensor],
    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
        transformer_state_dict = AnimaPipeline._sub_state_dict(state_dict, "transformer.")
        text_conditioner_state_dict = AnimaPipeline._sub_state_dict(state_dict, "text_conditioner.")
        if transformer_state_dict and text_conditioner_state_dict:
            return transformer_state_dict, text_conditioner_state_dict

        text_conditioner_state_dict = AnimaPipeline._sub_state_dict(state_dict, "net.llm_adapter.")
        if text_conditioner_state_dict:
            transformer_state_dict = {
                key: value for key, value in state_dict.items() if not key.startswith("net.llm_adapter.")
            }
            return transformer_state_dict, text_conditioner_state_dict

        raise ValueError(
            "Anima native loader could not find text-conditioner weights. Expected either "
            "`text_conditioner.*` Diffusers keys or `net.llm_adapter.*` original Anima keys."
        )

    def _load_native_denoiser_components(
        self,
        state_dict: dict[str, torch.Tensor] | None = None,
    ) -> tuple[AnimaTransformer3DModel, AnimaTextConditioner]:
        if state_dict is None:
            model_path = self.od_config.model
            if not isinstance(model_path, str) or not os.path.isfile(model_path):
                raise ValueError(f"AnimaPipeline currently requires a local file path, got: {model_path}")

            state_dict = load_file(model_path, device="cpu")
        transformer_state_dict, text_conditioner_state_dict = self._split_single_file_state_dict(state_dict)
        del state_dict
        transformer_state_dict = self._convert_original_transformer_state_dict(transformer_state_dict)

        config_path = Path(self.od_config.model).with_suffix(".json")
        config = json.loads(config_path.read_text()) if config_path.is_file() else {}
        text_conditioner_config = {**ANIMA_TEXT_CONDITIONER_CONFIG, **config.get("text_conditioner", {})}
        transformer_config = {**ANIMA_TRANSFORMER_CONFIG, **config.get("transformer", {})}

        native_text_conditioner = AnimaTextConditioner(**cast(dict[str, Any], text_conditioner_config))
        native_text_conditioner.load_state_dict(text_conditioner_state_dict, strict=True)
        native_text_conditioner.to(device=self.device, dtype=self.od_config.dtype)
        del text_conditioner_state_dict

        native_transformer = AnimaTransformer3DModel(**cast(dict[str, Any], transformer_config))
        native_transformer.load_state_dict(transformer_state_dict, strict=True)
        native_transformer.to(device=self.device, dtype=self.od_config.dtype)
        del transformer_state_dict

        return native_transformer, native_text_conditioner

    def _component_path(self, name: str, default_model: str) -> str:
        custom_args = self.od_config.custom_pipeline_args or {}
        return str(custom_args.get(f"{name}_model", custom_args.get(f"{name}_path", default_model)))

    @staticmethod
    def _from_pretrained_kwargs(model: str, subfolder: str | None = None) -> dict[str, Any]:
        kwargs: dict[str, Any] = {"local_files_only": os.path.exists(model)}
        if subfolder is not None and os.path.isdir(os.path.join(model, subfolder)):
            kwargs["subfolder"] = subfolder
        return kwargs

    def _load_outer_components(self) -> tuple[Any, Any, Any, Any, Any]:
        custom_args = self.od_config.custom_pipeline_args or {}
        components_model_value = custom_args.get("components_model", custom_args.get("components_path"))
        if components_model_value is None:
            if not isinstance(self.od_config.model, str):
                raise ValueError("AnimaPipeline requires `model` or `components_model` to load outer components.")
            checkpoint_dir = os.path.dirname(os.path.abspath(self.od_config.model))
            components_model = checkpoint_dir if os.path.isdir(checkpoint_dir) else self.od_config.model
        else:
            components_model = str(components_model_value)

        text_encoder_model = self._component_path("text_encoder", components_model)
        vae_model = self._component_path("vae", components_model)
        tokenizer_model = self._component_path("tokenizer", components_model)
        t5_tokenizer_model = self._component_path("t5_tokenizer", components_model)

        text_encoder = Qwen3Model.from_pretrained(
            text_encoder_model,
            torch_dtype=self.od_config.dtype,
            **self._from_pretrained_kwargs(text_encoder_model, "text_encoder"),
        ).to(self.device)
        vae = DistributedAutoencoderKLQwenImage.from_pretrained(
            vae_model,
            torch_dtype=self.od_config.dtype,
            **self._from_pretrained_kwargs(vae_model, "vae"),
        ).to(self.device)
        tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_model,
            **self._from_pretrained_kwargs(tokenizer_model, "tokenizer"),
        )
        t5_tokenizer = T5TokenizerFast.from_pretrained(
            t5_tokenizer_model,
            **self._from_pretrained_kwargs(t5_tokenizer_model, "t5_tokenizer"),
        )

        scheduler_model = self._component_path("scheduler", components_model)
        has_scheduler_config = os.path.isfile(os.path.join(scheduler_model, "scheduler_config.json")) or os.path.isdir(
            os.path.join(scheduler_model, "scheduler")
        )
        if has_scheduler_config:
            scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
                scheduler_model,
                **self._from_pretrained_kwargs(scheduler_model, "scheduler"),
            )
        else:
            scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0)

        return text_encoder, tokenizer, t5_tokenizer, vae, scheduler

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]] | None = None) -> set[str]:
        logger.info("Loading Anima transformer/text-conditioner directly from single-file checkpoint.")
        state_dict = dict(weights) if weights is not None else None
        loaded = set(state_dict.keys()) if state_dict is not None else set()
        if not state_dict:
            state_dict = None
        with set_current_diffusion_config(self.od_config):
            native_transformer, native_text_conditioner = self._load_native_denoiser_components(state_dict)
        text_encoder, tokenizer, t5_tokenizer, vae, scheduler = self._load_outer_components()

        self.transformer = native_transformer
        self.text_conditioner = native_text_conditioner
        self.text_encoder = text_encoder
        self.tokenizer = tokenizer
        self.t5_tokenizer = t5_tokenizer
        self.vae = vae
        self.scheduler = scheduler
        self.vae_scale_factor = _anima_vae_scale_factor_from_vae(self.vae)
        self._setup_profiler()
        return loaded

    def _setup_profiler(self) -> None:
        if self._profiler_initialized:
            return
        self.setup_diffusion_pipeline_profiler(
            profiler_targets=["vae.decode", "diffuse", "text_encoder.forward"],
            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
        )
        self._profiler_initialized = True

    def _extract_prompts(self, prompts: list[Any]) -> tuple[str | list[str] | None, str | list[str] | None]:
        """Extract prompt and negative_prompt from OmniPromptType list."""
        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None
        if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts):
            negative_prompt = None
        elif prompts:
            negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts]
        else:
            negative_prompt = None
        return prompt, negative_prompt

    def _get_qwen_prompt_embeds(
        self,
        prompt: str | list[str],
        max_sequence_length: int,
        device: torch.device,
        dtype: torch.dtype,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        prompt = [prompt] if isinstance(prompt, str) else prompt

        text_inputs = self.tokenizer(
            prompt,
            padding="longest",
            max_length=max_sequence_length,
            truncation=True,
            return_tensors="pt",
        )
        text_input_ids = text_inputs.input_ids.to(device)
        prompt_attention_mask = text_inputs.attention_mask.to(device)
        if text_input_ids.shape[-1] == 0:
            text_input_ids = text_input_ids.new_zeros((text_input_ids.shape[0], 1))
            prompt_attention_mask = prompt_attention_mask.new_zeros((prompt_attention_mask.shape[0], 1))

        prompt_embeds = self.text_encoder(
            input_ids=text_input_ids,
            attention_mask=prompt_attention_mask,
            output_hidden_states=False,
        ).last_hidden_state
        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
        prompt_embeds = prompt_embeds * prompt_attention_mask.to(prompt_embeds).unsqueeze(-1)

        return prompt_embeds, prompt_attention_mask

    def _get_t5_prompt_ids(
        self,
        prompt: str | list[str],
        max_sequence_length: int,
        device: torch.device,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        prompt = [prompt] if isinstance(prompt, str) else prompt

        text_inputs = self.t5_tokenizer(
            prompt,
            padding="longest",
            max_length=max_sequence_length,
            truncation=True,
            return_tensors="pt",
        )
        return text_inputs.input_ids.to(device), text_inputs.attention_mask.to(device)

    def encode_prompt(
        self,
        prompt: str | list[str],
        negative_prompt: str | list[str] | None = None,
        prepare_unconditional_embeds: bool = True,
        max_sequence_length: int = 512,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> dict[str, torch.Tensor | None]:
        device = device or self.device
        dtype = dtype or self.text_encoder.dtype

        prompt = [prompt] if isinstance(prompt, str) else prompt
        batch_size = len(prompt)

        prompt_embeds, prompt_attention_mask = self._get_qwen_prompt_embeds(
            prompt=prompt,
            max_sequence_length=max_sequence_length,
            device=device,
            dtype=dtype,
        )
        t5_input_ids, t5_attention_mask = self._get_t5_prompt_ids(
            prompt=prompt,
            max_sequence_length=max_sequence_length,
            device=device,
        )

        negative_prompt_embeds = None
        negative_prompt_attention_mask = None
        negative_t5_input_ids = None
        negative_t5_attention_mask = None
        if prepare_unconditional_embeds:
            if negative_prompt is None:
                negative_prompt = ""
            if isinstance(negative_prompt, str):
                negative_prompt = [negative_prompt] * batch_size
            elif len(negative_prompt) != batch_size:
                raise ValueError(
                    f"`negative_prompt` length ({len(negative_prompt)}) must match prompt batch size ({batch_size})."
                )

            negative_prompt_embeds, negative_prompt_attention_mask = self._get_qwen_prompt_embeds(
                prompt=negative_prompt,
                max_sequence_length=max_sequence_length,
                device=device,
                dtype=dtype,
            )
            negative_t5_input_ids, negative_t5_attention_mask = self._get_t5_prompt_ids(
                prompt=negative_prompt,
                max_sequence_length=max_sequence_length,
                device=device,
            )

        return {
            "qwen_prompt_embeds": prompt_embeds,
            "qwen_attention_mask": prompt_attention_mask,
            "t5_input_ids": t5_input_ids,
            "t5_attention_mask": t5_attention_mask,
            "negative_qwen_prompt_embeds": negative_prompt_embeds,
            "negative_qwen_attention_mask": negative_prompt_attention_mask,
            "negative_t5_input_ids": negative_t5_input_ids,
            "negative_t5_attention_mask": negative_t5_attention_mask,
        }

    def condition_prompt_embeds(
        self,
        qwen_prompt_embeds: torch.Tensor,
        qwen_attention_mask: torch.Tensor,
        t5_input_ids: torch.Tensor,
        t5_attention_mask: torch.Tensor,
        device: torch.device | None = None,
        conditioning_dtype: torch.dtype | None = None,
        output_dtype: torch.dtype | None = None,
    ) -> torch.Tensor:
        device = device or self.device
        conditioning_dtype = conditioning_dtype or self.text_conditioner.dtype
        output_dtype = output_dtype or self.transformer.dtype

        prompt_embeds = self.text_conditioner(
            source_hidden_states=qwen_prompt_embeds.to(device=device, dtype=conditioning_dtype),
            target_input_ids=t5_input_ids.to(device),
            target_attention_mask=t5_attention_mask.to(device),
            source_attention_mask=qwen_attention_mask.to(device),
        )
        return prompt_embeds.to(dtype=output_dtype, device=device)

    def prepare_latents(
        self,
        batch_size: int,
        num_channels_latents: int,
        height: int,
        width: int,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
        generator: torch.Generator | list[torch.Generator] | None = None,
        latents: torch.Tensor | None = None,
    ) -> torch.Tensor:
        device = device or self.device
        dtype = dtype or self.transformer.dtype
        if latents is not None:
            return latents.to(device=device, dtype=dtype)

        latent_height = height // self.vae_scale_factor
        latent_width = width // self.vae_scale_factor
        shape = (batch_size, num_channels_latents, 1, latent_height, latent_width)

        if isinstance(generator, list) and len(generator) != batch_size:
            raise ValueError(
                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
            )

        # Draw latents in float32 then cast
        latents = randn_tensor(shape, generator=generator, device=device, dtype=torch.float32)
        sigma_max = (
            getattr(self.scheduler.config, "sigma_max", 1.0)
            if hasattr(self, "scheduler") and hasattr(self.scheduler, "config")
            else 1.0
        )
        return latents.to(dtype=dtype) * sigma_max

    def prepare_timesteps(
        self,
        num_inference_steps: int,
        sigmas: list[float] | None = None,
        device: torch.device | None = None,
    ) -> tuple[torch.Tensor, int]:
        device = device or self.device
        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
        return retrieve_timesteps(
            self.scheduler,
            device=device,
            sigmas=sigmas,
        )

    def diffuse(
        self,
        prompt_embeds: torch.Tensor,
        negative_prompt_embeds: torch.Tensor | None,
        latents: torch.Tensor,
        padding_mask: torch.Tensor,
        timesteps: torch.Tensor,
        do_true_cfg: bool,
        true_cfg_scale: float,
    ) -> torch.Tensor:
        self.scheduler.set_begin_index(0)

        with self.progress_bar(total=len(timesteps)) as pbar:
            for i, t in enumerate(timesteps):
                if self.interrupt:
                    continue
                self._current_timestep = t

                # Broadcast timestep to match batch size
                timestep = t.expand(latents.shape[0]).to(dtype=self.transformer.dtype, device=latents.device)
                timestep = timestep / 1000.0  # normalize to 0..1 range

                # Predict noise
                latent_model_input = latents.to(self.transformer.dtype)

                # Positive pass
                noise_pred = self.transformer(
                    hidden_states=latent_model_input,
                    timestep=timestep,
                    encoder_hidden_states=prompt_embeds,
                    padding_mask=padding_mask,
                    return_dict=False,
                )[0]

                # Negative / uncond pass
                if do_true_cfg and negative_prompt_embeds is not None:
                    noise_pred_uncond = self.transformer(
                        hidden_states=latent_model_input,
                        timestep=timestep,
                        encoder_hidden_states=negative_prompt_embeds,
                        padding_mask=padding_mask,
                        return_dict=False,
                    )[0]
                    noise_pred = noise_pred_uncond + true_cfg_scale * (noise_pred - noise_pred_uncond)

                # Scheduler step
                latents_dtype = latents.dtype
                latents = self.scheduler.step(noise_pred.to(latents_dtype), t, latents, return_dict=False)[0]
                if latents.dtype != latents_dtype:
                    latents = latents.to(latents_dtype)

                pbar.update()

        return latents

    def decode_latents(
        self,
        latents: torch.Tensor,
        output_type: str = "pil",
    ) -> DiffusionOutput:
        """Decode final latents."""
        if output_type == "latent":
            return DiffusionOutput(
                output=latents,
                stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
            )

        latents = latents.to(self.vae.dtype)
        latents_mean = (
            torch.tensor(self.vae.config.latents_mean)
            .view(1, self.vae.config.z_dim, 1, 1, 1)
            .to(latents.device, latents.dtype)
        )
        latents_std = (
            torch.tensor(self.vae.config.latents_std)
            .view(1, self.vae.config.z_dim, 1, 1, 1)
            .to(latents.device, latents.dtype)
        )
        latents = latents * latents_std + latents_mean
        image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
        return DiffusionOutput(
            output=image,
            stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None,
        )

    def forward(
        self,
        req: DiffusionRequestBatch,
        prompt: str | list[str] | None = None,
        negative_prompt: str | list[str] | None = None,
        true_cfg_scale: float = 4.0,
        height: int | None = None,
        width: int | None = None,
        num_inference_steps: int = 50,
        sigmas: list[float] | None = None,
        guidance_scale: float = 4.0,
        num_images_per_prompt: int = 1,
        generator: torch.Generator | list[torch.Generator] | None = None,
        latents: torch.Tensor | None = None,
        output_type: str | None = "pil",
        max_sequence_length: int = 512,
    ) -> list[DiffusionOutput]:
        if req.num_reqs != 1:
            raise ValueError("AnimaPipeline supports one request per forward.")
        extracted_prompt, negative_prompt = self._extract_prompts(req.prompts)
        prompt = extracted_prompt or prompt
        if prompt is None:
            prompt = ""

        height = req.sampling_params.height or height or 1024
        width = req.sampling_params.width or width or 1024
        height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2)
        if height is None or width is None:
            raise ValueError("AnimaPipeline requires resolved height and width.")
        height = int(height)
        width = int(width)
        if req.sampling_params.output_type is not None:
            output_type = req.sampling_params.output_type
        output_type = output_type or "pil"

        num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
        sigmas = req.sampling_params.sigmas or sigmas
        max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length
        generator = req.sampling_params.generator or generator
        if req.sampling_params.guidance_scale_provided:
            guidance_scale = req.sampling_params.guidance_scale
            if req.sampling_params.true_cfg_scale is None:
                true_cfg_scale = guidance_scale
        if req.sampling_params.true_cfg_scale is not None:
            true_cfg_scale = req.sampling_params.true_cfg_scale
        num_images_per_prompt = (
            req.sampling_params.num_outputs_per_prompt
            if req.sampling_params.num_outputs_per_prompt > 0
            else num_images_per_prompt
        )

        self._guidance_scale = guidance_scale
        self._interrupt = False

        if isinstance(prompt, str):
            batch_size = 1
        else:
            batch_size = len(prompt)

        do_true_cfg = guidance_scale > 1.0

        # Encode prompts
        enc_out = self.encode_prompt(
            prompt=prompt,
            negative_prompt=negative_prompt,
            prepare_unconditional_embeds=do_true_cfg,
            max_sequence_length=max_sequence_length,
            device=self.device,
            dtype=self.text_encoder.dtype,
        )
        qwen_prompt_embeds = enc_out["qwen_prompt_embeds"]
        qwen_attention_mask = enc_out["qwen_attention_mask"]
        t5_input_ids = enc_out["t5_input_ids"]
        t5_attention_mask = enc_out["t5_attention_mask"]
        if (
            qwen_prompt_embeds is None
            or qwen_attention_mask is None
            or t5_input_ids is None
            or t5_attention_mask is None
        ):
            raise ValueError("Anima prompt encoding did not produce all required positive embeddings.")

        # Condition prompt embeds
        prompt_embeds = self.condition_prompt_embeds(
            qwen_prompt_embeds=qwen_prompt_embeds,
            qwen_attention_mask=qwen_attention_mask,
            t5_input_ids=t5_input_ids,
            t5_attention_mask=t5_attention_mask,
            device=self.device,
        )

        negative_prompt_embeds = None
        negative_qwen_prompt_embeds = enc_out["negative_qwen_prompt_embeds"]
        if do_true_cfg and negative_qwen_prompt_embeds is not None:
            negative_qwen_attention_mask = enc_out["negative_qwen_attention_mask"]
            negative_t5_input_ids = enc_out["negative_t5_input_ids"]
            negative_t5_attention_mask = enc_out["negative_t5_attention_mask"]
            if (
                negative_qwen_attention_mask is None
                or negative_t5_input_ids is None
                or negative_t5_attention_mask is None
            ):
                raise ValueError("Anima prompt encoding did not produce all required negative embeddings.")
            negative_prompt_embeds = self.condition_prompt_embeds(
                qwen_prompt_embeds=negative_qwen_prompt_embeds,
                qwen_attention_mask=negative_qwen_attention_mask,
                t5_input_ids=negative_t5_input_ids,
                t5_attention_mask=negative_t5_attention_mask,
                device=self.device,
            )

        # Repeat for num_images_per_prompt
        _, seq_len, _ = prompt_embeds.shape
        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)

        if negative_prompt_embeds is not None:
            _, seq_len, _ = negative_prompt_embeds.shape
            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)

        # Prepare latents
        latents = self.prepare_latents(
            batch_size=batch_size * num_images_per_prompt,
            num_channels_latents=16,  # Cosmos/Anima has 16 channels
            height=height,
            width=width,
            dtype=torch.float32,
            device=self.device,
            generator=generator,
            latents=latents,
        )

        # Prepare padding mask
        padding_mask = latents.new_zeros(1, 1, height, width, dtype=self.transformer.dtype)

        # Prepare timesteps
        timesteps, num_inference_steps = self.prepare_timesteps(
            num_inference_steps=num_inference_steps,
            sigmas=sigmas,
            device=self.device,
        )
        self._num_timesteps = len(timesteps)

        # Denoise loop
        latents = self.diffuse(
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            latents=latents,
            padding_mask=padding_mask,
            timesteps=timesteps,
            do_true_cfg=do_true_cfg,
            true_cfg_scale=true_cfg_scale,
        )

        self._current_timestep = None
        return [self.decode_latents(latents, output_type=output_type)]

    @property
    def guidance_scale(self) -> float:
        return self._guidance_scale

    @property
    def num_timesteps(self) -> int:
        return self._num_timesteps

    @property
    def current_timestep(self) -> torch.Tensor | None:
        return self._current_timestep

    @property
    def interrupt(self) -> bool:
        return self._interrupt