vllm-omni/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

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

"""
End-to-end test for MammothModa2 text-to-image generation.

Verifies that the AR->DiT pipeline produces a valid PIL image. When an
optional golden fixture is present, the test also compares fixed pixel samples;
the portable golden/stability oracle remains tracked in #7090.

Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
Deploy 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 numpy as np
import pytest
import torch
from PIL import Image
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
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.transformers_utils.repo_utils import hf_api

# ---------------------------------------------------------------------------
# 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_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")

_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)

# Optional 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(hf_api().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


def _iter_images(outputs: list[object]):
    """Yield images from shared diffusion ``OmniRequestOutput`` objects."""
    for out in outputs:
        ro_list = out if isinstance(out, list) else [out]
        for ro in ro_list:
            images = getattr(ro, "images", None)
            if isinstance(images, list):
                yield from images


def _pil_to_tensor(image: Image.Image) -> torch.Tensor:
    array = np.asarray(image, dtype=np.float32) / 255.0
    return torch.from_numpy(array).permute(2, 0, 1)


@pytest.mark.cpu
def test_diffusion_output_exposes_images_at_top_level():
    image = Image.new("RGB", (16, 16), "black")
    output = OmniRequestOutput.from_diffusion(request_id="diffusion-test", images=[image])

    assert output.outputs == []
    assert list(_iter_images([output])) == [image]


@pytest.mark.slow
@pytest.mark.diffusion
@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
@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 one RGB PIL image at the requested size.
      - When the optional fixture exists, fixed pixel samples match its 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 = OmniDiffusionSamplingParams(
        height=height,
        width=width,
        seed=42,
        guidance_scale=1.0,
        num_inference_steps=2,
        extra_args={"cfg_range": [0.0, 1.0]},
    )

    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],
                        "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"

    images = list(_iter_images(outputs))
    assert len(images) == 1, f"Expected exactly one image, got {len(images)}"

    image = images[0]
    assert isinstance(image, Image.Image), f"Expected PIL image, got {type(image)}"
    assert image.mode == "RGB"
    assert image.size == (width, height)

    sampled = _sample_pixels(_pil_to_tensor(image))

    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}"