vllm-omni/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

from __future__ import annotations

from collections.abc import Iterable
from typing import Any

import torch
from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
from diffusers.utils.torch_utils import randn_tensor
from torch import nn
from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
from vllm.config import VllmConfig
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper

from vllm_omni.model_executor.models.output_templates import OmniOutput
from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config

from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DModel
from .rope_real import RotaryPosEmbedReal
from .schedulers import FlowMatchEulerDiscreteScheduler


class MammothModa2DiTPipeline(nn.Module):
    """
    MammothModa2 DiT + VAE generation stage (non-autoregressive).

    This stage expects "image condition token hidden states" from the upstream AR stage,
    and outputs image tensors via diffusion transformer + VAE decode.

    """

    have_multimodal_outputs = True

    # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage.
    hf_to_vllm_mapper = WeightsMapper(
        orig_to_new_prefix={
            "llm_model.": None,
        }
    )

    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        del prefix

        hf_config = vllm_config.model_config.hf_config
        if not isinstance(hf_config, Mammothmoda2Config):
            raise TypeError(f"Expected Mammothmoda2Config, got {type(hf_config)}")

        self.config = hf_config

        # --- Build DiT / VAE modules (names must match checkpoint keys) ---
        if self.config.gen_vae_config is None or self.config.gen_dit_config is None:
            raise ValueError("Mammothmoda2Config.gen_vae_config / gen_dit_config must not be None")

        self.gen_vae = AutoencoderKL.from_config(self.config.gen_vae_config)
        self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config)

        llm_hidden_size = int(getattr(self.config.llm_config, "hidden_size", 0) or 0)
        if llm_hidden_size <= 0:
            raise ValueError("Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.hidden_size")
        self._reinit_caption_embedder(llm_hidden_size)

        # Optional: image condition refiner (Q-Former)
        if self.config.gen_image_condition_refiner_config is not None:
            self.gen_image_condition_refiner = SimpleQFormerImageRefiner(
                hidden_size=llm_hidden_size,
                **self.config.gen_image_condition_refiner_config,
            )
        else:
            self.gen_image_condition_refiner = None

        # Precompute rotary freqs for diffusion transformer
        # IMPORTANT: follow upstream mammothmoda: use top-level `config.gen_axes_*`
        # (the checkpoint's `gen_dit_config.axes_lens` can be as small as 1024,
        # which is insufficient for vLLM dummy-run/cudagraph warmup).
        self.gen_freqs_cis = RotaryPosEmbedReal.get_freqs_real(
            tuple(self.config.gen_axes_dim_rope),
            tuple(self.config.gen_axes_lens),
            theta=10000,
        )

        # vLLM PP interface compatibility
        self.make_empty_intermediate_tensors = lambda: None

        self._llm_hidden_size = llm_hidden_size

    def _reinit_caption_embedder(self, in_features: int) -> None:
        # Align with upstream Mammothmoda2Model's `reinit_caption_embedder`:
        # Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features).
        out_features = int(getattr(self.gen_transformer, "hidden_size", 0) or self.gen_transformer.config.hidden_size)
        self.gen_transformer.time_caption_embed.caption_embedder = nn.Sequential(
            Qwen2RMSNorm(in_features, eps=1e-5),
            nn.Linear(in_features, out_features, bias=True),
        )

    def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]:
        if num_reqs <= 0:
            raise ValueError(f"num_reqs must be positive, got {num_reqs}")
        if num_reqs > 1:
            raise NotImplementedError(
                f"get_dummy_runtime_additional_information does not support num_reqs > 1, got {num_reqs}"
            )
        text_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32)
        image_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32)
        negative_prompt_embeds = torch.zeros((0, self._llm_hidden_size), dtype=torch.float32)
        info = {
            "text_prompt_embeds": text_prompt_embeds,
            "image_prompt_embeds": image_prompt_embeds,
            "negative_prompt_embeds": negative_prompt_embeds,
            "negative_prompt_attention_mask": [],
            "image_height": [512],
            "image_width": [512],
            "text_guidance_scale": [1.0],
            "cfg_range": [0.0, 1.0],
            "num_inference_steps": [1],
        }
        return [info for _ in range(num_reqs)]

    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        # DiT stage does not consume token embeddings; return a dummy tensor.
        try:
            dtype = next(self.parameters()).dtype
        except StopIteration:
            dtype = torch.float32
        return torch.zeros(
            (input_ids.numel(), self._llm_hidden_size),
            device=input_ids.device,
            dtype=dtype,
        )

    @torch.inference_mode()
    def forward(
        self,
        *,
        inputs_embeds: torch.Tensor | None = None,
        **kwargs: Any,  # noqa: ARG002
    ) -> OmniOutput:
        runtime_addi = kwargs.get("runtime_additional_information", None)
        info = runtime_addi[0]
        text_cond = info["text_prompt_embeds"]
        image_cond = info["image_prompt_embeds"]
        negative_cond = info.get("negative_prompt_embeds")
        negative_attention_mask = info.get("negative_prompt_attention_mask")
        image_hw = info["image_height"][0], info["image_width"][0]
        text_guidance_scale = info["text_guidance_scale"][0]
        cfg_range = info["cfg_range"][0], info["cfg_range"][1]
        num_inference_steps = info["num_inference_steps"][0]

        # Move to model device/dtype.
        model_device = next(self.parameters()).device
        if self.gen_image_condition_refiner is not None:
            target_dtype = next(self.gen_image_condition_refiner.parameters()).dtype
        else:
            target_dtype = next(self.gen_transformer.parameters()).dtype

        def _ensure_2d(x: torch.Tensor, name: str) -> torch.Tensor:
            if x.ndim == 3 and x.shape[0] == 1:
                x = x[0]
            if x.ndim != 2:
                raise ValueError(f"Expected {name} to be 2D [T,H], got shape={tuple(x.shape)}")
            return x

        text_cond = _ensure_2d(text_cond, "text_prompt_embeds")
        image_cond = _ensure_2d(image_cond, "image_prompt_embeds")
        text_cond = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()
        image_cond = image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()

        text_embeds = text_cond.unsqueeze(0)  # [1, T_text, H]
        text_attention_mask = torch.ones(
            (1, text_embeds.shape[1]),
            dtype=torch.bool,
            device=text_embeds.device,
        )

        image_embeds = image_cond.unsqueeze(0)  # [1, T_img, H]
        image_attention_mask = torch.ones(
            (1, image_embeds.shape[1]),
            dtype=torch.bool,
            device=image_embeds.device,
        )

        # Apply optional refiner ONLY on image condition tokens.
        if self.gen_image_condition_refiner is not None and image_embeds.shape[1] > 0:
            image_embeds = self.gen_image_condition_refiner(image_embeds, ~image_attention_mask.bool())
            image_attention_mask = torch.ones(
                image_embeds.shape[:2],
                dtype=torch.bool,
                device=image_embeds.device,
            )

        prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1)
        prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)

        # Prepare negative prompt (for CFG). If none provided, fall back to unconditional.
        negative_prompt_embeds = None
        negative_prompt_attention_mask = None
        if text_guidance_scale > 1.0:
            if negative_cond is not None:
                negative_cond = _ensure_2d(negative_cond, "negative_prompt_embeds")
                negative_prompt_embeds = (
                    negative_cond.to(device=model_device, dtype=target_dtype, non_blocking=True)
                    .contiguous()
                    .unsqueeze(0)
                )
                if isinstance(negative_attention_mask, torch.Tensor):
                    neg_mask = negative_attention_mask
                elif isinstance(negative_attention_mask, list):
                    neg_mask = torch.tensor(negative_attention_mask, dtype=torch.bool)
                else:
                    neg_mask = None
                if neg_mask is None:
                    negative_prompt_attention_mask = torch.ones(
                        (1, negative_prompt_embeds.shape[1]),
                        dtype=torch.bool,
                        device=negative_prompt_embeds.device,
                    )
                else:
                    neg_mask = neg_mask.to(device=negative_prompt_embeds.device, dtype=torch.bool)
                    if neg_mask.ndim == 1:
                        neg_mask = neg_mask.unsqueeze(0)
                    negative_prompt_attention_mask = neg_mask
            else:
                hidden_size = int(prompt_embeds.shape[-1])
                negative_prompt_embeds = torch.zeros(
                    (1, 0, hidden_size),
                    dtype=target_dtype,
                    device=prompt_embeds.device,
                )
                negative_prompt_attention_mask = torch.zeros(
                    (1, 0),
                    dtype=torch.bool,
                    device=prompt_embeds.device,
                )

        # Output image size (px), passed from stage input processor.
        height, width = image_hw
        if height <= 0 or width <= 0:
            raise ValueError(f"Invalid image size: {height}x{width}")
        if height % 16 != 0 or width % 16 != 0:
            raise ValueError(f"Image size must be multiples of 16, got {height}x{width}")
        vae_scale_factor = 16

        latent_channels = int(self.gen_transformer.config.in_channels)
        shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
        latents = randn_tensor(shape, device=prompt_embeds.device, dtype=prompt_embeds.dtype)

        scheduler = FlowMatchEulerDiscreteScheduler()

        scheduler.set_timesteps(
            num_inference_steps=num_inference_steps,
            device=prompt_embeds.device,
            num_tokens=latents.shape[-2] * latents.shape[-1],
        )

        # Run diffusion loop (CFG supported when text_guidance_scale > 1.0)
        total_steps = max(1, len(scheduler.timesteps))
        for i, t in enumerate(scheduler.timesteps):
            timestep = t.expand(latents.shape[0]).to(latents.dtype)
            model_pred = self.gen_transformer(
                hidden_states=latents,
                timestep=timestep,
                text_hidden_states=prompt_embeds,
                text_attention_mask=prompt_attention_mask,
                ref_image_hidden_states=None,
                freqs_cis=self.gen_freqs_cis,
            )
            guidance_scale = text_guidance_scale if cfg_range[0] <= i / total_steps <= cfg_range[1] else 1.0
            if guidance_scale > 1.0 and negative_prompt_embeds is not None:
                model_pred_uncond = self.gen_transformer(
                    hidden_states=latents,
                    timestep=timestep,
                    text_hidden_states=negative_prompt_embeds,
                    text_attention_mask=negative_prompt_attention_mask,
                    ref_image_hidden_states=None,
                    freqs_cis=self.gen_freqs_cis,
                )
                model_pred = model_pred_uncond + guidance_scale * (model_pred - model_pred_uncond)
            latents = scheduler.step(model_pred, t, latents, return_dict=False)[0]
            latents = latents.to(dtype=prompt_embeds.dtype)

        # VAE decode
        if self.gen_vae.config.scaling_factor is not None:
            latents = latents / self.gen_vae.config.scaling_factor
        if self.gen_vae.config.shift_factor is not None:
            latents = latents + self.gen_vae.config.shift_factor
        image = self.gen_vae.decode(latents, return_dict=False)[0]

        return OmniOutput(
            text_hidden_states=inputs_embeds,  # placeholder, not used by runner
            multimodal_outputs=image,
            intermediate_tensors=None,
        )

    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:  # noqa: ARG002
        return None

    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        loader = AutoWeightsLoader(self)
        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)