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

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

from collections.abc import Iterable
from dataclasses import dataclass
from typing import ClassVar

import torch
from diffusers.image_processor import VaeImageProcessor
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.logger import init_logger
from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper

from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
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

logger = init_logger(__name__)


def _first_request_value(value: object) -> object:
    if isinstance(value, (list, tuple)):
        return value[0] if value else None
    return value


def get_mammoth_moda2_post_process_func(
    _od_config: OmniDiffusionConfig,
):
    image_processor = VaeImageProcessor()

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

    return post_process_func


def _build_mammoth_config(od_config: OmniDiffusionConfig) -> Mammothmoda2Config:
    raw_config = od_config.tf_model_config.to_dict()
    if not raw_config:
        raise ValueError("MammothModa2 diffusion stage requires the root checkpoint config")
    return Mammothmoda2Config(**raw_config)


def _root_weight_source(
    od_config: OmniDiffusionConfig,
) -> DiffusersPipelineLoader.ComponentSource:
    if not od_config.model:
        raise ValueError("MammothModa2 diffusion stage requires a model path")
    return DiffusersPipelineLoader.ComponentSource(
        model_or_path=od_config.model,
        subfolder=None,
        revision=od_config.revision,
        prefix="",
        fall_back_to_pt=True,
    )


@dataclass(frozen=True)
class _MammothRequest:
    request_id: str
    full_hidden_states: torch.Tensor
    full_token_ids: list[int]
    answer_start_index: int
    height: int
    width: int
    text_guidance_scale: float
    cfg_range: tuple[float, float]
    num_inference_steps: int
    seed: int | None
    generator: torch.Generator | list[torch.Generator] | None


class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
    """
    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.

    """

    _dit_modules: ClassVar[list[str]] = ["gen_transformer"]
    _encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"]
    _vae_modules: ClassVar[list[str]] = ["gen_vae"]

    supports_request_batch = False
    supports_step_execution = False

    # 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,
            "gen_tokenizer.": None,
        }
    )

    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
        super().__init__()
        del prefix
        self.od_config = od_config
        self.device = get_local_device()
        self.config = _build_mammoth_config(od_config)
        self.weights_sources = [_root_weight_source(od_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_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
        llm_hidden_size = 0
        text_config = self.config.get_text_config()
        if text_config is None:
            logger.warning("No text config; failed to infer llm_hidden_size.")
        elif not hasattr(text_config, "hidden_size"):
            logger.warning("Text config exists, but has no hidden_size attribute; failed to infer llm_hidden_size.")
        else:
            llm_hidden_size = int(text_config.hidden_size or 0)
        if llm_hidden_size <= 0:
            raise ValueError(
                "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size"
            )
        self._reinit_caption_embedder(llm_hidden_size)

        # Optional image condition Q-Former. Preview stores it as a standalone
        # module; Dev stores it under the DiT timestep/caption embedder.
        llm_model_type = getattr(self.config.llm_config, "model_type", "")
        refiner_config = self.config.gen_image_condition_refiner_config
        if refiner_config is not None and llm_model_type == "mammothmoda2_qwen3_vl":
            dit_hidden_size = int(self.gen_transformer.hidden_size)
            self.gen_transformer.time_caption_embed.image_embedder = SimpleQFormerImageRefiner(
                hidden_size=llm_hidden_size,
                output_hidden_size=dit_hidden_size,
                num_heads=max(1, dit_hidden_size // 128),
                **refiner_config,
            )
            self.gen_image_condition_refiner = None
        elif refiner_config is not None:
            self.gen_image_condition_refiner = SimpleQFormerImageRefiner(
                hidden_size=llm_hidden_size,
                **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,
        )

        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 _parse_request(self, req: DiffusionRequestBatch) -> _MammothRequest:
        if req.num_reqs != 1:
            raise ValueError("MammothModa2 diffusion requires exactly one request")
        request_id = req.request_id
        sampling = req.sampling_params
        if sampling.num_outputs_per_prompt != 1:
            raise ValueError(f"MammothModa2 requires num_outputs_per_prompt=1 for request {request_id}")

        prompt = req.prompts[0]
        prompt = prompt if isinstance(prompt, dict) else {}
        info = prompt.get("additional_information")
        if req.is_dummy_run():
            full_hidden_states = torch.zeros((2, self._llm_hidden_size), dtype=torch.float32, device="cpu")
            full_token_ids = [0, int(self.config.llm_config.gen_vocab_start_index)]
            answer_start_index = 1
        else:
            if not isinstance(info, dict):
                raise ValueError(f"Missing additional_information AR conditions for request {request_id}")
            full_hidden_states = info.get("full_hidden_states")
            full_token_ids = info.get("full_token_ids")
            if not isinstance(full_hidden_states, torch.Tensor) or not isinstance(full_token_ids, list):
                raise ValueError(f"Expected full_hidden_states tensor and full_token_ids list for request {request_id}")
            try:
                answer_start_index = int(info.get("answer_start_index"))
            except (TypeError, ValueError, OverflowError) as exc:
                raise ValueError(f"Invalid answer_start_index for request {request_id}") from exc
            if full_hidden_states.ndim != 2:
                raise ValueError(f"Expected 2D full_hidden_states for request {request_id}")
            if full_hidden_states.shape[0] != len(full_token_ids):
                raise ValueError(f"AR hidden-state/token-count mismatch for request {request_id}")
            if not 0 <= answer_start_index <= len(full_token_ids):
                raise ValueError(f"answer_start_index outside token range for request {request_id}")
            try:
                full_token_ids = [int(token_id) for token_id in full_token_ids]
            except (TypeError, ValueError, OverflowError) as exc:
                raise ValueError(f"Invalid full_token_ids for request {request_id}") from exc

        dimensions = []
        for name in ("height", "width"):
            value = DiffusionRequestBatch.get_prompt_field(prompt, name)
            if value is None:
                value = getattr(sampling, name)
            if value is None:
                value = 1024
            try:
                dimensions.append(int(value))
            except (TypeError, ValueError, OverflowError) as exc:
                raise ValueError(f"Invalid image size {name}={value!r} for request {request_id}") from exc
        height, width = dimensions
        if height <= 0 or width <= 0:
            raise ValueError(f"Invalid image size: {height}x{width} for request {request_id}")
        if height % 16 != 0 or width % 16 != 0:
            raise ValueError(f"Image size must be multiples of 16, got {height}x{width} for request {request_id}")

        request_info = info if isinstance(info, dict) else {}
        extra_args = sampling.extra_args or {}
        guidance = extra_args.get("text_guidance_scale")
        if guidance is None:
            guidance = sampling.guidance_scale if sampling.guidance_scale_provided else None
        if guidance is None:
            guidance = _first_request_value(request_info.get("text_guidance_scale"))
        if guidance is None:
            guidance = 9.0
        try:
            text_guidance_scale = float(guidance)
        except (TypeError, ValueError, OverflowError) as exc:
            raise ValueError(f"Invalid text_guidance_scale for request {request_id}") from exc
        raw_num_inference_steps = extra_args.get("num_inference_steps")
        if raw_num_inference_steps is None:
            raw_num_inference_steps = sampling.num_inference_steps
        if raw_num_inference_steps is None:
            raw_num_inference_steps = _first_request_value(request_info.get("num_inference_steps"))
        if raw_num_inference_steps is None:
            raw_num_inference_steps = 50
        try:
            num_inference_steps = int(raw_num_inference_steps)
        except (TypeError, ValueError, OverflowError) as exc:
            raise ValueError(f"Invalid num_inference_steps for request {request_id}") from exc
        if num_inference_steps <= 0:
            raise ValueError(f"num_inference_steps must be positive for request {request_id}")
        cfg_range = extra_args.get("cfg_range")
        if cfg_range is None:
            cfg_range = request_info.get("cfg_range")
        if cfg_range is None:
            cfg_range = [0.0, 1.0]
        if not isinstance(cfg_range, (list, tuple)) or len(cfg_range) != 2:
            raise ValueError(f"cfg_range requires two values for request {request_id}")
        try:
            cfg_start, cfg_end = float(cfg_range[0]), float(cfg_range[1])
        except (TypeError, ValueError, OverflowError) as exc:
            raise ValueError(f"cfg_range requires two values convertible to floats for request {request_id}") from exc
        if not 0 <= cfg_start <= cfg_end <= 1:
            raise ValueError(f"cfg_range must satisfy 0 <= start <= end <= 1 for request {request_id}")

        generator = sampling.generator
        if isinstance(generator, list) and len(generator) != 1:
            raise ValueError(
                f"MammothModa2 single-output request mode requires exactly one generator for request {request_id}"
            )

        return _MammothRequest(
            request_id=request_id,
            full_hidden_states=full_hidden_states,
            full_token_ids=full_token_ids,
            answer_start_index=answer_start_index,
            height=height,
            width=width,
            text_guidance_scale=text_guidance_scale,
            cfg_range=(cfg_start, cfg_end),
            num_inference_steps=num_inference_steps,
            seed=sampling.seed,
            generator=generator,
        )

    def _split_ar_conditions(
        self,
        *,
        full_hidden_states: torch.Tensor,
        full_token_ids: list[int],
        answer_start_index: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Split AR-stage hidden states into text / image condition embeds.

        The token ids that distinguish question (text) tokens, generated visual
        tokens, and multi-modal placeholder tokens are read from the model config
        (``gen_vocab_start_index`` and the vision placeholder token ids), so the
        caller no longer needs to pass them. Mirrors the masking the bespoke
        MammothModa2 example performed via ar2dit.
        """
        gen_vocab_start_index = int(self.config.llm_config.gen_vocab_start_index)
        visual_ids = [
            int(self.config.image_token_id),
            int(self.config.video_token_id),
            int(self.config.vision_start_token_id),
            int(self.config.vision_end_token_id),
        ]

        device = full_hidden_states.device
        token_ids = torch.tensor(full_token_ids, dtype=torch.long, device=device)
        positions = torch.arange(token_ids.shape[0], device=device)
        questions_mask = positions < answer_start_index
        answers_mask = ~questions_mask
        gen_token_mask = token_ids >= gen_vocab_start_index
        visual_token_mask = torch.isin(token_ids, torch.tensor(visual_ids, dtype=torch.long, device=device))
        text_mask = questions_mask & ~(visual_token_mask | gen_token_mask)
        image_mask = answers_mask & gen_token_mask

        # Keep the transferred representation compact until the final device/dtype
        # conversion in ``forward``.  The masks are dtype-independent, and an
        # unconditional host-side float32 expansion here would otherwise double both
        # the selected-condition staging footprint and BF16 H2D traffic.
        text_cond = full_hidden_states[text_mask].contiguous()
        image_cond = full_hidden_states[image_mask].contiguous()
        return text_cond, image_cond

    @torch.inference_mode()
    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
        request = self._parse_request(req)
        text_cond, image_cond = self._split_ar_conditions(
            full_hidden_states=request.full_hidden_states,
            full_token_ids=request.full_token_ids,
            answer_start_index=request.answer_start_index,
        )
        if image_cond.shape[0] == 0:
            answer_token_ids = request.full_token_ids[request.answer_start_index :]
            raise ValueError(
                "MammothModa2 AR stage produced no visual-token hidden states for "
                f"request {request.request_id}; the DiT stage requires at least one generated visual token. "
                f"Generated token ids: {answer_token_ids[:32]}"
            )

        # 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

        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,
            )

        nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None)
        if nested_image_embedder is None:
            prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1)
            prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
            ar_image_embeds = None
            ar_image_attention_mask = None
        else:
            prompt_embeds = text_embeds
            prompt_attention_mask = text_attention_mask
            ar_image_embeds = image_embeds
            ar_image_attention_mask = image_attention_mask

        # Empty unconditional prompt for classifier-free guidance.
        negative_prompt_embeds = None
        negative_prompt_attention_mask = None
        if request.text_guidance_scale > 1.0:
            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,
            )

        generator = request.generator
        if generator is None and request.seed is not None:
            generator = torch.Generator(device=model_device).manual_seed(request.seed)

        # Output image size (px), passed from stage input processor.
        height, width = request.height, request.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, generator=generator, device=prompt_embeds.device, dtype=prompt_embeds.dtype)

        scheduler = FlowMatchEulerDiscreteScheduler()

        scheduler.set_timesteps(
            num_inference_steps=request.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,
                ar_image_hidden_states=ar_image_embeds,
                ar_image_attention_mask=ar_image_attention_mask,
                freqs_cis=self.gen_freqs_cis,
            )
            guidance_scale = (
                request.text_guidance_scale if request.cfg_range[0] <= i / total_steps <= request.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 DiffusionOutput(output=image)

    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)