[Fix] Restore MammothModa2 support
Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
This commit is contained in:
hsliuustc0106 committed
1 parent
c6df39e9f7
commit
fc5bafb2af
35 files changed
+4671
-13
No files matched your search
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
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``.
|
||||
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 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_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
|
||||
|
||||
_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_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.slow,
|
||||
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 = 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"
|
||||
Reference in new issue
Block a user