anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Model] Migrate MammothModa2-Preview example to standard task example + model_extras (#4691)

Browse Source 
Signed-off-by: mo saif <moisf56@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Mohammed Saif AbedandClaude authored and GitHub committed 2026-06-30 04:34:06 +00:00
1 parent 11f6a0b72e
commit 2fc00efa49
12 files changed
+415 -326
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/user_guide/examples/offline_inference/mammothmodal2_preview.md
+5 -16
View File
@@ -7,22 +7,15 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
Download model
```bash
hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
hf download 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
```
Text-to-image now runs through the shared offline image example
(`examples/offline_inference/text_to_image/text_to_image.py`). See the recipe
`recipes/MammothModa2/MammothModa2-Preview.md` for the full command and the
`extra_body` knobs (`text_guidance_scale`, `cfg_range`, `num_inference_steps`).
### Image Summary
@@ -40,7 +33,3 @@ python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_s
``````py
--8<-- "examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py"
``````
??? abstract "run_mammothmoda2_t2i.py"
``````py
--8<-- "examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py"
``````
examples/offline_inference/mammothmodal2_preview/README.md
+9 -16
View File
@@ -1,24 +1,17 @@
# MammothModa2-Preview
## Run examples (MammothModa2-Preview)
> 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
examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py
-248
View File
@@ -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 ``<base>_0<ext>``, ``<base>_1<ext>``, …
"""
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()
examples/offline_inference/text_to_image/text_to_image.py
+37
View File
@@ -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()
recipes/MammothModa2/MammothModa2-Preview.md
+110
View File
@@ -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)"
```
recipes/README.md
+1
View File
@@ -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 |
tests/model_extras/test_model_extras.py
+42
View File
@@ -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],
},
}
vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py
+63 -5
View File
@@ -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
vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml
+9
View File
@@ -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
vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py
+55 -41
View File
@@ -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)],
vllm_omni/model_extras/mammothmodal2_preview.py
+70
View File
@@ -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
vllm_omni/model_extras/registry.py
+14
View File
@@ -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,
},
}