[Model] Migrate MammothModa2 DiT to the shared diffusion runtime (#7134)

Signed-off-by: levius <2114377220@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>

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
 

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()

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"

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

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))

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

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

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

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

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:

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

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)

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

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},
         ),
     ),
 )

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

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
+            },
+        },
+    }