diff --git a/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md b/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md index 85227b748..6ab31298c 100644 --- a/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md +++ b/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md @@ -7,22 +7,15 @@ Source Text-to-image (T2I) generation has moved to the shared offline image example. +> See the recipe: [`recipes/MammothModa2/MammothModa2-Preview.md`](../../../recipes/MammothModa2/MammothModa2-Preview.md). +> +> This directory retains the image-understanding (image → text) example below, +> which is outside the scope of the diffusion `model_extras` example migration +> (#4548 / RFC #4539: AR / Omni understanding examples keep dedicated scripts). + +## Run example (MammothModa2-Preview) Download model ```bash -hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview -``` - -### Text-to-Image (T2I) - -```bash -python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \ - --model ./MammothModa2-Preview \ - --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \ - --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ - --height 1024 \ - --width 1024 \ - --num-inference-steps 50 \ - --text-guidance-scale 4.0 \ - --out output.png +hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview ``` ### Image Summary diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py deleted file mode 100644 index a4c41fee1..000000000 --- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py +++ /dev/null @@ -1,248 +0,0 @@ -""" -Offline inference example for MammothModa2 Text-to-Image (T2I) generation. -This script uses the vllm_omni.Omni pipeline with a multi-stage configuration. - -Workflow: -1. Stage 0 (AR): Generates visual tokens and their corresponding hidden states. -2. Stage 1 (DiT): Consumes the hidden states as conditions to perform diffusion - and VAE decoding to produce the final image. - -Example Usage: - uv run python examples/offline_inference/run_mammothmoda2_t2i.py \ - --model path/to/MammothModa2-Preview \ - --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \ - --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ - --out output.png -""" - -from __future__ import annotations - -import argparse -import json -import logging -import os -from pathlib import Path -from typing import NamedTuple - -import torch -from PIL import Image -from vllm.sampling_params import SamplingParams - -from vllm_omni import Omni - -logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") -logger = logging.getLogger(__name__) - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- -_PATCH_SIZE = 16 # AR image grid patch size (pixels per token) - - -class T2IGenConfig(NamedTuple): - eol_token_id: int - visual_token_start_id: int - visual_token_end_id: int - top_k: int # AR sampling top-k (covers the full visual generation vocabulary) - # Qwen2.5-VL special vision tokens: <|image_pad|>, <|video_pad|>, <|vision_start|>, <|vision_end|> - visual_ids: list[int] - - -def load_t2i_generation_config(model_dir: str) -> T2IGenConfig: - """Load T2I token IDs from t2i_generation_config.json and config.json. - - Supports both local directory paths and HuggingFace Hub model IDs. - """ - model_path = Path(model_dir) - - def _read_json(filename: str) -> dict: - local = model_path / filename - if local.exists(): - with local.open(encoding="utf-8") as f: - return json.load(f) - # Fall back to HuggingFace Hub when model_dir is a repo ID. - try: - from huggingface_hub import hf_hub_download - except ImportError as exc: - raise ImportError( - "huggingface_hub is required to load configs from HF Hub. Install it with: pip install huggingface_hub" - ) from exc - cached = hf_hub_download(repo_id=model_dir, filename=filename) - with open(cached, encoding="utf-8") as f: - return json.load(f) - - gen_cfg = _read_json("t2i_generation_config.json") - llm_cfg = _read_json("config.json").get("llm_config", {}) - - return T2IGenConfig( - eol_token_id=int(gen_cfg["eol_token_id"]), - visual_token_start_id=int(gen_cfg["visual_token_start_id"]), - visual_token_end_id=int(gen_cfg["visual_token_end_id"]), - top_k=int(gen_cfg["top_k"]), - visual_ids=[ - int(llm_cfg["image_token_id"]), - int(llm_cfg["video_token_id"]), - int(llm_cfg["vision_start_token_id"]), - int(llm_cfg["vision_end_token_id"]), - ], - ) - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description="Run MammothModa2 T2I (AR -> DiT) with vLLM-Omni.") - p.add_argument("--model", type=str, required=True, help="Path to the model directory.") - p.add_argument("--stage-config", type=str, required=True, help="Path to the multi-stage YAML configuration.") - p.add_argument( - "--prompt", - type=str, - action="append", - default=None, - help=( - "Text prompt for image generation. Can be provided multiple times " - "to generate multiple images with shared height/width/CFG settings." - ), - ) - p.add_argument("--height", type=int, default=1024, help="Output image height (must be a multiple of 16).") - p.add_argument("--width", type=int, default=1024, help="Output image width (must be a multiple of 16).") - p.add_argument("--num-inference-steps", type=int, default=50, help="Number of diffusion steps for the DiT stage.") - p.add_argument( - "--text-guidance-scale", type=float, default=9.0, help="Classifier-Free Guidance (CFG) scale for DiT." - ) - p.add_argument( - "--cfg-range", - type=float, - nargs=2, - default=(0.0, 1.0), - help="Relative step range [start, end] where CFG is active.", - ) - p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.") - p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.") - args = p.parse_args() - if not args.prompt: - args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."] - return args - - -def tensor_to_pil(image: torch.Tensor) -> Image.Image: - """Convert a normalized torch tensor [-1, 1] to a PIL Image.""" - if image.ndim == 4: - image = image[0] - image = image.detach().to("cpu") - image = (image / 2 + 0.5).clamp(0, 1) - image = (image * 255).to(torch.uint8) - image = image.permute(1, 2, 0).contiguous().numpy() - return Image.fromarray(image) - - -def _format_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str: - """Build the AR-stage prompt string including the image grid header.""" - return ( - "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n" - f"<|im_start|>user\n{user_prompt}<|im_end|>\n" - "<|im_start|>assistant\n" - f"<|image start|>{ar_width}*{ar_height}<|image token|>" - ) - - -def _collect_images(outputs: list) -> list[torch.Tensor]: - """Extract all image tensors produced by the final (DiT) stage.""" - images: list[torch.Tensor] = [] - for out in outputs: - ro_item = getattr(out, "request_output", out) - for completion in getattr(ro_item, "outputs", None) or []: - mm = getattr(completion, "multimodal_output", None) - if not isinstance(mm, dict) or "image" not in mm: - raise RuntimeError(f"Missing image in multimodal output: {mm}") - payload = mm["image"] - for tensor in payload if isinstance(payload, list) else [payload]: - if not isinstance(tensor, torch.Tensor): - raise TypeError(f"Expected image tensor, got {type(tensor)}") - images.append(tensor) - return images - - -def _save_images(images: list[torch.Tensor], out_path: str) -> list[str]: - """Save image tensors to disk. - - Single image: written to *out_path* exactly. - Multiple images: suffixed as ``_0``, ``_1``, … - """ - if not images: - raise RuntimeError("No images to save.") - base, ext = os.path.splitext(out_path) - ext = ext or ".png" - paths = [] - for i, tensor in enumerate(images): - path = out_path if len(images) == 1 else f"{base}_{i}{ext}" - tensor_to_pil(tensor).save(path) - paths.append(path) - return paths - - -def main() -> None: - args = parse_args() - os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) - - if args.height <= 0 or args.width <= 0: - raise ValueError(f"Height and width must be positive, got {args.height}x{args.width}") - if args.height % _PATCH_SIZE != 0 or args.width % _PATCH_SIZE != 0: - raise ValueError(f"Height and width must be multiples of {_PATCH_SIZE}, got {args.height}x{args.width}") - - ar_height = args.height // _PATCH_SIZE - ar_width = args.width // _PATCH_SIZE - gen_cfg = load_t2i_generation_config(args.model) - expected_grid_tokens = ar_height * (ar_width + 1) - - logger.info("Initializing Omni pipeline...") - omni = Omni(model=args.model, stage_configs_path=args.stage_config, trust_remote_code=args.trust_remote_code) - try: - ar_sampling = SamplingParams( - temperature=1.0, - top_p=1.0, - top_k=gen_cfg.top_k, - max_tokens=max(1, expected_grid_tokens + 1), # +1 for hidden state of eoi - detokenize=False, - ) - dit_sampling = SamplingParams( - temperature=0.0, - top_p=1.0, - top_k=-1, - max_tokens=1, - detokenize=False, - ) - - additional_information = { - "omni_task": ["t2i"], - "ar_width": [ar_width], - "ar_height": [ar_height], - "eol_token_id": [gen_cfg.eol_token_id], - "visual_token_start_id": [gen_cfg.visual_token_start_id], - "visual_token_end_id": [gen_cfg.visual_token_end_id], - "image_height": [args.height], - "image_width": [args.width], - "num_inference_steps": [args.num_inference_steps], - "text_guidance_scale": [args.text_guidance_scale], - "cfg_range": [args.cfg_range[0], args.cfg_range[1]], - "visual_ids": gen_cfg.visual_ids, - } - inputs = [ - { - "prompt": _format_prompt(p, ar_width, ar_height), - "additional_information": dict(additional_information), - } - for p in args.prompt - ] - - logger.info("Starting generation...") - # omni.generate() returns a Generator; consume it to run the full pipeline. - outputs = list(omni.generate(inputs, [ar_sampling, dit_sampling])) - - logger.info("Post-processing and saving image(s)...") - for path in _save_images(_collect_images(outputs), args.out): - logger.info(f"Saved: {path}") - finally: - omni.close() - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index ed51e85b6..69f9e4fcb 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -589,6 +589,12 @@ def main(): if images: break + # Fallback: generation-stage pipelines (e.g. MammothModa2's AR->DiT) return the + # generated image as a tensor under multimodal_output instead of populating the + # `images` field that diffusion-stage pipelines fill. + if not images: + images = _images_from_multimodal_output(outputs) + if not images: raise ValueError("No images found in request_output") @@ -606,5 +612,36 @@ def main(): print(f"Saved generated image to {save_path}") +def _images_from_multimodal_output(outputs: list[Any]) -> list[Any]: + """Extract PIL images from multimodal_output tensors. + + Generation-stage pipelines (e.g. MammothModa2's AR->DiT) return the generated + image as a tensor (normalized to [-1, 1], CHW) under ``multimodal_output`` + rather than populating the ``images`` field. Convert any such tensors to PIL. + """ + from PIL import Image + + pil_images: list[Any] = [] + for output in outputs: + req_out = getattr(output, "request_output", output) + for completion in getattr(req_out, "outputs", None) or []: + # multimodal_output is a MultimodalPayload (a Mapping) keyed by modality, + # matching how omni examples (ming_flash_omni / magi_human / dynin) read it. + mm = getattr(completion, "multimodal_output", None) or {} + if "image" not in mm: + continue + payload = mm["image"] + for tensor in payload if isinstance(payload, list) else [payload]: + if not isinstance(tensor, torch.Tensor): + continue + img = tensor.detach().to("cpu", dtype=torch.float32) + if img.ndim == 4: + img = img[0] + img = (img / 2 + 0.5).clamp(0, 1).mul(255).to(torch.uint8) + img = img.permute(1, 2, 0).contiguous().numpy() + pil_images.append(Image.fromarray(img)) + return pil_images + + if __name__ == "__main__": main() diff --git a/recipes/MammothModa2/MammothModa2-Preview.md b/recipes/MammothModa2/MammothModa2-Preview.md new file mode 100644 index 000000000..289efece0 --- /dev/null +++ b/recipes/MammothModa2/MammothModa2-Preview.md @@ -0,0 +1,110 @@ +# MammothModa2-Preview + +> MammothModa2-Preview text-to-image generation through the shared offline image example + +## Summary + +- Vendor: ByteDance Research +- Model: `bytedance-research/MammothModa2-Preview` +- Task: Text-to-image generation (AR → DiT two-stage pipeline) +- Mode: Offline inference +- Maintainer: Community + +## When to use this recipe + +Use this recipe to run MammothModa2-Preview text-to-image through the shared +offline image example (`text_to_image.py`) instead of a model-specific script. +The generic example formats the AR prompt, drives the AR → DiT stage pipeline, +and forwards MammothModa2-specific generation parameters through the +pipeline-declared `extra_body` contract. + +MammothModa2's DiT stage consumes its inputs through the multi-stage kwargs +interface (not `OmniDiffusionRequest`), so its generation knobs +(`text_guidance_scale`, `cfg_range`, `num_inference_steps`) are passed via +`--extra-body` rather than the standard `--num-inference-steps` / `--cfg-scale` +flags. Image size uses the standard `--height` / `--width` flags. + +## References + +- Upstream model: + [`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview) +- Related offline example: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) +- Declared parameters: + [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py) +- Stage config: + [`vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml`](../../vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml) + +## Hardware Support + +The default stage config runs both the AR and DiT stages on a single GPU +(`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`, +stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU +after rebalancing the split so the AR weights (~23 GB) leave room for the KV +cache — see the note under *1x L40S 48GB*. + +## GPU + +### 1x L40S 48GB + +> **48 GB config adjustment:** the committed +> `vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml` uses +> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S, +> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16` +> before running. (On an ~80 GB GPU, leave the defaults unchanged.) + +#### Environment + +- OS: Linux +- Python: Match the repository requirements for your checkout +- Driver / runtime: NVIDIA CUDA environment with one L40S 48 GB (verified) or an + ~80 GB GPU for the default config +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Offline Commands + +Download the model: + +```bash +hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview +``` + +Run text-to-image with the shared offline example from the repository root. The +stage config sets `trust_remote_code`, so no extra flag is needed. Forward the +MammothModa2 generation parameters as a JSON object through `--extra-body`: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model ./MammothModa2-Preview \ + --stage-configs-path vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \ + --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ + --height 1024 \ + --width 1024 \ + --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \ + --output mammoth_t2i.png +``` + +The `--extra-body` JSON forwards MammothModa2-specific parameters into +`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's +declared `extra_body_params` (see +[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)), +so unknown keys for MammothModa2 are silently dropped: + +- `text_guidance_scale` — classifier-free guidance scale for the DiT stage + (default `9.0`; CFG is active only when `> 1.0`). +- `cfg_range` — relative step range `[start, end]` over which CFG is applied + (default `[0.0, 1.0]`). +- `num_inference_steps` — number of DiT denoising steps (default `50`). + +`--height` and `--width` must be multiples of 16. + +#### Verification + +The example writes the generated image to the `--output` path. Confirm the file +exists and is a valid image: + +```bash +ls -lh mammoth_t2i.png +python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)" +``` diff --git a/recipes/README.md b/recipes/README.md index 63dd74174..77a78ef7d 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -46,6 +46,7 @@ recipes/ | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB | | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) | +| [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU | | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB | | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 | | [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 | diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 7eaf3edf6..c08e03f55 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -335,3 +335,45 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None: "cfg_text_scale": 4.0, "think": False, } + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("MammothModa2DiTPipeline") == frozenset( + { + "text_guidance_scale", + "cfg_range", + "num_inference_steps", + } + ) + assert get_extra_output_params("MammothModa2DiTPipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("MammothModa2DiTPipeline") is True + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_mammothmoda2_text_to_image_prompt_builder() -> None: + # Image dims are converted to the AR grid (width/16 x height/16); the negative + # prompt is ignored (MammothModa2 t2i uses CFG, not an explicit negative path). + assert build_text_to_image_prompt( + "MammothModa2DiTPipeline", + prompt="a cat", + negative_prompt="blurry", + height=512, + width=768, + ) == { + "prompt": ( + "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n" + "<|im_start|>user\na cat<|im_end|>\n" + "<|im_start|>assistant\n" + "<|image start|>48*32<|image token|>" + ), + "additional_information": { + "omni_task": ["t2i"], + "ar_width": [48], + "ar_height": [32], + "image_height": [512], + "image_width": [768], + }, + } diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index 39fbb92b4..0c312aa7e 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py @@ -142,6 +142,43 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): 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, @@ -151,14 +188,35 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): ) -> 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"] + + # 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] - 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] + + # 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 diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml index f4f523ac5..f5e449253 100644 --- a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml +++ b/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml @@ -36,6 +36,15 @@ stage_args: trust_remote_code: true engine_output_type: image enable_prefix_caching: false + # Declares extra_args so has_sampling_extra_args is enabled, routing the t2i knobs + # (text_guidance_scale, cfg_range, num_inference_steps) declared in + # vllm_omni/model_extras/mammothmodal2_preview.py from extra_body into the DiT + # forward. Values here are the defaults applied when a request omits a knob. + default_sampling_params: + extra_args: + text_guidance_scale: 9.0 + cfg_range: [0.0, 1.0] + num_inference_steps: 50 engine_input_source: [0] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit final_output: true diff --git a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py index 5ab4b6f48..7b23f8e49 100644 --- a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py +++ b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py @@ -3,12 +3,34 @@ from collections.abc import Mapping from typing import Any -import torch from vllm.inputs import TextPrompt from vllm_omni.inputs.data import OmniTokensPrompt +def _as_dict(prompt: Any) -> dict[str, Any]: + """Coerce an original-stage prompt to a dict. + + It may arrive as a dict, a NamedTuple/object, or a bare string depending on + the calling flow (the shared text_to_image example vs the bespoke script). + """ + if isinstance(prompt, dict): + return prompt + if hasattr(prompt, "_asdict"): + return prompt._asdict() + if hasattr(prompt, "__dict__"): + return vars(prompt) + return {} + + +def _coerce_dim(value: Any, default: int) -> int: + try: + iv = int(value) + except (TypeError, ValueError): + return default + return iv if iv > 0 else default + + def ar2dit( source_outputs: list[Any], prompts: OmniTokensPrompt | TextPrompt | list[OmniTokensPrompt | TextPrompt] | None = None, @@ -17,21 +39,33 @@ def ar2dit( """Convert AR stage outputs to DiT stage inputs.""" ar_outputs = source_outputs - # Normalize prompts to list + # The shared text_to_image example forwards a single prompt (not a list); normalize + # so a lone dict isn't iterated as its keys. Mirrors glm_image.ar2diffusion. if not isinstance(prompts, list): prompts = [prompts] if prompts is not None else [{}] dit_inputs: list[OmniTokensPrompt] = [] - for ar_output, prompt in zip(ar_outputs, prompts): - addi_info = prompt["additional_information"] - image_height = addi_info["image_height"][0] - image_width = addi_info["image_width"][0] - text_guidance_scale = addi_info["text_guidance_scale"][0] - cfg_range = addi_info["cfg_range"] - num_inference_steps = addi_info["num_inference_steps"][0] - gen_vocab_start_index = addi_info["visual_token_start_id"][0] - # ["<|image_pad|>", "<|video_pad|>", "<|vision_start|>", "<|vision_end|>"] - visual_ids = addi_info["visual_ids"] + for i, ar_output in enumerate(ar_outputs): + prompt_dict = _as_dict(prompts[i] if i < len(prompts) else {}) + addi_info = prompt_dict.get("additional_information") or {} + mm_kwargs = prompt_dict.get("mm_processor_kwargs") or {} + + # Image size: prefer mm_processor_kwargs target_h/target_w (set by the serving + # layer), fall back to additional_information, then a 1024 default. + image_height = _coerce_dim( + mm_kwargs.get("target_h"), + _coerce_dim((addi_info.get("image_height") or [None])[0], 1024), + ) + image_width = _coerce_dim( + mm_kwargs.get("target_w"), + _coerce_dim((addi_info.get("image_width") or [None])[0], 1024), + ) + + # Sampling knobs arrive on the DiT stage via extra_body -> extra_args; these are + # defensive fallbacks (defaults mirror the former bespoke script's argparse). + text_guidance_scale = (addi_info.get("text_guidance_scale") or [9.0])[0] + cfg_range = addi_info.get("cfg_range") or [0.0, 1.0] + num_inference_steps = (addi_info.get("num_inference_steps") or [50])[0] prompt_token_ids = ar_output.prompt_token_ids # exclude the last token because it has no corresponding hidden state @@ -52,36 +86,16 @@ def ar2dit( f"Hidden states length mismatch: expected {len(prompt_token_ids) + len(gen_token_ids)}, got {hidden_total}" ) - mask_device = full_hidden_states.device - full_token_ids_t = torch.tensor(full_token_ids, dtype=torch.long, device=mask_device) - attention_mask = torch.ones_like(full_token_ids_t, dtype=torch.bool) - - pos = torch.arange(full_token_ids_t.shape[0], device=mask_device) - answer_start_index = len(prompt_token_ids) - questions_mask = pos < answer_start_index - answers_mask = ~questions_mask - - gen_token_mask = full_token_ids_t >= gen_vocab_start_index - - visual_token_mask = torch.isin( - full_token_ids_t, - torch.tensor(visual_ids, dtype=torch.long, device=mask_device), - ) - - text_condition_token_mask = questions_mask & ~(visual_token_mask | gen_token_mask) & attention_mask - image_condition_token_mask = answers_mask & gen_token_mask & attention_mask - - text_condition = full_hidden_states[text_condition_token_mask] - image_condition = full_hidden_states[image_condition_token_mask] - - text_prompt_embeds = text_condition.to(dtype=torch.float32).contiguous() - image_prompt_embeds = image_condition.to(dtype=torch.float32).contiguous() - + # The text/image condition split is performed in the DiT pipeline, which sources + # the distinguishing token ids (gen_vocab_start_index, vision placeholder ids) + # from the model config. Pass through the raw AR hidden states + token ids and + # the question/answer boundary so the pipeline can reconstruct the masks. additional_information = { - "text_prompt_embeds": text_prompt_embeds, - "text_prompt_embeds_shape": list(text_prompt_embeds.shape), - "image_prompt_embeds": image_prompt_embeds, - "image_prompt_embeds_shape": list(image_prompt_embeds.shape), + # float32 so the tensor crosses the stage boundary (the serializer uses + # numpy, which has no bf16); the DiT re-casts to the model dtype. + "full_hidden_states": full_hidden_states.float().contiguous(), + "full_token_ids": full_token_ids, + "answer_start_index": [len(prompt_token_ids)], "image_height": [int(image_height)], "image_width": [int(image_width)], "text_guidance_scale": [float(text_guidance_scale)], diff --git a/vllm_omni/model_extras/mammothmodal2_preview.py b/vllm_omni/model_extras/mammothmodal2_preview.py new file mode 100644 index 000000000..c17e428a8 --- /dev/null +++ b/vllm_omni/model_extras/mammothmodal2_preview.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any + +# AR image-grid patch size (pixels per visual token). Mirrors the constant the +# bespoke MammothModa2 example used to convert image dimensions -> AR grid dims. +_PATCH_SIZE = 16 + +_AR_SYSTEM_PROMPT = "You are a helpful image generator." + + +def build_text_to_image_prompt( + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + """Build the MammothModa2 AR-stage prompt for text-to-image generation. + + Reproduces the prompt string and the structural ``additional_information`` + that the former bespoke example (``run_mammothmoda2_t2i.py``) constructed, + deriving the AR image grid from ``height`` / ``width``. + + Model-specific sampling knobs (``text_guidance_scale``, ``cfg_range``, + ``num_inference_steps``) flow separately via ``extra_body`` -> ``extra_args``. + Config-derived token ids (``eol_token_id``, ``visual_token_start_id``, + ``visual_token_end_id``, ``visual_ids``) are sourced inside the AR stage + rather than passed by the caller -- see the ar2dit stage input processor. + + MammothModa2 t2i uses classifier-free guidance via ``text_guidance_scale`` + and has no explicit negative-prompt path, so ``negative_prompt`` is accepted + for signature compatibility but not injected. + """ + h = height or 1024 + w = width or 1024 + ar_height = h // _PATCH_SIZE + ar_width = w // _PATCH_SIZE + + return { + "prompt": ( + f"<|im_start|>system\n{_AR_SYSTEM_PROMPT}<|im_end|>\n" + f"<|im_start|>user\n{prompt}<|im_end|>\n" + "<|im_start|>assistant\n" + f"<|image start|>{ar_width}*{ar_height}<|image token|>" + ), + "additional_information": { + "omni_task": ["t2i"], + "ar_width": [ar_width], + "ar_height": [ar_height], + "image_height": [h], + "image_width": [w], + }, + } + + +MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS = frozenset( + { + "text_guidance_scale", + "cfg_range", + # MammothModa2's DiT stage consumes inputs via the kwargs interface rather + # than OmniDiffusionRequest, so the standard --num-inference-steps flag does + # not reach it; it is routed through extra_body like the CFG knobs. + "num_inference_steps", + } +) +MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS = frozenset() +MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index 69608764a..eecb76522 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -34,6 +34,14 @@ from vllm_omni.model_extras.magi_human import ( MAGI_HUMAN_EXTRA_BODY_PARAMS, MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, ) +from vllm_omni.model_extras.mammothmodal2_preview import ( + MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS, + MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS, + MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, +) +from vllm_omni.model_extras.mammothmodal2_preview import ( + build_text_to_image_prompt as build_mammothmoda2_text_to_image_prompt, +) from vllm_omni.model_extras.sensenova_u1 import ( SENSENOVA_U1_EXTRA_BODY_PARAMS, SENSENOVA_U1_EXTRA_OUTPUT_PARAMS, @@ -143,6 +151,12 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = { "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS, "image_to_video_prompt_builder": build_vace_image_to_video_prompt, }, + "MammothModa2DiTPipeline": { + "extra_body_params": MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS, + "extra_output_params": MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS, + "init_extra_args_for_non_diffusion_stages": MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, + "text_to_image_prompt_builder": build_mammothmoda2_text_to_image_prompt, + }, }