""" End-to-end test for MammothModa2 text-to-image generation. Verifies that the AR->DiT pipeline produces an image tensor whose pixel values match a golden reference. Model Hub repo id: ``bytedance-research/MammothModa2-Preview``. Stage config: ``get_deploy_config_path("mammoth_moda2.yaml")`` → ``vllm_omni/deploy/mammoth_moda2.yaml`` Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json`` Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py`` """ from __future__ import annotations import json import os from pathlib import Path import pytest import torch from huggingface_hub import snapshot_download from vllm.sampling_params import SamplingParams from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- _IMAGE_TOKEN_ID = 151655 # "<|image_pad|>" _VIDEO_TOKEN_ID = 151656 # "<|video_pad|>" _VISION_START_TOKEN_ID = 151652 # "<|vision_start|>" _VISION_END_TOKEN_ID = 151653 # "<|vision_end|>" _AR_PATCH_SIZE = 16 MODEL_PATH = "bytedance-research/MammothModa2-Preview" T2I_STAGE_CONFIG = get_deploy_config_path("mammoth_moda2.yaml") _OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_STAGE_CONFIG) # Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate. _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json" # Fixed sampling coordinates: (channel, row_fraction, col_fraction) # Covers corners, centre, and mid-edges across all 3 channels. _PIXEL_SAMPLE_COORDS = [ (0, 0.0, 0.0), (0, 0.5, 0.5), (0, 1.0, 1.0), (0, 0.25, 0.75), (1, 0.0, 1.0), (1, 0.5, 0.0), (1, 0.75, 0.25), (1, 1.0, 0.5), (2, 0.0, 0.5), (2, 0.5, 1.0), (2, 0.75, 0.75), (2, 1.0, 0.0), ] # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _load_t2i_gen_config(repo_id: str) -> dict: weights_dir = Path(snapshot_download(repo_id)) cfg_path = weights_dir / "t2i_generation_config.json" if not cfg_path.exists(): pytest.skip(f"t2i_generation_config.json not found at {cfg_path}") with cfg_path.open() as f: return json.load(f) def _format_t2i_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str: 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 _sample_pixels(img_tensor: torch.Tensor) -> list[float]: """Sample pixel values at fixed fractional coordinates from a (C, H, W) tensor.""" t = img_tensor.float().clamp(0.0, 1.0) if t.ndim == 4: t = t[0] # unbatch C, H, W = t.shape values = [] for c, rh, rw in _PIXEL_SAMPLE_COORDS: ri = min(int(rh * (H - 1)), H - 1) ci = min(int(rw * (W - 1)), W - 1) values.append(round(float(t[c, ri, ci]), 6)) return values # --------------------------------------------------------------------------- # End-to-end test # --------------------------------------------------------------------------- pytestmark = [ pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] @pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3201") @hardware_test(res={"cuda": "H100"}) def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner): """ End-to-end text-to-image generation with MammothModa2 (AR -> DiT). Verifies: - Omni pipeline initialises with the two-stage YAML config. - DiT stage outputs an image tensor with the correct shape. - A fixed set of pixel values matches a golden reference (regenerate with ``UPDATE_GOLDEN=1``). """ gen_cfg = _load_t2i_gen_config(MODEL_PATH) eol_token_id = int(gen_cfg["eol_token_id"]) visual_start = int(gen_cfg["visual_token_start_id"]) visual_end = int(gen_cfg["visual_token_end_id"]) height, width = 256, 256 # small for CI speed ar_height, ar_width = height // _AR_PATCH_SIZE, width // _AR_PATCH_SIZE expected_grid_tokens = ar_height * (ar_width + 1) prompt_text = "A cat sitting on a laptop keyboard" formatted_prompt = _format_t2i_prompt(prompt_text, ar_width, ar_height) omni = omni_runner.omni ar_sampling = SamplingParams( temperature=0.0, top_k=1, max_tokens=max(1, expected_grid_tokens + 1), detokenize=False, ) dit_sampling = SamplingParams(temperature=0.0, max_tokens=1, detokenize=False) outputs = list( omni.generate( [ { "prompt": formatted_prompt, "additional_information": { "omni_task": ["t2i"], "ar_width": [ar_width], "ar_height": [ar_height], "eol_token_id": [eol_token_id], "visual_token_start_id": [visual_start], "visual_token_end_id": [visual_end], "image_height": [height], "image_width": [width], "num_inference_steps": [2], "text_guidance_scale": [1.0], "cfg_range": [0.0, 1.0], "visual_ids": [ _IMAGE_TOKEN_ID, _VIDEO_TOKEN_ID, _VISION_START_TOKEN_ID, _VISION_END_TOKEN_ID, ], }, } ], [ar_sampling, dit_sampling], ) ) assert len(outputs) > 0, "Pipeline produced no outputs" found_image = False for out in outputs: ro_list = getattr(out, "request_output", out) if not isinstance(ro_list, list): ro_list = [ro_list] for ro in ro_list: completion_outputs = getattr(ro, "outputs", None) if not isinstance(completion_outputs, list): continue for completion in completion_outputs: mm = getattr(completion, "multimodal_output", None) if not (isinstance(mm, dict) and "image" in mm): continue img_list = mm["image"] if isinstance(mm["image"], list) else [mm["image"]] for img_tensor in img_list: assert isinstance(img_tensor, torch.Tensor), f"Expected image tensor, got {type(img_tensor)}" assert img_tensor.ndim in (3, 4), f"Expected 3D or 4D image tensor, got {img_tensor.ndim}D" sampled = _sample_pixels(img_tensor) if os.environ.get("UPDATE_GOLDEN"): _GOLDEN_T2I_PATH.parent.mkdir(parents=True, exist_ok=True) _GOLDEN_T2I_PATH.write_text(json.dumps({"pixels": sampled}, indent=2)) print(f"\nGolden file written to {_GOLDEN_T2I_PATH}") elif _GOLDEN_T2I_PATH.exists(): golden = json.loads(_GOLDEN_T2I_PATH.read_text())["pixels"] for i, (got, exp) in enumerate(zip(sampled, golden)): assert abs(got - exp) < 1e-4, f"Pixel {i} mismatch: got {got}, expected {exp}" found_image = True assert found_image, "No image tensor found in pipeline output"