from __future__ import annotations from collections.abc import Iterable from typing import Any, ClassVar 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.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery 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 logger = init_logger(__name__) 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"] 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, "gen_tokenizer.": 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_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, ) # 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, ) 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 text_cond = full_hidden_states[text_mask].to(dtype=torch.float32).contiguous() image_cond = full_hidden_states[image_mask].to(dtype=torch.float32).contiguous() return text_cond, image_cond @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] # Sampling knobs are declared in vllm_omni/model_extras/mammothmodal2_preview.py # and routed via extra_body -> sampling_params.extra_args (surfaced here as # ``sampling_extra_args``). Fall back to runtime_additional_information for the # legacy bespoke-example path during the transition. extra_args_list = kwargs.get("sampling_extra_args") or [] extra_args = extra_args_list[0] if extra_args_list else {} text_guidance_scale = float(extra_args.get("text_guidance_scale", info["text_guidance_scale"][0])) cfg_range_val = extra_args.get("cfg_range", info["cfg_range"]) cfg_range = float(cfg_range_val[0]), float(cfg_range_val[1]) num_inference_steps = int(extra_args.get("num_inference_steps", info["num_inference_steps"][0])) 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] # Split the AR hidden states into text / image conditions. The token ids that # drive the split are sourced from the model config (see _split_ar_conditions), # formerly supplied by the bespoke example via additional_information. Legacy # fallback: ar2dit may have already produced the split conditions. if "text_prompt_embeds" in info: text_cond = info["text_prompt_embeds"] image_cond = info["image_prompt_embeds"] else: text_cond, image_cond = self._split_ar_conditions( full_hidden_states=info["full_hidden_states"], full_token_ids=info["full_token_ids"], answer_start_index=int(info["answer_start_index"][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") if image_cond.shape[0] == 0: answer_token_ids = info.get("full_token_ids", [])[int(info.get("answer_start_index", [0])[0]) :] raise ValueError( "MammothModa2 AR stage produced no visual-token hidden states; " "the DiT stage requires at least one generated visual token. " f"Generated token ids: {answer_token_ids[:32]}" ) 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 # 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, ar_image_hidden_states=ar_image_embeds, ar_image_attention_mask=ar_image_attention_mask, 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)