diff --git a/recipes/MammothModa2/MammothModa2.md b/recipes/MammothModa2/MammothModa2.md index 02f8e30bf..26ec9ed1f 100644 --- a/recipes/MammothModa2/MammothModa2.md +++ b/recipes/MammothModa2/MammothModa2.md @@ -18,11 +18,18 @@ 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. +MammothModa2's DiT stage runs in the shared diffusion runtime in request mode. +The first integration intentionally supports one request and one image per +forward only (`max_num_seqs: 1`, `num_outputs_per_prompt: 1`). Request-level +batching, step execution, continuous batching, cache acceleration, +compilation, quantization, parallelism, and offload are not enabled by this +recipe. + +Image size, seed, guidance, and denoising steps use the standard diffusion +request fields. `cfg_range` remains a MammothModa2-specific `extra_body` +parameter. For compatibility, the runtime also accepts the former +`text_guidance_scale` and `num_inference_steps` keys in `extra_body`; when +present and non-null, those keys take precedence over the standard fields. ## References @@ -41,11 +48,11 @@ flags. Image size uses the standard `--height` / `--width` flags. ## Hardware Support -The default deploy 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*. +The default deploy config places both the AR and DiT stages on one GPU +(`devices: "0"`). Its committed `gpu_memory_utilization` split is 0.5 for +stage 0 and 0.3 for stage 1. The A800 validation section below also shows a +two-GPU placement with one stage per GPU for attributable timing and memory; +the measured results are summarized below. ## GPU @@ -102,12 +109,13 @@ not implied. > 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.) +### 1x NVIDIA A800 80GB + #### 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 +- Driver / runtime: NVIDIA CUDA environment with one A800 80 GB - vLLM version: Match the repository requirements for your checkout - vLLM-Omni version or commit: Use the commit you are deploying from @@ -120,8 +128,7 @@ hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-P ``` Run text-to-image with the shared offline example from the repository root. The -deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the -MammothModa2 generation parameters as a JSON object through `--extra-body`: +deploy config sets `trust_remote_code`, so no extra flag is needed: ```bash python examples/offline_inference/text_to_image/text_to_image.py \ @@ -130,23 +137,25 @@ python examples/offline_inference/text_to_image/text_to_image.py \ --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}' \ + --seed 42 \ + --guidance-scale 4.0 \ + --num-inference-steps 50 \ + --extra-body '{"cfg_range": [0.0, 1.0]}' \ --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 +The standard diffusion request fields are `height`, `width`, `seed`, +`guidance_scale`, and `num_inference_steps`; use their corresponding CLI flags +shown above. `--height` and `--width` must be multiples of 16. + +`cfg_range` is the only recommended MammothModa2 field in `--extra-body`; it +sets the relative step range `[start, end]` over which CFG is applied (default +`[0.0, 1.0]`). For compatibility, `text_guidance_scale` and +`num_inference_steps` remain accepted `extra_body` aliases and, when non-null, +take precedence over the standard request fields. Model extras are filtered +against the 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. +so unknown MammothModa2 extras may be dropped. Run text-to-text through the shared understanding example. It recognizes the MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`: @@ -177,7 +186,49 @@ ls -lh mammoth_t2i.png python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)" ``` -### 1x AMD MI300X, MammothModa2 Preview +### 2x NVIDIA A800 80GB validation + +Use one A800 per stage so AR and DiT memory and timing are attributable. The +per-stage override changes placement only; both stages remain single-rank. + +```bash +VLLM_LOGGING_LEVEL=DEBUG vllm serve ./MammothModa2-Preview --omni \ + --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \ + --stage-overrides '{"0":{"devices":"0"},"1":{"devices":"1"}}' \ + --port 8099 \ + --log-stats +``` + +Startup logs should identify stage 1 as `StageDiffusionClient` and resolve it +to `MammothModa2DiTPipeline`. `DiffusionEngine` step timing is a DEBUG-level, +per-request message, so it appears only after sending a text-to-image request +with `VLLM_LOGGING_LEVEL=DEBUG`; it is not a startup marker. Seeing the legacy +generation model runner for stage 1 is a failed migration. + +#### Migration benchmark + +The request-mode migration was checked on 2x NVIDIA A800 80GB PCIe with AR on +GPU 0 and DiT on GPU 1. Each revision ran one warmup followed by 10 serial +measured requests in the same initialized process. Both used BF16 eager mode, +1024x1024 output, 50 denoising steps, guidance scale 4.0, seed 42, and no +diffusion cache. The baseline was the pre-migration revision `caed3061`; the +candidate was `19de562a`. Lower latency is better. + +| Metric | Baseline p50 | Baseline p95 | Candidate p50 | Candidate p95 | +| --- | ---: | ---: | ---: | ---: | +| End-to-end latency | 105.36 s | 106.01 s | 104.94 s | 105.79 s | +| AR stage latency | 86.97 s | 87.61 s | 86.26 s | 87.11 s | +| DiT stage latency | 18.27 s | 18.41 s | 18.60 s | 18.62 s | + +Peak sampled device memory was 39,209 MiB on the AR GPU for both revisions. +The DiT GPU used 11,089 MiB for the baseline and 10,967 MiB for the candidate. +The candidate's shared runtime reported 372.02 ms p50 per denoising step and a +5.89 ms p50 AR-to-diffusion adapter time. All measured requests completed and +both revisions produced valid, prompt-aligned 1024x1024 RGB images. The small +latency differences are regression evidence, not a statistically significant +speedup claim. + +### 1x AMD MI300X, MammothModa2 Preview (pre-migration baseline) #### Environment diff --git a/tests/config/test_endpoint_policy.py b/tests/config/test_endpoint_policy.py index aeb1ffa51..fa42b80c7 100644 --- a/tests/config/test_endpoint_policy.py +++ b/tests/config/test_endpoint_policy.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for endpoint restrictions logic.""" import pytest @@ -44,6 +44,32 @@ def test_restricted_completions_returns_400(): assert body["error"]["type"] == "BadRequestError" +def test_restricted_image_edits_returns_400(): + """Ensure pipelines can reject image edits independently of generation.""" + app = FastAPI() + + @app.post("/v1/images/edits") + async def existing_handler(): + return {"ok": True} + + @app.post("/v1/images/generations") + async def generations_handler(): + return {"ok": True} + + restrictions = (EndpointRestriction(OmniServingCapability.IMAGE_EDITS, REJECTION_REASON),) + shutdown_unsupported_routes(app, restrictions) + + client = TestClient(app) + resp = client.post("/v1/images/edits") + assert resp.status_code == 400 + assert resp.json()["error"]["message"] == REJECTION_REASON + + # The restriction is endpoint-specific and must not remove generations. + resp = client.post("/v1/images/generations") + assert resp.status_code == 200 + assert resp.json() == {"ok": True} + + def test_unrestricted_completions_not_blocked(): """Ensure that if we don't shutdown any routes, the route stays.""" app = _make_app_with_server_route() diff --git a/tests/config/test_mammoth_moda2_shared_runtime.py b/tests/config/test_mammoth_moda2_shared_runtime.py new file mode 100644 index 000000000..f5c11b2a5 --- /dev/null +++ b/tests/config/test_mammoth_moda2_shared_runtime.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.config.endpoint_policy import OmniServingCapability +from vllm_omni.config.stage_config import StageExecutionType, load_deploy_config +from vllm_omni.diffusion.registry import _DIFFUSION_MODELS +from vllm_omni.model_executor.models.mammoth_moda2.pipeline import ( + MAMMOTH_MODA2_AR_PIPELINE, + MAMMOTH_MODA2_PIPELINE, +) +from vllm_omni.model_executor.models.registry import _OMNI_MODELS + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_mammothmoda2_generation_stage_uses_shared_diffusion_runtime() -> None: + ar_stage, dit_stage = MAMMOTH_MODA2_PIPELINE.stages + assert ar_stage.execution_type is StageExecutionType.LLM_AR + assert dit_stage.execution_type is StageExecutionType.DIFFUSION + assert dit_stage.model_arch == "MammothModa2DiTPipeline" + assert dit_stage.custom_process_input_func.endswith(".ar2diffusion") + assert dit_stage.omni_kv_config == {"need_recv_cache": False} + assert dit_stage.input_sources == (0,) + assert dit_stage.final_output is True + assert dit_stage.final_output_type == "image" + + +def test_mammothmoda2_dit_is_registered_only_with_diffusion_runtime() -> None: + assert _DIFFUSION_MODELS["MammothModa2DiTPipeline"] == ( + "mammoth_moda2", + "pipeline_mammothmoda2_dit", + "MammothModa2DiTPipeline", + ) + assert "MammothModa2DiTPipeline" not in _OMNI_MODELS + + +def test_mammothmoda2_deploy_does_not_shadow_request_sampling_fields() -> None: + deploy = load_deploy_config(get_deploy_config_path("mammoth_moda2.yaml")) + + assert deploy.stages[1].default_sampling_params is None + + +def test_mammothmoda2_rejects_unsupported_image_edits() -> None: + restrictions = { + restriction.capability: restriction.reason for restriction in MAMMOTH_MODA2_PIPELINE.endpoint_restrictions + } + + assert set(restrictions) == {OmniServingCapability.IMAGE_EDITS} + assert "img2img" in restrictions[OmniServingCapability.IMAGE_EDITS] + assert MAMMOTH_MODA2_AR_PIPELINE.endpoint_restrictions == () + + +def test_mammothmoda2_ar_only_topology_is_unchanged() -> None: + assert len(MAMMOTH_MODA2_AR_PIPELINE.stages) == 1 + stage = MAMMOTH_MODA2_AR_PIPELINE.stages[0] + assert stage.execution_type is StageExecutionType.LLM_AR + assert stage.final_output_type == "text" diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index c0ac38924..3d356ad91 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -95,6 +95,20 @@ def _from_pipeline_key( ) +def test_mammothmoda2_diffusion_stage_projects_native_backend_config() -> None: + config = _from_pipeline_key( + "mammoth_moda2", + cli_overrides={"model": "/models/MammothModa2-Preview"}, + ) + stage = config.stage_by_id(1) + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.diffusion_config.model_class_name == "MammothModa2DiTPipeline" + assert stage.diffusion_config.model == "/models/MammothModa2-Preview" + assert stage.diffusion_config.step_execution is False + assert stage.scheduler_config.max_num_seqs == 1 + assert stage.connector_config.omni_kv_config == {"need_recv_cache": False} + + def test_minimax_h3_text_encoder_tp_targets_only_structured_stage_zero() -> None: config = _from_pipeline_key( "minimax_h3_disaggregated", diff --git a/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py b/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py new file mode 100644 index 000000000..776ca5bc6 --- /dev/null +++ b/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py @@ -0,0 +1,392 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from dataclasses import dataclass + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, TransformerConfig +from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit +from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import ( + MammothModa2DiTPipeline, + _build_mammoth_config, + _root_weight_source, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _raw_config() -> dict: + return { + "model_type": "mammothmoda2", + "llm_config": { + "model_type": "mammothmoda2_qwen2_5_vl", + "text_config": { + "model_type": "mammothmoda2_qwen2_5_vl_text", + "hidden_size": 8, + "gen_vocab_start_index": 100, + }, + }, + "gen_vae_config": {"block_out_channels": [8, 8]}, + "gen_dit_config": {"hidden_size": 8, "in_channels": 4}, + } + + +def _od_config() -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model="/models/MammothModa2-Preview", + model_class_name="MammothModa2DiTPipeline", + tf_model_config=TransformerConfig.from_dict(_raw_config()), + ) + + +def test_build_mammoth_config_uses_shared_transformer_projection() -> None: + config = _build_mammoth_config(_od_config()) + assert config.model_type == "mammothmoda2" + assert config.get_text_config().hidden_size == 8 + assert config.gen_dit_config == {"hidden_size": 8, "in_channels": 4} + + +def test_build_mammoth_config_rejects_empty_shared_projection() -> None: + config = _od_config() + config.tf_model_config = TransformerConfig() + with pytest.raises(ValueError, match="root checkpoint config"): + _build_mammoth_config(config) + + +def test_root_weight_source_loads_combined_checkpoint_once() -> None: + source = _root_weight_source(_od_config()) + assert source.model_or_path == "/models/MammothModa2-Preview" + assert source.subfolder is None + assert source.prefix == "" + assert source.fall_back_to_pt is True + + +def test_root_weight_source_forwards_revision() -> None: + config = _od_config() + config.revision = "rev-7" + assert _root_weight_source(config).revision == "rev-7" + + +def test_pipeline_declares_native_components_and_single_request_mode_only() -> None: + assert MammothModa2DiTPipeline._dit_modules == ["gen_transformer"] + assert MammothModa2DiTPipeline._encoder_modules == ["gen_image_condition_refiner"] + assert MammothModa2DiTPipeline._vae_modules == ["gen_vae"] + assert MammothModa2DiTPipeline.supports_request_batch is False + assert MammothModa2DiTPipeline.supports_step_execution is False + + +def test_mammoth_postprocess_denormalizes_nonnegative_raw_vae_output() -> None: + factory = getattr(pipeline_mammothmoda2_dit, "get_mammoth_moda2_post_process_func", None) + assert factory is not None + + images = factory(_od_config())(torch.zeros(1, 3, 2, 2)) + + assert len(images) == 1 + assert images[0].getpixel((0, 0)) == (128, 128, 128) + + +def test_mammoth_postprocess_is_registered() -> None: + from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS + + assert _DIFFUSION_POST_PROCESS_FUNCS["MammothModa2DiTPipeline"] == "get_mammoth_moda2_post_process_func" + + +def test_root_weight_source_rejects_missing_model_path() -> None: + config = _od_config() + config.model = None + with pytest.raises(ValueError, match="model path"): + _root_weight_source(config) + + +def _pipeline_shell() -> MammothModa2DiTPipeline: + pipeline = object.__new__(MammothModa2DiTPipeline) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.config = _build_mammoth_config(_od_config()) + pipeline._llm_hidden_size = 8 + return pipeline + + +def _batch( + *, + request_id: str = "req-7", + prompt: object | None = None, + sampling: OmniDiffusionSamplingParams | None = None, +) -> DiffusionRequestBatch: + if prompt is None: + prompt = { + "prompt": "", + "height": 32, + "width": 48, + "additional_information": { + "full_hidden_states": torch.arange(32, dtype=torch.float32).reshape(4, 8), + "full_token_ids": [10, 11, 100, 101], + "answer_start_index": 2, + }, + } + if sampling is None: + sampling = OmniDiffusionSamplingParams( + height=32, + width=48, + seed=42, + guidance_scale=4.0, + num_inference_steps=7, + extra_args={"cfg_range": [0.2, 0.8]}, + ) + return DiffusionRequestBatch([OmniDiffusionRequest(prompt=prompt, sampling_params=sampling, request_id=request_id)]) + + +def test_parse_request_uses_standard_sampling_fields() -> None: + parsed = _pipeline_shell()._parse_request(_batch()) + assert parsed.request_id == "req-7" + assert (parsed.height, parsed.width) == (32, 48) + assert parsed.num_inference_steps == 7 + assert parsed.text_guidance_scale == 4.0 + assert parsed.cfg_range == (0.2, 0.8) + assert parsed.seed == 42 + assert parsed.answer_start_index == 2 + + +def test_parse_request_prefers_legacy_sampling_overrides() -> None: + sampling = OmniDiffusionSamplingParams( + guidance_scale=3.0, + num_inference_steps=5, + seed=9, + extra_args={"text_guidance_scale": 6.0, "num_inference_steps": 11, "cfg_range": [0.0, 0.5]}, + ) + parsed = _pipeline_shell()._parse_request(_batch(sampling=sampling)) + assert parsed.text_guidance_scale == 6.0 + assert parsed.num_inference_steps == 11 + assert parsed.cfg_range == (0.0, 0.5) + + +def test_parse_request_falls_back_to_request_level_sampling_values() -> None: + prompt = _batch().prompts[0] + prompt["additional_information"].update( + text_guidance_scale=[1.5], + num_inference_steps=[3], + cfg_range=[0.25, 0.75], + ) + + parsed = _pipeline_shell()._parse_request(_batch(prompt=prompt, sampling=OmniDiffusionSamplingParams())) + + assert parsed.text_guidance_scale == 1.5 + assert parsed.num_inference_steps == 3 + assert parsed.cfg_range == (0.25, 0.75) + + +def test_parse_request_standard_fields_precede_request_level_fallbacks() -> None: + prompt = _batch().prompts[0] + prompt["additional_information"].update( + text_guidance_scale=[1.5], + num_inference_steps=[3], + ) + sampling = OmniDiffusionSamplingParams(guidance_scale=4.0, num_inference_steps=7) + + parsed = _pipeline_shell()._parse_request(_batch(prompt=prompt, sampling=sampling)) + + assert parsed.text_guidance_scale == 4.0 + assert parsed.num_inference_steps == 7 + + +def test_parse_request_rejects_multiple_requests_and_outputs() -> None: + pipeline = _pipeline_shell() + batch = _batch() + batch.requests.append(_batch(request_id="req-8").requests[0]) + with pytest.raises(ValueError, match="exactly one request"): + pipeline._parse_request(batch) + with pytest.raises(ValueError, match="num_outputs_per_prompt=1"): + pipeline._parse_request(_batch(sampling=OmniDiffusionSamplingParams(num_outputs_per_prompt=2))) + + +@pytest.mark.parametrize( + ("cfg_range", "message"), + [ + ([0.5], "two values"), + ([-0.1, 0.5], "0 <= start <= end <= 1"), + ([0.7, 0.2], "0 <= start <= end <= 1"), + ([0.2, 1.1], "0 <= start <= end <= 1"), + ], +) +def test_parse_request_rejects_invalid_cfg_range(cfg_range, message) -> None: + batch = _batch(sampling=OmniDiffusionSamplingParams(extra_args={"cfg_range": cfg_range})) + with pytest.raises(ValueError, match=message): + _pipeline_shell()._parse_request(batch) + + +def test_parse_request_requires_real_ar_conditions() -> None: + with pytest.raises(ValueError, match="req-missing"): + _pipeline_shell()._parse_request(_batch(request_id="req-missing", prompt={"prompt": "draw a cat"})) + + +def test_parse_request_rejects_hidden_state_token_count_mismatch() -> None: + batch = _batch() + batch.prompts[0]["additional_information"]["full_hidden_states"] = torch.zeros(3, 8) + with pytest.raises(ValueError, match="hidden-state/token-count mismatch"): + _pipeline_shell()._parse_request(batch) + + +@pytest.mark.parametrize( + ("height", "width", "message"), + [ + (0, 32, "Invalid image size"), + (32, -1, "Invalid image size"), + (30, 32, "multiples of 16"), + (32, 31, "multiples of 16"), + ], +) +def test_parse_request_rejects_invalid_dimensions(height, width, message) -> None: + batch = _batch() + batch.prompts[0].update(height=height, width=width) + with pytest.raises(ValueError, match=message): + _pipeline_shell()._parse_request(batch) + + +def test_parse_request_rejects_explicit_zero_steps() -> None: + with pytest.raises(ValueError, match="num_inference_steps must be positive"): + _pipeline_shell()._parse_request(_batch(sampling=OmniDiffusionSamplingParams(num_inference_steps=0))) + + +@pytest.mark.parametrize(("standard_steps", "expected_steps"), [(7, 7), (None, 50)]) +def test_parse_request_falls_through_null_legacy_step_count(standard_steps, expected_steps) -> None: + sampling = OmniDiffusionSamplingParams(num_inference_steps=standard_steps, extra_args={"num_inference_steps": None}) + parsed = _pipeline_shell()._parse_request(_batch(sampling=sampling)) + assert parsed.num_inference_steps == expected_steps + + +@pytest.mark.parametrize(("standard_guidance", "expected_guidance"), [(4.0, 4.0), (None, 9.0)]) +def test_parse_request_falls_through_null_legacy_guidance(standard_guidance, expected_guidance) -> None: + sampling = OmniDiffusionSamplingParams(guidance_scale=standard_guidance, extra_args={"text_guidance_scale": None}) + parsed = _pipeline_shell()._parse_request(_batch(sampling=sampling)) + assert parsed.text_guidance_scale == expected_guidance + + +def test_parse_request_defaults_null_cfg_range() -> None: + sampling = OmniDiffusionSamplingParams(extra_args={"cfg_range": None}) + parsed = _pipeline_shell()._parse_request(_batch(sampling=sampling)) + assert parsed.cfg_range == (0.0, 1.0) + + +@pytest.mark.parametrize("generators", [[], [torch.Generator(), torch.Generator()]]) +def test_parse_request_rejects_invalid_generator_list_cardinality(generators) -> None: + sampling = OmniDiffusionSamplingParams(generator=generators) + with pytest.raises(ValueError, match="exactly one generator.*req-generators"): + _pipeline_shell()._parse_request(_batch(request_id="req-generators", sampling=sampling)) + + +def test_parse_request_synthesizes_dummy_ar_conditions() -> None: + batch = _batch( + request_id="dummy_req_id", + prompt={"prompt": "dummy run"}, + sampling=OmniDiffusionSamplingParams(height=512, width=512, seed=1, guidance_scale=0.0, num_inference_steps=2), + ) + parsed = _pipeline_shell()._parse_request(batch) + assert parsed.full_hidden_states.shape == (2, 8) + assert parsed.full_token_ids == [0, 100] + assert parsed.answer_start_index == 1 + + +@dataclass +class _FakeTransformerConfig: + in_channels: int = 4 + + +@dataclass +class _FakeTimeCaptionEmbed: + image_embedder: nn.Module | None = None + + +class _FakeTransformer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(1)) + self.config = _FakeTransformerConfig() + self.time_caption_embed = _FakeTimeCaptionEmbed() + self.calls = 0 + + def forward(self, *, hidden_states, **kwargs): + self.calls += 1 + return torch.zeros_like(hidden_states) + + +@dataclass +class _FakeVaeConfig: + scaling_factor: float | None = None + shift_factor: float | None = None + + +class _FakeVae(nn.Module): + def __init__(self) -> None: + super().__init__() + self.anchor = nn.Parameter(torch.zeros(1)) + self.config = _FakeVaeConfig() + + def decode(self, latents, return_dict=False): + assert return_dict is False + return (torch.zeros(1, 3, 32, 48, dtype=latents.dtype),) + + +class _FakeScheduler: + def __init__(self) -> None: + self.timesteps = torch.tensor([]) + self.requested_steps = None + + def set_timesteps(self, *, num_inference_steps, device, num_tokens): + self.requested_steps = num_inference_steps + self.timesteps = torch.arange(num_inference_steps, device=device, dtype=torch.float32) + + def step(self, model_pred, timestep, latents, return_dict=False): + assert return_dict is False + return (latents - model_pred,) + + +def test_forward_returns_diffusion_output_with_request_sampling(mocker) -> None: + pipeline = _pipeline_shell() + pipeline.gen_transformer = _FakeTransformer() + pipeline.gen_image_condition_refiner = None + pipeline.gen_vae = _FakeVae() + pipeline.gen_freqs_cis = torch.zeros(1) + scheduler = _FakeScheduler() + captured = {} + + def fake_randn_tensor(shape, *, generator, device, dtype): + captured["shape"] = shape + captured["seed"] = generator.initial_seed() + return torch.zeros(shape, device=device, dtype=dtype) + + module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit" + mocker.patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler) + mocker.patch(f"{module}.randn_tensor", side_effect=fake_randn_tensor) + result = pipeline.forward( + _batch( + sampling=OmniDiffusionSamplingParams( + height=32, width=48, seed=42, guidance_scale=1.0, num_inference_steps=2 + ) + ) + ) + + assert isinstance(result, DiffusionOutput) + assert result.output.shape == (1, 3, 32, 48) + assert captured["shape"] == (1, 4, 4, 6) + assert captured["seed"] == 42 + assert scheduler.requested_steps == 2 + assert pipeline.gen_transformer.calls == 2 + + +def test_forward_rejects_missing_visual_tokens_before_model_access() -> None: + prompt = { + "prompt": "", + "additional_information": { + "full_hidden_states": torch.zeros(3, 8), + "full_token_ids": [10, 11, 12], + "answer_start_index": 2, + }, + } + with pytest.raises(ValueError, match="no visual-token hidden states.*req-empty"): + _pipeline_shell().forward(_batch(request_id="req-empty", prompt=prompt)) diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py index faedd0e09..a5f7aac1d 100644 --- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py +++ b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py @@ -4,8 +4,9 @@ """ 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. +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`` @@ -20,13 +21,17 @@ 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 # --------------------------------------------------------------------------- @@ -43,7 +48,7 @@ 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. +# 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. @@ -98,17 +103,33 @@ def _sample_pixels(img_tensor: torch.Tensor) -> list[float]: return values -# --------------------------------------------------------------------------- -# End-to-end test -# --------------------------------------------------------------------------- -pytestmark = [ - pytest.mark.slow, - pytest.mark.diffusion, - pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), -] +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 -@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/7718") +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): """ @@ -116,9 +137,9 @@ def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner): 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``). + - 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"]) @@ -139,7 +160,14 @@ def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner): max_tokens=max(1, expected_grid_tokens + 1), detokenize=False, ) - dit_sampling = SamplingParams(temperature=0.0, max_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( @@ -155,9 +183,6 @@ def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner): "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, @@ -173,35 +198,21 @@ def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner): 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" + images = list(_iter_images(outputs)) + assert len(images) == 1, f"Expected exactly one image, got {len(images)}" - sampled = _sample_pixels(img_tensor) + image = images[0] + assert isinstance(image, Image.Image), f"Expected PIL image, got {type(image)}" + assert image.mode == "RGB" + assert image.size == (width, height) - 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}" + sampled = _sample_pixels(_pil_to_tensor(image)) - found_image = True - - assert found_image, "No image tensor found in pipeline output" + 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}" diff --git a/tests/model_executor/stage_input_processors/test_mammoth_moda2.py b/tests/model_executor/stage_input_processors/test_mammoth_moda2.py index 992908c36..c404b6d58 100644 --- a/tests/model_executor/stage_input_processors/test_mammoth_moda2.py +++ b/tests/model_executor/stage_input_processors/test_mammoth_moda2.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""CPU regression tests for MammothModa2's completed-AR to DiT bridge.""" +"""CPU regressions for MammothModa2's completed-AR to diffusion bridge.""" -from types import SimpleNamespace +from dataclasses import dataclass import pytest import torch @@ -16,90 +16,129 @@ from vllm_omni.engine.serialization import ( ) from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient from vllm_omni.model_executor.models.mammoth_moda2.pipeline import MAMMOTH_MODA2_PIPELINE -from vllm_omni.model_executor.stage_input_processors.mammoth_moda2 import ar2dit +from vllm_omni.model_executor.stage_input_processors.mammoth_moda2 import ar2diffusion pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -_PROMPT_TOKEN_IDS = [7, 8, 9] -_GENERATED_TOKEN_IDS = [101, 102, 103] +@dataclass +class _CompletionOutputStub: + cumulative_token_ids: list[int] + multimodal_output: dict[str, torch.Tensor] -def _hidden_states(rows: int = 5, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor: - return torch.arange(rows * 4, dtype=dtype).reshape(rows, 4) +@dataclass +class _AROutputStub: + request_id: str + prompt_token_ids: list[int] + outputs: list[_CompletionOutputStub] -def _ar_output(hidden_states: torch.Tensor | None = None) -> SimpleNamespace: - multimodal_output = {} if hidden_states is None else {"latent": hidden_states} - return SimpleNamespace( - request_id="mammoth-request", - prompt_token_ids=list(_PROMPT_TOKEN_IDS), - outputs=[ - SimpleNamespace( - cumulative_token_ids=list(_GENERATED_TOKEN_IDS), - multimodal_output=multimodal_output, - ) - ], +@dataclass +class _TextConfigStub: + gen_vocab_start_index: int + + +@dataclass +class _MammothConfigStub: + llm_config: _TextConfigStub + image_token_id: int + video_token_id: int + vision_start_token_id: int + vision_end_token_id: int + + +def _source_output(*, include_latent: bool = True) -> _AROutputStub: + multimodal_output = {"latent": torch.arange(32, dtype=torch.float32).reshape(4, 8)} if include_latent else {} + completion = _CompletionOutputStub( + cumulative_token_ids=[100, 101, 102], + multimodal_output=multimodal_output, + ) + return _AROutputStub( + request_id="req-7", + prompt_token_ids=[10, 11], + outputs=[completion], ) -def _prompt( - *, - target_h: int | None = None, - target_w: int | None = None, -) -> dict[str, object]: - mm_processor_kwargs: dict[str, int] = {} - if target_h is not None: - mm_processor_kwargs["target_h"] = target_h - if target_w is not None: - mm_processor_kwargs["target_w"] = target_w - return { - "additional_information": { - "image_height": [384], - "image_width": [640], - "text_guidance_scale": [4.5], - "cfg_range": [0.1, 0.9], - "num_inference_steps": [20], +def test_ar2diffusion_builds_one_prompt_with_raw_ar_conditions() -> None: + result = ar2diffusion( + [_source_output()], + {"prompt": "a cat", "mm_processor_kwargs": {"target_h": 512, "target_w": 768}}, + ) + assert not isinstance(result, list) + assert result["prompt"] == "" + assert result["height"] == 512 + assert result["width"] == 768 + info = result["additional_information"] + assert info["full_token_ids"] == [10, 11, 100, 101] + assert info["answer_start_index"] == 2 + torch.testing.assert_close( + info["full_hidden_states"], + torch.arange(32, dtype=torch.float32).reshape(4, 8), + ) + assert info["full_hidden_states"].is_contiguous() + + +def test_ar2diffusion_uses_prompt_dimension_fallbacks() -> None: + result = ar2diffusion( + [_source_output()], + {"additional_information": {"image_height": [256], "image_width": [384]}}, + ) + assert (result["height"], result["width"]) == (256, 384) + + +def test_ar2diffusion_preserves_request_level_sampling_fallbacks() -> None: + result = ar2diffusion( + [_source_output()], + { + "additional_information": { + "text_guidance_scale": [1.5], + "num_inference_steps": [3], + "cfg_range": [0.25, 0.75], + } }, - "mm_processor_kwargs": mm_processor_kwargs, - } + ) + + info = result["additional_information"] + assert info["text_guidance_scale"] == [1.5] + assert info["num_inference_steps"] == [3] + assert info["cfg_range"] == [0.25, 0.75] -def _assert_complete_dit_input(dit_inputs: list[dict[str, object]], hidden_states: torch.Tensor) -> None: - assert len(dit_inputs) == 1 - dit_input = dit_inputs[0] - assert dit_input["prompt_token_ids"] == [0] - - additional_information = dit_input["additional_information"] - assert isinstance(additional_information, dict) - full_hidden_states = additional_information["full_hidden_states"] - assert isinstance(full_hidden_states, torch.Tensor) - assert full_hidden_states.dtype == torch.bfloat16 - assert full_hidden_states.is_contiguous() - assert torch.equal(full_hidden_states, hidden_states) - assert additional_information["full_token_ids"] == [7, 8, 9, 101, 102] - assert additional_information["answer_start_index"] == [3] - assert additional_information["image_height"] == [384] - assert additional_information["image_width"] == [640] - assert additional_information["text_guidance_scale"] == [4.5] - assert additional_information["cfg_range"] == [0.1, 0.9] - assert additional_information["num_inference_steps"] == [20] +def test_ar2diffusion_unwraps_the_orchestrator_prompt_list() -> None: + result = ar2diffusion( + [_source_output()], + [{"mm_processor_kwargs": {"target_h": 640, "target_w": 960}}], + ) + assert (result["height"], result["width"]) == (640, 960) -def test_ar2dit_builds_one_complete_dit_input() -> None: - hidden_states = _hidden_states() +def test_ar2diffusion_rejects_multiple_source_requests() -> None: + with pytest.raises(ValueError, match="exactly one AR output"): + ar2diffusion([_source_output(), _source_output()], {}) - dit_inputs = ar2dit([_ar_output(hidden_states)], _prompt()) - _assert_complete_dit_input(dit_inputs, hidden_states) +def test_ar2diffusion_reports_missing_latent_with_request_id() -> None: + with pytest.raises(ValueError, match="req-7"): + ar2diffusion([_source_output(include_latent=False)], {}) + + +def test_ar2diffusion_rejects_hidden_state_length_mismatch() -> None: + source = _source_output() + source.outputs[0].multimodal_output["latent"] = torch.zeros(3, 8) + with pytest.raises(ValueError, match="Hidden states length mismatch"): + ar2diffusion([source], {}) @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) -def test_ar2dit_preserves_low_precision_through_engine_core_payload(dtype: torch.dtype) -> None: - hidden_states = _hidden_states(dtype=dtype) - dit_input = ar2dit([_ar_output(hidden_states)], _prompt())[0] +def test_ar2diffusion_preserves_low_precision_through_engine_core_payload(dtype: torch.dtype) -> None: + source = _source_output() + hidden_states = torch.arange(32, dtype=dtype).reshape(4, 8) + source.outputs[0].multimodal_output["latent"] = hidden_states + diffusion_input = ar2diffusion([source], {}) - wire_payload = serialize_additional_information(dit_input["additional_information"]) + wire_payload = serialize_additional_information(diffusion_input["additional_information"]) assert wire_payload is not None restored = deserialize_additional_information(wire_payload) restored_hidden_states = restored["full_hidden_states"] @@ -112,21 +151,19 @@ def test_ar2dit_preserves_low_precision_through_engine_core_payload(dtype: torch @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) def test_dit_condition_split_preserves_compact_transfer_dtype(dtype: torch.dtype) -> None: - # Construct only the config read by _split_ar_conditions. This keeps the - # test CPU-only and avoids loading the DiT/VAE weights. pipeline = object.__new__(MammothModa2DiTPipeline) object.__setattr__( pipeline, "config", - SimpleNamespace( - llm_config=SimpleNamespace(gen_vocab_start_index=100), + _MammothConfigStub( + llm_config=_TextConfigStub(gen_vocab_start_index=100), image_token_id=20, video_token_id=21, vision_start_token_id=22, vision_end_token_id=23, ), ) - hidden_states = _hidden_states(dtype=dtype) + hidden_states = torch.arange(20, dtype=dtype).reshape(5, 4) text_cond, image_cond = pipeline._split_ar_conditions( full_hidden_states=hidden_states, @@ -142,44 +179,29 @@ def test_dit_condition_split_preserves_compact_transfer_dtype(dtype: torch.dtype assert torch.equal(image_cond, hidden_states[[3, 4]]) -def test_ar2dit_prefers_processor_image_dimensions() -> None: - hidden_states = _hidden_states() - - dit_inputs = ar2dit([_ar_output(hidden_states)], _prompt(target_h=512, target_w=768)) - - additional_information = dit_inputs[0]["additional_information"] - assert additional_information["image_height"] == [512] - assert additional_information["image_width"] == [768] - - -def test_ar2dit_rejects_missing_latent_output() -> None: - with pytest.raises(ValueError, match="missing latent multimodal output"): - ar2dit([_ar_output()], _prompt()) - - -def test_ar2dit_rejects_hidden_token_row_mismatch() -> None: - with pytest.raises(AssertionError, match="Hidden states length mismatch"): - ar2dit([_ar_output(_hidden_states(rows=4))], _prompt()) - - def test_mammoth_pipeline_uses_standard_completed_ar_forwarding() -> None: stage0, stage1 = MAMMOTH_MODA2_PIPELINE.stages assert stage0.custom_process_next_stage_input_func is None - assert stage1.custom_process_input_func.endswith(".ar2dit") + assert stage1.custom_process_input_func.endswith(".ar2diffusion") assert stage1.sync_process_input_func is None assert stage1.requires_full_payload_input is False def test_stage_client_forwards_completed_ar_output_to_mammoth_adapter() -> None: - hidden_states = _hidden_states() + source = _source_output() client = object.__new__(StageEngineCoreClient) - client.custom_process_input_func = ar2dit + client.custom_process_input_func = ar2diffusion client.requires_multimodal_data = False - dit_inputs = client.process_engine_inputs( - [_ar_output(hidden_states)], - _prompt(), + diffusion_input = client.process_engine_inputs( + [source], + {"mm_processor_kwargs": {"target_h": 512, "target_w": 768}}, ) - _assert_complete_dit_input(dit_inputs, hidden_states) + assert diffusion_input["height"] == 512 + assert diffusion_input["width"] == 768 + info = diffusion_input["additional_information"] + assert info["full_token_ids"] == [10, 11, 100, 101] + assert info["answer_start_index"] == 2 + assert torch.equal(info["full_hidden_states"], source.outputs[0].multimodal_output["latent"]) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index ed266d734..f7f1855cc 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -28,6 +28,9 @@ EXCLUDED_MODELS = [ "GlmImagePipeline", "ZImagePipeline", "OvisImagePipeline", + # Requires upstream AR hidden states and token boundaries; covered by + # dedicated MammothModa2 bridge/pipeline tests and real-checkpoint E2E. + "MammothModa2DiTPipeline", "WanPipeline", "WanDMDPipeline", "WanVACEPipeline", diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index 603ca65ff..ec2d4a704 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -397,7 +397,7 @@ class TestLoadCustomFuncSelection(unittest.TestCase): def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self): incompatible_paths = [ "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav", - "vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit", + "vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2diffusion", "vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow", "vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion", ] diff --git a/vllm_omni/config/endpoint_policy.py b/vllm_omni/config/endpoint_policy.py index f0aca95eb..24dea97e2 100644 --- a/vllm_omni/config/endpoint_policy.py +++ b/vllm_omni/config/endpoint_policy.py @@ -24,6 +24,7 @@ class OmniServingCapability(Enum): """Serving capabilities that pipelines can shut down.""" COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"})) + IMAGE_EDITS = RouteTarget("/v1/images/edits", frozenset({"POST"})) @property def path(self) -> str: diff --git a/vllm_omni/deploy/mammoth_moda2.yaml b/vllm_omni/deploy/mammoth_moda2.yaml index afae38ee0..76027424c 100644 --- a/vllm_omni/deploy/mammoth_moda2.yaml +++ b/vllm_omni/deploy/mammoth_moda2.yaml @@ -20,8 +20,3 @@ stages: enforce_eager: true trust_remote_code: true enable_prefix_caching: false - default_sampling_params: - extra_args: - text_guidance_scale: 9.0 - cfg_range: [0.0, 1.0] - num_inference_steps: 50 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 c96e5548c..21977f488 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py @@ -3,19 +3,23 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any, ClassVar +from dataclasses import dataclass +from typing import ClassVar import torch +from diffusers.image_processor import VaeImageProcessor from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm -from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DModel @@ -25,6 +29,59 @@ from .schedulers import FlowMatchEulerDiscreteScheduler logger = init_logger(__name__) +def _first_request_value(value: object) -> object: + if isinstance(value, (list, tuple)): + return value[0] if value else None + return value + + +def get_mammoth_moda2_post_process_func( + _od_config: OmniDiffusionConfig, +): + image_processor = VaeImageProcessor() + + def post_process_func(images: torch.Tensor): + return image_processor.postprocess(images) + + return post_process_func + + +def _build_mammoth_config(od_config: OmniDiffusionConfig) -> Mammothmoda2Config: + raw_config = od_config.tf_model_config.to_dict() + if not raw_config: + raise ValueError("MammothModa2 diffusion stage requires the root checkpoint config") + return Mammothmoda2Config(**raw_config) + + +def _root_weight_source( + od_config: OmniDiffusionConfig, +) -> DiffusersPipelineLoader.ComponentSource: + if not od_config.model: + raise ValueError("MammothModa2 diffusion stage requires a model path") + return DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder=None, + revision=od_config.revision, + prefix="", + fall_back_to_pt=True, + ) + + +@dataclass(frozen=True) +class _MammothRequest: + request_id: str + full_hidden_states: torch.Tensor + full_token_ids: list[int] + answer_start_index: int + height: int + width: int + text_guidance_scale: float + cfg_range: tuple[float, float] + num_inference_steps: int + seed: int | None + generator: torch.Generator | list[torch.Generator] | None + + class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): """ MammothModa2 DiT + VAE generation stage (non-autoregressive). @@ -38,7 +95,8 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): _encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"] _vae_modules: ClassVar[list[str]] = ["gen_vae"] - have_multimodal_outputs = True + supports_request_batch = False + supports_step_execution = False # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage. hf_to_vllm_mapper = WeightsMapper( @@ -48,15 +106,13 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): } ) - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__() del prefix - - hf_config = vllm_config.model_config.hf_config - if not isinstance(hf_config, Mammothmoda2Config): - raise TypeError(f"Expected Mammothmoda2Config, got {type(hf_config)}") - - self.config = hf_config + self.od_config = od_config + self.device = get_local_device() + self.config = _build_mammoth_config(od_config) + self.weights_sources = [_root_weight_source(od_config)] # --- Build DiT / VAE modules (names must match checkpoint keys) --- if self.config.gen_vae_config is None or self.config.gen_dit_config is None: @@ -111,9 +167,6 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): theta=10000, ) - # vLLM PP interface compatibility - self.make_empty_intermediate_tensors = lambda: None - self._llm_hidden_size = llm_hidden_size def _reinit_caption_embedder(self, in_features: int) -> None: @@ -125,39 +178,118 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): nn.Linear(in_features, out_features, bias=True), ) - def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]: - if num_reqs <= 0: - raise ValueError(f"num_reqs must be positive, got {num_reqs}") - if num_reqs > 1: - raise NotImplementedError( - f"get_dummy_runtime_additional_information does not support num_reqs > 1, got {num_reqs}" - ) - text_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32) - image_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32) - negative_prompt_embeds = torch.zeros((0, self._llm_hidden_size), dtype=torch.float32) - info = { - "text_prompt_embeds": text_prompt_embeds, - "image_prompt_embeds": image_prompt_embeds, - "negative_prompt_embeds": negative_prompt_embeds, - "negative_prompt_attention_mask": [], - "image_height": [512], - "image_width": [512], - "text_guidance_scale": [1.0], - "cfg_range": [0.0, 1.0], - "num_inference_steps": [1], - } - return [info for _ in range(num_reqs)] + def _parse_request(self, req: DiffusionRequestBatch) -> _MammothRequest: + if req.num_reqs != 1: + raise ValueError("MammothModa2 diffusion requires exactly one request") + request_id = req.request_id + sampling = req.sampling_params + if sampling.num_outputs_per_prompt != 1: + raise ValueError(f"MammothModa2 requires num_outputs_per_prompt=1 for request {request_id}") - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - # DiT stage does not consume token embeddings; return a dummy tensor. + prompt = req.prompts[0] + prompt = prompt if isinstance(prompt, dict) else {} + info = prompt.get("additional_information") + if req.is_dummy_run(): + full_hidden_states = torch.zeros((2, self._llm_hidden_size), dtype=torch.float32, device="cpu") + full_token_ids = [0, int(self.config.llm_config.gen_vocab_start_index)] + answer_start_index = 1 + else: + if not isinstance(info, dict): + raise ValueError(f"Missing additional_information AR conditions for request {request_id}") + full_hidden_states = info.get("full_hidden_states") + full_token_ids = info.get("full_token_ids") + if not isinstance(full_hidden_states, torch.Tensor) or not isinstance(full_token_ids, list): + raise ValueError(f"Expected full_hidden_states tensor and full_token_ids list for request {request_id}") + try: + answer_start_index = int(info.get("answer_start_index")) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"Invalid answer_start_index for request {request_id}") from exc + if full_hidden_states.ndim != 2: + raise ValueError(f"Expected 2D full_hidden_states for request {request_id}") + if full_hidden_states.shape[0] != len(full_token_ids): + raise ValueError(f"AR hidden-state/token-count mismatch for request {request_id}") + if not 0 <= answer_start_index <= len(full_token_ids): + raise ValueError(f"answer_start_index outside token range for request {request_id}") + try: + full_token_ids = [int(token_id) for token_id in full_token_ids] + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"Invalid full_token_ids for request {request_id}") from exc + + dimensions = [] + for name in ("height", "width"): + value = DiffusionRequestBatch.get_prompt_field(prompt, name) + if value is None: + value = getattr(sampling, name) + if value is None: + value = 1024 + try: + dimensions.append(int(value)) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"Invalid image size {name}={value!r} for request {request_id}") from exc + height, width = dimensions + if height <= 0 or width <= 0: + raise ValueError(f"Invalid image size: {height}x{width} for request {request_id}") + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"Image size must be multiples of 16, got {height}x{width} for request {request_id}") + + request_info = info if isinstance(info, dict) else {} + extra_args = sampling.extra_args or {} + guidance = extra_args.get("text_guidance_scale") + if guidance is None: + guidance = sampling.guidance_scale if sampling.guidance_scale_provided else None + if guidance is None: + guidance = _first_request_value(request_info.get("text_guidance_scale")) + if guidance is None: + guidance = 9.0 try: - dtype = next(self.parameters()).dtype - except StopIteration: - dtype = torch.float32 - return torch.zeros( - (input_ids.numel(), self._llm_hidden_size), - device=input_ids.device, - dtype=dtype, + text_guidance_scale = float(guidance) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"Invalid text_guidance_scale for request {request_id}") from exc + raw_num_inference_steps = extra_args.get("num_inference_steps") + if raw_num_inference_steps is None: + raw_num_inference_steps = sampling.num_inference_steps + if raw_num_inference_steps is None: + raw_num_inference_steps = _first_request_value(request_info.get("num_inference_steps")) + if raw_num_inference_steps is None: + raw_num_inference_steps = 50 + try: + num_inference_steps = int(raw_num_inference_steps) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"Invalid num_inference_steps for request {request_id}") from exc + if num_inference_steps <= 0: + raise ValueError(f"num_inference_steps must be positive for request {request_id}") + cfg_range = extra_args.get("cfg_range") + if cfg_range is None: + cfg_range = request_info.get("cfg_range") + if cfg_range is None: + cfg_range = [0.0, 1.0] + if not isinstance(cfg_range, (list, tuple)) or len(cfg_range) != 2: + raise ValueError(f"cfg_range requires two values for request {request_id}") + try: + cfg_start, cfg_end = float(cfg_range[0]), float(cfg_range[1]) + except (TypeError, ValueError, OverflowError) as exc: + raise ValueError(f"cfg_range requires two values convertible to floats for request {request_id}") from exc + if not 0 <= cfg_start <= cfg_end <= 1: + raise ValueError(f"cfg_range must satisfy 0 <= start <= end <= 1 for request {request_id}") + + generator = sampling.generator + if isinstance(generator, list) and len(generator) != 1: + raise ValueError( + f"MammothModa2 single-output request mode requires exactly one generator for request {request_id}" + ) + + return _MammothRequest( + request_id=request_id, + full_hidden_states=full_hidden_states, + full_token_ids=full_token_ids, + answer_start_index=answer_start_index, + height=height, + width=width, + text_guidance_scale=text_guidance_scale, + cfg_range=(cfg_start, cfg_end), + num_inference_steps=num_inference_steps, + seed=sampling.seed, + generator=generator, ) def _split_ar_conditions( @@ -202,42 +334,19 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): return text_cond, image_cond @torch.inference_mode() - def forward( - self, - *, - inputs_embeds: torch.Tensor | None = None, - **kwargs: Any, # noqa: ARG002 - ) -> OmniOutput: - runtime_addi = kwargs.get("runtime_additional_information", None) - info = runtime_addi[0] - - # 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] - - # 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]), + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + request = self._parse_request(req) + text_cond, image_cond = self._split_ar_conditions( + full_hidden_states=request.full_hidden_states, + full_token_ids=request.full_token_ids, + answer_start_index=request.answer_start_index, + ) + if image_cond.shape[0] == 0: + answer_token_ids = request.full_token_ids[request.answer_start_index :] + raise ValueError( + "MammothModa2 AR stage produced no visual-token hidden states for " + f"request {request.request_id}; the DiT stage requires at least one generated visual token. " + f"Generated token ids: {answer_token_ids[:32]}" ) # Move to model device/dtype. @@ -247,22 +356,6 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): else: target_dtype = next(self.gen_transformer.parameters()).dtype - def _ensure_2d(x: torch.Tensor, name: str) -> torch.Tensor: - if x.ndim == 3 and x.shape[0] == 1: - x = x[0] - if x.ndim != 2: - raise ValueError(f"Expected {name} to be 2D [T,H], got shape={tuple(x.shape)}") - return x - - text_cond = _ensure_2d(text_cond, "text_prompt_embeds") - image_cond = _ensure_2d(image_cond, "image_prompt_embeds") - if image_cond.shape[0] == 0: - answer_token_ids = info.get("full_token_ids", [])[int(info.get("answer_start_index", [0])[0]) :] - raise ValueError( - "MammothModa2 AR stage produced no visual-token hidden states; " - "the DiT stage requires at least one generated visual token. " - f"Generated token ids: {answer_token_ids[:32]}" - ) text_cond = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous() image_cond = image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous() @@ -301,63 +394,38 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): ar_image_embeds = image_embeds ar_image_attention_mask = image_attention_mask - # Prepare negative prompt (for CFG). If none provided, fall back to unconditional. + # Empty unconditional prompt for classifier-free guidance. negative_prompt_embeds = None negative_prompt_attention_mask = None - if text_guidance_scale > 1.0: - if negative_cond is not None: - negative_cond = _ensure_2d(negative_cond, "negative_prompt_embeds") - negative_prompt_embeds = ( - negative_cond.to(device=model_device, dtype=target_dtype, non_blocking=True) - .contiguous() - .unsqueeze(0) - ) - if isinstance(negative_attention_mask, torch.Tensor): - neg_mask = negative_attention_mask - elif isinstance(negative_attention_mask, list): - neg_mask = torch.tensor(negative_attention_mask, dtype=torch.bool) - else: - neg_mask = None - if neg_mask is None: - negative_prompt_attention_mask = torch.ones( - (1, negative_prompt_embeds.shape[1]), - dtype=torch.bool, - device=negative_prompt_embeds.device, - ) - else: - neg_mask = neg_mask.to(device=negative_prompt_embeds.device, dtype=torch.bool) - if neg_mask.ndim == 1: - neg_mask = neg_mask.unsqueeze(0) - negative_prompt_attention_mask = neg_mask - else: - hidden_size = int(prompt_embeds.shape[-1]) - negative_prompt_embeds = torch.zeros( - (1, 0, hidden_size), - dtype=target_dtype, - device=prompt_embeds.device, - ) - negative_prompt_attention_mask = torch.zeros( - (1, 0), - dtype=torch.bool, - device=prompt_embeds.device, - ) + if request.text_guidance_scale > 1.0: + hidden_size = int(prompt_embeds.shape[-1]) + negative_prompt_embeds = torch.zeros( + (1, 0, hidden_size), + dtype=target_dtype, + device=prompt_embeds.device, + ) + negative_prompt_attention_mask = torch.zeros( + (1, 0), + dtype=torch.bool, + device=prompt_embeds.device, + ) + + generator = request.generator + if generator is None and request.seed is not None: + generator = torch.Generator(device=model_device).manual_seed(request.seed) # Output image size (px), passed from stage input processor. - height, width = image_hw - if height <= 0 or width <= 0: - raise ValueError(f"Invalid image size: {height}x{width}") - if height % 16 != 0 or width % 16 != 0: - raise ValueError(f"Image size must be multiples of 16, got {height}x{width}") + height, width = request.height, request.width vae_scale_factor = 16 latent_channels = int(self.gen_transformer.config.in_channels) shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor) - latents = randn_tensor(shape, device=prompt_embeds.device, dtype=prompt_embeds.dtype) + latents = randn_tensor(shape, generator=generator, device=prompt_embeds.device, dtype=prompt_embeds.dtype) scheduler = FlowMatchEulerDiscreteScheduler() scheduler.set_timesteps( - num_inference_steps=num_inference_steps, + num_inference_steps=request.num_inference_steps, device=prompt_embeds.device, num_tokens=latents.shape[-2] * latents.shape[-1], ) @@ -376,7 +444,9 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): ar_image_attention_mask=ar_image_attention_mask, freqs_cis=self.gen_freqs_cis, ) - guidance_scale = text_guidance_scale if cfg_range[0] <= i / total_steps <= cfg_range[1] else 1.0 + guidance_scale = ( + request.text_guidance_scale if request.cfg_range[0] <= i / total_steps <= request.cfg_range[1] else 1.0 + ) if guidance_scale > 1.0 and negative_prompt_embeds is not None: model_pred_uncond = self.gen_transformer( hidden_states=latents, @@ -397,14 +467,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): latents = latents + self.gen_vae.config.shift_factor image = self.gen_vae.decode(latents, return_dict=False)[0] - return OmniOutput( - text_hidden_states=inputs_embeds, # placeholder, not used by runner - multimodal_outputs=image, - intermediate_tensors=None, - ) - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: # noqa: ARG002 - return None + return DiffusionOutput(output=image) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index f94952b1d..fadb712df 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -56,6 +56,11 @@ _DIFFUSION_MODELS = { "pipeline_ovis_image", "OvisImagePipeline", ), + "MammothModa2DiTPipeline": ( + "mammoth_moda2", + "pipeline_mammothmoda2_dit", + "MammothModa2DiTPipeline", + ), "WanPipeline": ( "wan2_2", "pipeline_wan2_2", @@ -564,6 +569,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "GlmImagePipeline": "get_glm_image_post_process_func", "ZImagePipeline": "get_post_process_func", "OvisImagePipeline": "get_ovis_image_post_process_func", + "MammothModa2DiTPipeline": "get_mammoth_moda2_post_process_func", "BooguImagePipeline": "get_boogu_image_post_process_func", "BooguImageTurboPipeline": "get_boogu_image_post_process_func", "WanPipeline": "get_wan22_post_process_func", diff --git a/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py b/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py index 830c63aff..be2e256fb 100644 --- a/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py +++ b/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MammothModa2 pipeline topology (frozen). Stage 0: AR — multimodal understanding + latent generation @@ -8,6 +8,7 @@ Stage 1: DiT — latent → image For text/image understanding tasks (text output only), use MAMMOTH_MODA2_AR_PIPELINE. """ +from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability from vllm_omni.config.stage_config import ( PipelineConfig, StageExecutionType, @@ -21,6 +22,13 @@ MAMMOTH_MODA2_PIPELINE = PipelineConfig( default_deploy_config_name="mammoth_moda2.yaml", model_arch="MammothModa2ForConditionalGeneration", hf_architectures=("Mammothmoda2Model", "MammothModa2ForConditionalGeneration"), + endpoint_restrictions=( + EndpointRestriction( + OmniServingCapability.IMAGE_EDITS, + "MammothModa2 image editing (img2img) is not supported. " + "Use /v1/images/generations for text-to-image requests.", + ), + ), stages=( StagePipelineConfig( stage_id=0, @@ -35,14 +43,15 @@ MAMMOTH_MODA2_PIPELINE = PipelineConfig( StagePipelineConfig( stage_id=1, model_stage="dit", - execution_type=StageExecutionType.LLM_GENERATION, + execution_type=StageExecutionType.DIFFUSION, input_sources=(0,), final_output=True, final_output_type="image", owns_tokenizer=False, requires_multimodal_data=False, - engine_output_type="image", - custom_process_input_func=f"{_PROC}.ar2dit", + model_arch="MammothModa2DiTPipeline", + custom_process_input_func=f"{_PROC}.ar2diffusion", + omni_kv_config={"need_recv_cache": False}, ), ), ) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 1a519917c..d1638ae9d 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -142,11 +142,6 @@ _OMNI_MODELS = { "mammoth_moda2", "MammothModa2Qwen3ForCausalLM", ), - "MammothModa2DiTPipeline": ( - "mammoth_moda2", - "pipeline_mammothmoda2_dit", - "MammothModa2DiTPipeline", - ), "MammothModa2ForConditionalGeneration": ( "mammoth_moda2", "mammoth_moda2", 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 fab22b739..5972dfa02 100644 --- a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py +++ b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py @@ -1,21 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Stage input processor for MammothModa2 (AR -> DiT).""" +"""Stage input processor for MammothModa2 (AR -> diffusion).""" from collections.abc import Mapping from typing import Any -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"): @@ -27,93 +18,73 @@ def _as_dict(prompt: Any) -> dict[str, Any]: def _coerce_dim(value: Any, default: int) -> int: try: - iv = int(value) + resolved = int(value) except (TypeError, ValueError): return default - return iv if iv > 0 else default + return resolved if resolved > 0 else default -def ar2dit( +def ar2diffusion( source_outputs: list[Any], - prompts: OmniTokensPrompt | TextPrompt | list[OmniTokensPrompt | TextPrompt] | None = None, - _requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - """Convert AR stage outputs to DiT stage inputs.""" - ar_outputs = source_outputs - - # 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 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), + prompt: Any | None = None, + requires_multimodal_data: bool = False, +) -> dict[str, Any]: + del requires_multimodal_data + if len(source_outputs) != 1: + raise ValueError( + f"MammothModa2 request-mode diffusion expects exactly one AR output, got {len(source_outputs)}" ) - # 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] + ar_output = source_outputs[0] + if isinstance(prompt, list): + prompt = prompt[0] if prompt else {} + prompt_dict = _as_dict(prompt) + additional = prompt_dict.get("additional_information") or {} + mm_kwargs = prompt_dict.get("mm_processor_kwargs") or {} + height = _coerce_dim( + mm_kwargs.get("target_h"), + _coerce_dim((additional.get("image_height") or [None])[0], 1024), + ) + width = _coerce_dim( + mm_kwargs.get("target_w"), + _coerce_dim((additional.get("image_width") or [None])[0], 1024), + ) - prompt_token_ids = ar_output.prompt_token_ids - # exclude the last token because it has no corresponding hidden state - completion_output = ar_output.outputs[0] - gen_token_ids = completion_output.cumulative_token_ids[:-1] - full_token_ids = prompt_token_ids + gen_token_ids - - mm_output = getattr(completion_output, "multimodal_output", None) - if not isinstance(mm_output, Mapping) or "latent" not in mm_output: - raise ValueError( - "AR stage output missing latent multimodal output. " - f"request_id={getattr(ar_output, 'request_id', None)}, " - f"completion_has_mm={hasattr(completion_output, 'multimodal_output')}" - ) - full_hidden_states = mm_output["latent"] - hidden_total = int(full_hidden_states.shape[0]) - assert hidden_total == len(prompt_token_ids) + len(gen_token_ids), ( - f"Hidden states length mismatch: expected {len(prompt_token_ids) + len(gen_token_ids)}, got {hidden_total}" + completion = ar_output.outputs[0] + generated_token_ids = list(completion.cumulative_token_ids[:-1]) + prompt_token_ids = list(ar_output.prompt_token_ids) + full_token_ids = prompt_token_ids + generated_token_ids + multimodal_output = getattr(completion, "multimodal_output", None) + if not isinstance(multimodal_output, Mapping) or "latent" not in multimodal_output: + raise ValueError( + "MammothModa2 AR stage output is missing latent multimodal output; " + f"request_id={getattr(ar_output, 'request_id', None)}" ) - # 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 = { - # The EngineCore payload serializer preserves bfloat16 as raw bytes. - # Keep the AR representation compact across the CPU/IPC boundary; - # the DiT selects its conditioning rows and casts them only when it - # consumes them. + full_hidden_states = multimodal_output["latent"] + hidden_total = int(full_hidden_states.shape[0]) + if hidden_total != len(full_token_ids): + raise ValueError( + "Hidden states length mismatch: " + f"expected {len(full_token_ids)}, got {hidden_total}; " + f"request_id={getattr(ar_output, 'request_id', None)}" + ) + + return { + "prompt": "", + "height": height, + "width": width, + "additional_information": { + # Keep #7102's compact BF16/FP16 payload across the EngineCore + # boundary. The diffusion pipeline casts only after selecting the + # text/image condition rows. "full_hidden_states": full_hidden_states.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)], - "cfg_range": [float(cfg_range[0]), float(cfg_range[1])], - "num_inference_steps": [int(num_inference_steps)], - } - - dit_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[0], - additional_information=additional_information, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) - - return dit_inputs + "answer_start_index": len(prompt_token_ids), + **{ + key: additional[key] + for key in ("text_guidance_scale", "num_inference_steps", "cfg_range") + if key in additional + }, + }, + }