Enable twine upload to PyPI in release pipeline

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.buildkite/test-nightly.yml

@@ -660,6 +660,87 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · GEBench Accuracy Test"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level full_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
+      - label: ":full_moon: Diffusion X2I(&A&T) · GEdit-Bench Accuracy Test"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level full_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
+          - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 1
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
+                        value: "120"
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
       - label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test"
         timeout_in_minutes: 180
         commands:

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.20.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.

docs/api/README.md

@@ -99,6 +99,7 @@ Configuration classes.
 - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
 - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
+- [vllm_omni.model_executor.models.cosyvoice3.config.CosyVoice3Config][]
 - [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechConfig][]
 - [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechFastARConfig][]
 - [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechSlowARConfig][]
@@ -115,14 +116,10 @@ Configuration classes.
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
-- [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
 - [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
-- [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
-- [vllm_omni.transformers_utils.configs.voxcpm.VoxCPMConfig][]
-- [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
 
 ## Workers
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
+uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
 ```
 
 #### Installation of vLLM-Omni
@@ -37,10 +37,10 @@ uv pip install onnxruntime-rocm
 If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
 
 ```bash
-uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
+uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
 ```
 
-The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.1. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source

docs/user_guide/diffusion_features.md

@@ -144,7 +144,7 @@ The following tables show which models support each feature:
 | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
 | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
 | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 

examples/offline_inference/text_to_speech/README.md

@@ -73,15 +73,13 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ```
 
 ### Voice cloning
-If `--ref-audio` is omitted, the script downloads the upstream
-[`zero_shot_prompt.wav`](https://github.com/FunAudioLLM/CosyVoice/blob/main/asset/zero_shot_prompt.wav) from the CosyVoice repo into the current directory.
-To use your own clip, pass `--ref-audio /path/to/reference.wav`, and modify `--prompt-text` correspondingly.
+Pass a reference audio. Note that CosyVoice3's `--prompt-text` is a system-style prompt for the GPT stage, not a reference transcript:
 ```bash
 python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
     --model pretrained_models/Fun-CosyVoice3-0.5B \
     --tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN \
-    --ref-audio /path/to/reference.wav \
-    --prompt-text "You are a helpful assistant.<|endofprompt|>Trascript in your ref audio clip"
+    --ref-audio prompt.wav \
+    --prompt-text "You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?"
 ```
 
 ### Streaming
@@ -319,39 +317,44 @@ pip install voxcpm soundfile
 export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
 ```
 
+If the native VoxCPM `config.json` does not contain HF metadata such as `model_type`, prepare a persistent HF-compatible config directory and point the stage configs to it via `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`:
+
+```bash
+export VOXCPM_MODEL=/path/to/voxcpm-model
+export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
+mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
+cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
+cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
+python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
+```
+
 ### Quick start
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model openbmb/VoxCPM-0.5B \
+    --model "$VOXCPM_MODEL" \
     --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
 ```
 
 ### Voice cloning
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model openbmb/VoxCPM-0.5B \
+    --model "$VOXCPM_MODEL" \
     --text "This sentence is synthesized with a cloned voice." \
     --ref-audio /path/to/reference.wav \
     --ref-text  "The exact transcript spoken in reference.wav."
 ```
 
 ### Streaming
-Pass `--streaming` together with the legacy async-chunk stage config:
+Pass the async-chunk stage config:
 ```bash
 python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
-    --model openbmb/VoxCPM-0.5B \
-    --streaming \
+    --model "$VOXCPM_MODEL" \
     --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
     --text "This is a split-stage VoxCPM streaming example running on vLLM Omni."
 ```
 
-### Persistent HF config (optional)
-The engine auto-patches a missing `model_type` in the checkpoint's `config.json` in-memory at startup (see `vllm_omni/engine/arg_utils.py`), so no manual setup is required.
-
-If you'd rather maintain a persistent patched config dir (e.g. to share across runs or to add fields the auto-patcher doesn't set), export `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` to its location — both the bundled deploy and async-chunk yamls read it via `${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,}`.
-
 ### Notes
-- Non-streaming auto-loads `vllm_omni/deploy/voxcpm.yaml` from the model directory name / HF `model_type`. Streaming uses the legacy `voxcpm_async_chunk.yaml` passed via `--stage-configs-path`.
+- `voxcpm.yaml` is the default non-streaming stage config; `voxcpm_async_chunk.yaml` enables streaming.
 - Streaming is currently single-request oriented.
 - `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
 - For online serving, see the [VoxCPM section in the online hub](../../online_serving/text_to_speech/README.md#voxcpm). For benchmark reporting, see [`benchmarks/voxcpm`](../../../benchmarks/voxcpm/README.md).

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -2,32 +2,18 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import argparse
 import os
-import urllib.request
-from pathlib import Path
 
 import numpy as np
 import soundfile as sf
 from vllm import SamplingParams
+from vllm.assets.audio import AudioAsset
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
-from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
-
-# Upstream zero-shot reference clip
-ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
-
-
-def _default_ref_audio() -> str:
-    # Download the upstream zero_shot_prompt.wav into the current dir
-    dest = Path("zero_shot_prompt.wav")
-    if not dest.exists() or dest.stat().st_size == 0:
-        print(f"Downloading default reference audio to {dest}")
-        urllib.request.urlretrieve(ZERO_SHOT_PROMPT_URL, dest)
-
-    return str(dest)
 
 
 def run_e2e():
@@ -50,15 +36,9 @@ def run_e2e():
     parser.add_argument(
         "--prompt-text",
         type=str,
-        default="You are a helpful assistant.<|endofprompt|>希望你以后，能够做的比我还好呦!",
-    )
-    parser.add_argument(
-        "--ref-audio",
-        type=str,
-        default=None,
-        help="Path to reference audio for voice cloning. "
-        "If unset, downloads the upstream CosyVoice3 zero-shot prompt audio clip",
+        default="You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?",
     )
+    parser.add_argument("--ref-audio", type=str, default="prompt.wav")
     parser.add_argument(
         "--tokenizer",
         type=str,
@@ -86,22 +66,24 @@ def run_e2e():
     sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
 
     print("Model initialized. Preparing inputs...")
-    ref_audio_path = args.ref_audio or _default_ref_audio()
-    if not os.path.exists(ref_audio_path):
-        raise FileNotFoundError(f"Audio file not found: {ref_audio_path}")
-    # Load at native sample rate
-    audio_signal, sr = load_audio(ref_audio_path, sr=None)
+    if args.ref_audio:
+        if not os.path.exists(args.ref_audio):
+            raise FileNotFoundError(f"Audio file not found: {args.ref_audio}")
+        # Load at native sample rate
+        audio_signal, sr = load_audio(args.ref_audio, sr=None)
 
-    # Validate sample rate before processing (similar to original CosyVoice)
-    min_sr = 16000
-    if sr < min_sr:
-        raise ValueError(
-            f"Audio sample rate {sr} Hz is too low. "
-            f"Minimum required: {min_sr} Hz. "
-            f"Please provide audio with sample rate >= {min_sr} Hz."
-        )
+        # Validate sample rate before processing (similar to original CosyVoice)
+        min_sr = 16000
+        if sr < min_sr:
+            raise ValueError(
+                f"Audio sample rate {sr} Hz is too low. "
+                f"Minimum required: {min_sr} Hz. "
+                f"Please provide audio with sample rate >= {min_sr} Hz."
+            )
 
-    audio_data = (audio_signal.astype(np.float32), sr)
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
 
     prompts = {
         "prompt": args.text,

examples/offline_inference/text_to_speech/voxcpm/end2end.py

@@ -11,9 +11,12 @@ import soundfile as sf
 import torch
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni import AsyncOmni, Omni
 from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
+DEFAULT_SYNC_STAGE_CONFIG = get_deploy_config_path("voxcpm.yaml")
+
 
 def _build_prompt(args) -> dict[str, Any]:
     additional_information: dict[str, list[Any]] = {
@@ -56,6 +59,10 @@ def _extract_sample_rate(mm: dict[str, Any]) -> int:
     return int(sr_raw)
 
 
+def _is_streaming_stage_config(stage_config_path: str) -> bool:
+    return "async_chunk" in Path(stage_config_path).stem
+
+
 def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path:
     output_dir.mkdir(parents=True, exist_ok=True)
     output_path = output_dir / f"output_{request_id}.wav"
@@ -80,7 +87,6 @@ async def _run_streaming(args) -> Path:
     omni = AsyncOmni(
         model=args.model,
         stage_configs_path=args.stage_configs_path,
-        deploy_config=args.deploy_config,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
@@ -127,7 +133,6 @@ def _run_sync(args) -> Path:
     omni = Omni(
         model=args.model,
         stage_configs_path=args.stage_configs_path,
-        deploy_config=args.deploy_config,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
     )
@@ -159,27 +164,11 @@ def _run_sync(args) -> Path:
 def parse_args():
     parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.")
     parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
-    parser.add_argument(
-        "--deploy-config",
-        type=str,
-        default=None,
-        help=("Override the deploy config path."),
-    )
     parser.add_argument(
         "--stage-configs-path",
         type=str,
-        default=None,
-        help=(
-            "Legacy stage_args yaml path. Required for streaming "
-            "(vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml); "
-            "leave unset for non-streaming to use the auto-resolved deploy config."
-        ),
-    )
-    parser.add_argument(
-        "--streaming",
-        action="store_true",
-        default=False,
-        help="Stream audio chunks as they arrive via AsyncOmni.",
+        default=DEFAULT_SYNC_STAGE_CONFIG,
+        help=("Stage config path. Use voxcpm.yaml for non-streaming or voxcpm_async_chunk.yaml for streaming."),
     )
     parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.")
     parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
@@ -205,10 +194,9 @@ def parse_args():
 
 
 def main(args) -> None:
-    route = "streaming" if args.streaming else "sync"
+    route = "streaming" if _is_streaming_stage_config(args.stage_configs_path) else "sync"
     print(f"Model: {args.model}")
-    print(f"Deploy config: {args.deploy_config or '<auto from HF model_type>'}")
-    print(f"Stage configs path: {args.stage_configs_path or '<unused>'}")
+    print(f"Stage config: {args.stage_configs_path}")
     print(f"Route: {route}")
     if route == "streaming":
         asyncio.run(_run_streaming(args))

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -5,7 +5,6 @@ from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
-from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
@@ -20,10 +19,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def mock_engine():
     with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock:
         engine_instance = MagicMock()
-        engine_instance.executor = SimpleNamespace(
-            register_failure_callback=MagicMock(),
-            check_health=MagicMock(),
-        )
         mock.make_engine.return_value = engine_instance
         yield engine_instance
 
@@ -102,41 +97,6 @@ def test_inline_shutdown(client, mock_engine):
     mock_engine.close.assert_called_once()
 
 
-def test_inline_registers_executor_failure_callback(client, mock_engine):
-    mock_engine.executor.register_failure_callback.assert_called_once()
-
-
-def test_inline_executor_failure_marks_engine_dead(client, mock_engine):
-    callback = mock_engine.executor.register_failure_callback.call_args.args[0]
-
-    callback()
-
-    assert client._engine_dead is True
-    with pytest.raises(EngineDeadError, match="inline diffusion engine is dead"):
-        client.get_diffusion_output_nowait()
-
-
-def test_inline_returns_queued_output_before_engine_dead(client, mock_engine):
-    callback = mock_engine.executor.register_failure_callback.call_args.args[0]
-    output = OmniRequestOutput.from_diffusion(request_id="req-queued", images=[])
-    client._output_queue.put_nowait(output)
-
-    callback()
-
-    assert client.get_diffusion_output_nowait() is output
-    with pytest.raises(EngineDeadError, match="inline diffusion engine is dead"):
-        client.get_diffusion_output_nowait()
-
-
-def test_inline_check_health_marks_engine_dead(client, mock_engine):
-    mock_engine.executor.check_health.side_effect = EngineDeadError()
-
-    with pytest.raises(EngineDeadError):
-        client.check_health()
-
-    assert client._engine_dead is True
-
-
 def test_inline_client_requires_replica_id(mock_engine):
     metadata = SimpleNamespace(
         stage_id=0,

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -700,232 +700,3 @@ def test_wire_round_trip_struct_to_dict_contract():
     assert set(decoded.keys()) == set(expected.keys())
     assert set(decoded["meta"].keys()) == set(expected["meta"].keys())
     assert set(decoded["codes"].keys()) == set(expected["codes"].keys())
-
-
-# ---------------------------------------------------------------
-# Deferred finish for upstream-completed requests
-# ---------------------------------------------------------------
-
-
-def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
-    """Build a mock scheduler with requests queued for deferred finish."""
-    adapter_mock = mocker.MagicMock()
-    adapter_mock.finished_requests = {r.request_id for r in pending_finish_reqs}
-    cleanup_calls = []
-    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
-
-    scheduler = mocker.MagicMock()
-    scheduler.chunk_transfer_adapter = adapter_mock
-    scheduler.connector = None
-    scheduler.ec_connector = None
-    scheduler.perf_metrics = None
-    scheduler.log_stats = False
-    scheduler.recompute_kv_load_failures = False
-    scheduler.structured_output_manager = mocker.MagicMock()
-    scheduler.structured_output_manager.should_advance.return_value = False
-    scheduler.finished_req_ids_dict = {}
-    scheduler.kv_cache_manager.take_events.return_value = None
-    scheduler.kv_event_publisher = mocker.MagicMock()
-    scheduler._pending_finish_reqs = list(pending_finish_reqs)
-
-    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
-    scheduler._free_request = mocker.MagicMock(return_value=None)
-    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
-    scheduler.running = list(running)
-    scheduler.waiting = mocker.MagicMock()
-    scheduler.waiting.remove_requests = mocker.MagicMock()
-    scheduler.make_stats = mocker.MagicMock(return_value=None)
-    scheduler.requests = {r.request_id: r for r in running}
-
-    scheduler_output = SimpleNamespace(
-        num_scheduled_tokens={},
-        scheduled_spec_decode_tokens={},
-        num_invalid_spec_tokens=0,
-    )
-    model_runner_output = SimpleNamespace(
-        sampled_token_ids=None,
-        logprobs=None,
-        prompt_logprobs_dict={},
-        pooler_output=None,
-        num_nans_in_logits=None,
-        kv_connector_output=None,
-        cudagraph_stats=None,
-        req_id_to_index={},
-    )
-    return scheduler, scheduler_output, model_runner_output, cleanup_calls
-
-
-def test_deferred_finish_emits_finished_output(mocker: MockerFixture):
-    """A request whose upstream completed with no remaining tokens should
-    emit a FINISHED EngineCoreOutput, free resources, and clean up adapter state."""
-    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
-
-    request = _HashableRequest(
-        request_id="req-df1",
-        external_req_id="ext-df1",
-        status=RequestStatus.RUNNING,
-        is_finished=lambda: False,
-        num_computed_tokens=16,
-        num_prompt_tokens=16,
-        prompt_token_ids=list(range(1, 17)),
-        num_output_placeholders=0,
-        sampling_params=None,
-        pooling_params=None,
-        stop_reason=None,
-        client_index=0,
-        take_events=lambda: [],
-        trace_headers=None,
-        has_encoder_inputs=False,
-        take_prefill_stats=lambda: None,
-        num_nans_in_logits=0,
-        get_finished_reason=lambda: "stop",
-    )
-    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
-        mocker,
-        running=[request],
-        pending_finish_reqs=[request],
-    )
-    scheduler._free_request.return_value = {"mock": "kv_params"}
-
-    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
-
-    assert request.status == RequestStatus.FINISHED_STOPPED
-    scheduler._handle_stopped_request.assert_called_once_with(request)
-    scheduler._free_request.assert_called_once_with(request)
-    assert len(cleanup_calls) == 1
-    assert cleanup_calls[0][0] == ("req-df1", "ext-df1")
-
-    eco = result[0]
-    assert len(eco.outputs) == 1
-    assert eco.outputs[0].request_id == "req-df1"
-    assert eco.outputs[0].finish_reason == "stop"
-    assert eco.outputs[0].kv_transfer_params == {"mock": "kv_params"}
-    assert scheduler._pending_finish_reqs == []
-
-
-def test_deferred_finish_empty_prompt(mocker: MockerFixture):
-    """A request that never received any tokens (finished immediately upstream)
-    should still emit a FINISHED output and clean up."""
-    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
-
-    request = _HashableRequest(
-        request_id="req-df2",
-        external_req_id="ext-df2",
-        status=RequestStatus.WAITING,
-        is_finished=lambda: False,
-        num_computed_tokens=0,
-        num_prompt_tokens=0,
-        prompt_token_ids=[],
-        num_output_placeholders=0,
-        sampling_params=None,
-        pooling_params=None,
-        stop_reason=None,
-        client_index=0,
-        take_events=lambda: [],
-        trace_headers=None,
-        has_encoder_inputs=False,
-        take_prefill_stats=lambda: None,
-        num_nans_in_logits=0,
-        get_finished_reason=lambda: "stop",
-    )
-    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
-        mocker,
-        running=[],
-        pending_finish_reqs=[request],
-    )
-
-    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
-
-    assert request.status == RequestStatus.FINISHED_STOPPED
-    scheduler._free_request.assert_called_once_with(request)
-    assert len(cleanup_calls) == 1
-    eco = result[0]
-    assert len(eco.outputs) == 1
-    assert eco.outputs[0].finish_reason == "stop"
-    assert scheduler._pending_finish_reqs == []
-
-
-def test_deferred_finish_skips_already_finished(mocker: MockerFixture):
-    """A request that was aborted between schedule() and update_from_output()
-    should be skipped without emitting output or freeing resources twice."""
-    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
-
-    request = _HashableRequest(
-        request_id="req-df3",
-        external_req_id="ext-df3",
-        status=RequestStatus.FINISHED_ABORTED,
-        is_finished=lambda: True,
-        num_computed_tokens=0,
-        num_prompt_tokens=0,
-        prompt_token_ids=[],
-        num_output_placeholders=0,
-        sampling_params=None,
-        pooling_params=None,
-        stop_reason=None,
-        client_index=0,
-        take_events=lambda: [],
-        trace_headers=None,
-        has_encoder_inputs=False,
-        take_prefill_stats=lambda: None,
-        num_nans_in_logits=0,
-        get_finished_reason=lambda: "stop",
-    )
-    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
-        mocker,
-        running=[],
-        pending_finish_reqs=[request],
-    )
-
-    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
-
-    scheduler._handle_stopped_request.assert_not_called()
-    scheduler._free_request.assert_not_called()
-    assert len(cleanup_calls) == 0
-    assert 0 not in result
-    assert scheduler._pending_finish_reqs == []
-
-
-def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture):
-    """When _handle_stopped_request returns False (resumable request), the
-    output must still be emitted so the client stream doesn't hang."""
-    from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
-
-    request = _HashableRequest(
-        request_id="req-df4",
-        external_req_id="ext-df4",
-        status=RequestStatus.RUNNING,
-        is_finished=lambda: False,
-        num_computed_tokens=16,
-        num_prompt_tokens=16,
-        prompt_token_ids=list(range(1, 17)),
-        num_output_placeholders=0,
-        sampling_params=None,
-        pooling_params=None,
-        stop_reason=None,
-        client_index=0,
-        take_events=lambda: [],
-        trace_headers=None,
-        has_encoder_inputs=False,
-        take_prefill_stats=lambda: None,
-        num_nans_in_logits=0,
-        get_finished_reason=lambda: "stop",
-    )
-    scheduler, sched_out, model_out, cleanup_calls = _build_deferred_finish_scheduler(
-        mocker,
-        running=[request],
-        pending_finish_reqs=[request],
-    )
-    scheduler._handle_stopped_request.return_value = False
-
-    result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
-
-    scheduler._handle_stopped_request.assert_called_once_with(request)
-    scheduler._free_request.assert_not_called()
-    assert len(cleanup_calls) == 0
-
-    eco = result[0]
-    assert len(eco.outputs) == 1
-    assert eco.outputs[0].request_id == "req-df4"
-    assert eco.outputs[0].finish_reason == "stop"
-    assert eco.outputs[0].kv_transfer_params is None
-    assert scheduler._pending_finish_reqs == []

tests/e2e/accuracy/conftest.py

@@ -185,12 +185,6 @@ def accuracy_artifact_root() -> Path:
     return root
 
 
-@pytest.fixture(scope="session")
-def accuracy_assets_root() -> Path:
-    root = Path(__file__).resolve().parent / "assets"
-    return root
-
-
 @pytest.fixture(scope="session")
 def qwen_bear_image(accuracy_artifact_root: Path):
     """Download the Qwen bear image from the URL and save it to the accuracy artifact root."""

tests/e2e/accuracy/helpers.py

@@ -1,11 +1,9 @@
 import os
-from io import BytesIO
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 import pytest
-import requests
 import torch
 from PIL import Image
 from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
@@ -127,74 +125,6 @@ def assert_image_sequence_similarity(
         )
 
 
-def assert_images_pixel_close(
-    *,
-    model_name: str,
-    vllm_image: Image.Image,
-    diffusers_image: Image.Image,
-    mean_threshold: float,
-    p99_threshold: float,
-    compare_mode: str = "RGB",
-) -> None:
-    """Assert full-image pixel closeness between online serving and diffusers output.
-
-    Match the threshold style used by diffusion parallelism tests: convert both
-    images to float32 RGB tensors in the [0, 1] range, then compare mean and p99
-    absolute channel differences.
-
-    Accuracy improves as both values move lower:
-    - ``mean_abs_diff`` tracks global average drift across the whole image.
-    - ``p99_abs_diff`` tracks tail drift while ignoring the noisiest 1% of
-      channel samples.
-
-    The printed mismatch ratios are diagnostics only. ``pixel_ratio`` counts
-    pixels where any channel exceeds a tolerance, while ``channel_ratio`` counts
-    individual channel samples above that tolerance.
-    """
-    assert vllm_image.size == diffusers_image.size, (
-        f"Online and diffusers output sizes mismatch: online={vllm_image.size}, diffusers={diffusers_image.size}"
-    )
-
-    vllm_array = np.asarray(vllm_image.convert(compare_mode), dtype=np.float32) / 255.0
-    diffusers_array = np.asarray(diffusers_image.convert(compare_mode), dtype=np.float32) / 255.0
-    channel_abs_diff = np.abs(vllm_array - diffusers_array)
-    mean_abs_diff = float(channel_abs_diff.mean())
-    p99_abs_diff = float(np.quantile(channel_abs_diff, 0.99))
-    percentiles = {
-        percentile: float(np.quantile(channel_abs_diff, percentile / 100.0)) for percentile in (50, 90, 95, 99, 99.9)
-    }
-
-    print(f"{model_name} pixel metrics:")
-    print(
-        f"  mean_abs_diff: value={mean_abs_diff:.6e}, threshold<={mean_threshold:.6e}, "
-        "range=[0, 1], lower_is_better=True"
-    )
-    print(
-        f"  p99_abs_diff: value={p99_abs_diff:.6e}, threshold<={p99_threshold:.6e}, range=[0, 1], lower_is_better=True"
-    )
-    print("  abs_diff_percentiles:")
-    for percentile, value in percentiles.items():
-        print(f"    p{percentile}: value={value:.6e}, range=[0, 1], lower_is_better=True")
-    print("  mismatch_ratios_by_channel_threshold:")
-    for tolerance in (0, 1, 2, 4, 8, 16, 32, 64, 128):
-        normalized_tolerance = tolerance / 255.0
-        pixel_mismatch = np.any(channel_abs_diff > normalized_tolerance, axis=-1)
-        pixel_mismatch_ratio = float(np.count_nonzero(pixel_mismatch) / pixel_mismatch.size)
-        channel_mismatch_ratio = float(
-            np.count_nonzero(channel_abs_diff > normalized_tolerance) / channel_abs_diff.size
-        )
-        print(
-            f"    threshold>{normalized_tolerance:.6e} ({tolerance}/255): pixel_ratio={pixel_mismatch_ratio:.8f}, "
-            f"channel_ratio={channel_mismatch_ratio:.8f}"
-        )
-
-    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
-        f"Image diff exceeded threshold for {model_name}: mean_abs_diff={mean_abs_diff:.6e}, "
-        f"p99_abs_diff={p99_abs_diff:.6e} (thresholds: mean<={mean_threshold:.6e}, "
-        f"p99<={p99_threshold:.6e})"
-    )
-
-
 def compute_image_ssim_psnr(
     *,
     prediction: Image.Image,
@@ -231,16 +161,6 @@ class CLIPScorer:
         similarity = (img_emb * txt_emb).sum(dim=-1)
         return float(similarity.item() * 100)
 
-    def image_image_score(self, image1: Image.Image, image2: Image.Image) -> float:
-        inputs = self._processor(text=[""], images=[image1, image2], return_tensors="pt", padding=True)
-        with torch.no_grad():
-            outputs = self._model(**inputs)
-        img_emb1, img_emb2 = outputs.image_embeds.chunk(2)
-        img_emb1 = img_emb1 / img_emb1.norm(p=2, dim=-1, keepdim=True)
-        img_emb2 = img_emb2 / img_emb2.norm(p=2, dim=-1, keepdim=True)
-        similarity = (img_emb1 * img_emb2).sum(dim=-1)
-        return float(similarity.item() * 100)
-
     def assert_score(self, *, model_name: str, image: Image.Image, text: str, threshold: float) -> None:
         value = self.score(image, text)
         print(f"{model_name} CLIP score:")
@@ -260,73 +180,3 @@ def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Te
     array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0
     tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
     return tensor
-
-
-class SemanticSimilarityScorer:
-    """Semantic similarity scorer for text-text and image-text comparisons"""
-
-    def __init__(
-        self,
-        text_model_name: str = "BAAI/bge-m3",
-    ):
-        """
-        Initialize semantic similarity scorer
-
-        Args:
-            text_model_name: Model for text-text similarity (BGE-M3)
-        """
-        self._text_model = None
-        self._text_model_name = text_model_name
-
-    def _load_text_model(self):
-        """Lazy load text similarity model (BGE-M3)"""
-        if self._text_model is None:
-            try:
-                from FlagEmbedding import BGEM3FlagModel
-
-                print(f"Loading text similarity model: {self._text_model_name}")
-                self._text_model = BGEM3FlagModel(self._text_model_name, use_fp16=True)
-            except ImportError:
-                raise ImportError("FlagEmbedding not installed. Install with: pip install FlagEmbedding")
-        return self._text_model
-
-    def text_similarity(self, text1: str, text2: str) -> float:
-        """
-        Compute semantic similarity between two texts using BGE-M3
-        """
-        results = {}
-        # 1. Text prefix matching (maximum matching character length)
-        match_count = 0
-        for char_vllm, char_ref in zip(text1, text2):
-            if char_vllm == char_ref:
-                match_count += 1
-            else:
-                break
-        results["text_prefix_match_count"] = match_count
-
-        # 2. Full CoT semantic similarity comparison
-        model = self._load_text_model()
-        emb1 = model.encode([text1], return_dense=True, return_sparse=False, return_colbert_vecs=False)["dense_vecs"][0]
-        emb2 = model.encode([text2], return_dense=True, return_sparse=False, return_colbert_vecs=False)["dense_vecs"][0]
-        similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
-        results["cot_semantic_sim"] = float(similarity)
-        return results
-
-    def release_models(self):
-        """Release loaded models to free memory"""
-        if self._text_model is not None:
-            del self._text_model
-            self._text_model = None
-
-
-def download_images(urls: list[str]) -> list[Image.Image]:
-    """Download and cache multiple images"""
-    images = []
-    for i, url in enumerate(urls):
-        print(f"Downloading image {i} from {url}")
-        response = requests.get(url, timeout=30)
-        response.raise_for_status()
-        image = Image.open(BytesIO(response.content)).convert("RGB")
-        images.append(image)
-
-    return images

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -5,8 +5,6 @@ from __future__ import annotations
 
 import copy
 import gc
-import importlib
-import os
 import tempfile
 import time
 from pathlib import Path
@@ -16,59 +14,28 @@ import torch
 import yaml
 from PIL import Image
 
-from tests.e2e.accuracy.helpers import (
-    CLIPScorer,
-    SemanticSimilarityScorer,
-    compute_image_ssim_psnr,
-    download_images,
-    model_output_dir,
-)
+from tests.e2e.accuracy.helpers import CLIPScorer, assert_similarity, model_output_dir
 from tests.helpers.runtime import OmniRunner
-from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
-
-os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
+from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens
 
 pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
 
-# ============================================================================
-# Configurable Parameters
-# ============================================================================
-AR_DEVICES = "0,1"
-DIT_DEVICES = "2,3"
 MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
-NUM_INFERENCE_STEPS = 50
-GUIDANCE_SCALE = 2.5
-
-# ============================================================================
-# Constants
-# ============================================================================
-MODEL_PATH = os.environ.get("HUNYUAN_MODEL_PATH", MODEL_NAME)
-# Test input
-PROMPT = "基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴"
-TEST_IMAGE_URLS = [
-    "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_0.png",
-    "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_1.png",
-]
 SEED = 42
-AR_TP_SIZE = len(AR_DEVICES.split(","))
-DIT_TP_SIZE = len(DIT_DEVICES.split(","))
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 0
+HEIGHT, WIDTH = 1024, 1024
+PSNR_THRESHOLD = 0.0
+SSIM_THRESHOLD = 0.0
+CLIP_SCORE_THRESHOLD = 20.0
+PROMPT = "A brown and white dog is running on the grass."
 
-# Precision thresholds
-THRESHOLDS = {
-    # AR text comparison
-    "text_prefix_match": 10,  # First 10 characters must match exactly
-    "cot_semantic_sim": 0.9,  # Full CoT semantic similarity
-    # Image comparison
-    "clip_score": 85,  # CLIP image semantic similarity
-    "ssim": 0.20,  # Structural similarity
-    "psnr": 11.0,  # Peak signal-to-noise ratio (dB)
-}
 # fmt: off
 _BASE_CONFIG = {
     "stage_args": [
         {
             "stage_id": 0, "stage_type": "llm",
-            "runtime": {"process": True, "devices": AR_DEVICES, "max_batch_size": 1, "requires_multimodal_data": True},
+            "runtime": {"process": True, "devices": "0,1", "max_batch_size": 1, "requires_multimodal_data": True},
             "engine_args": {
                 "model_stage": "AR", "model_arch": "HunyuanImage3ForCausalMM",
                 "worker_cls": "vllm_omni.worker.gpu_ar_worker.GPUARWorker",
@@ -83,11 +50,11 @@ _BASE_CONFIG = {
                 "temperature": 0.0, "top_p": 1, "top_k": -1, "max_tokens": 8192,
                 "stop_token_ids": [128025], "detokenize": True, "skip_special_tokens": False,
             },
-            "output_connectors": {"to_stage_1": "shared_memory_connector"},
+            "output_connectors": {"to_stage_1": "rdma_connector"},
         },
         {
             "stage_id": 1, "stage_type": "diffusion",
-            "runtime": {"process": True, "devices": DIT_DEVICES, "max_batch_size": 1, "requires_multimodal_data": True},
+            "runtime": {"process": True, "devices": "2,3", "max_batch_size": 1, "requires_multimodal_data": True},
             "engine_args": {
                 "model_stage": "dit", "model_arch": "HunyuanImage3ForCausalMM",
                 "enforce_eager": True, "trust_remote_code": True, "distributed_executor_backend": "mp",
@@ -97,32 +64,20 @@ _BASE_CONFIG = {
             "custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion",
             "final_output": True, "final_output_type": "image",
             "default_sampling_params": {"num_inference_steps": NUM_INFERENCE_STEPS, "guidance_scale": GUIDANCE_SCALE},
-            "input_connectors": {"from_stage_0": "shared_memory_connector"},
+            "input_connectors": {"from_stage_0": "rdma_connector"},
         },
     ],
     "runtime": {
         "enabled": True,
-        "connectors": {"shared_memory_connector": {
-            "name": "SharedMemoryConnector",
-            "extra": {"shm_threshold_bytes": 65536}
+        "connectors": {"rdma_connector": {
+            "name": "MooncakeTransferEngineConnector",
+            "extra": {"host": "auto", "zmq_port": 50051, "protocol": "rdma", "device_name": "", "memory_pool_size": 4294967296, "memory_pool_device": "cpu"},
         }},
         "edges": [{"from": 0, "to": 1}],
     },
 }
 # fmt: on
 
-# fmt: off
-COT_REF = ("首先，我分析所有输入图像：图像1是一个圆形的logo，设计现代且抽象。它由不同色调的蓝色（深蓝、中蓝、浅蓝）和白色构成，这些色块以流畅的曲线相互交织，形成一个动态的、类似旋涡或波浪的图案。整个logo是扁平化的矢量图形，背景为纯黑色。图像2展示了四个并排摆放的卡通动物造型冰箱贴，"
-           "它们被放置在灰色的织物背景上。这些冰箱贴的关键特征是其材质：它们具有光滑、高光的珐琅或烤漆质感，边缘有明显的金属包边，整体呈现出一种立体的、有厚度的实体感。用户的指令是“基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴”。这个指令要求将一个二维的平面设计（logo）"
-           "转化为一个具有特定物理属性（材质和立体感）的三维物体。核心任务是保留logo的视觉识别性，同时赋予其冰箱贴的实体质感。为了构建答案图像，我会将图一的圆形logo作为基础形状。然后，我会将图二中冰箱贴的材质特性应用到这个logo上。具体来说，logo中的每一个色块（深蓝、中蓝、浅蓝、白色）"
-           "都会被渲染成具有高光泽度的珐琅质感，表面会反射出柔和的环境光，形成自然的高光。logo中不同颜色区域之间的分界线，将被处理成纤细的、带有金属光泽的凸起边缘，这既能清晰地勾勒出图案，也符合珐琅工艺品的典型特征。整个冰箱贴会呈现出轻微的厚度和圆润的边缘，使其看起来像一个真实的、可触摸的物体。"
-           "最后，将这个制作完成的冰箱贴放置在图二所示的灰色织物背景上，并为其添加一个微妙的、柔和的阴影，以增强其立体感和与背景的融合度，最终呈现出一个精致、逼真的产品展示图。</think><recaption>这幅图像以产品摄影的精致风格，呈现了一枚根据`image_1`标志定制的圆形珐琅冰箱贴。最终图像使用`image_2`的分辨率。"
-           "冰箱贴居中放置在`image_2`的灰色织物背景上，其设计完美复刻了`image_1`中由深蓝、中蓝、浅蓝和白色构成的动态旋涡图案。整个冰箱贴被赋予了`image_2`中冰箱贴特有的高级质感：表面覆盖着一层光滑如镜的珐琅釉面，反射出柔和而清晰的高光；图案的每一个色块边缘都由纤细的抛光金属边框精确勾勒，增强了立体感。"
-           "柔和的顶光在冰箱贴的弧形边缘上形成平滑的过渡，并在其下方投下淡淡的、轮廓模糊的阴影，使其与织物背景无缝融合，营造出一种真实、静谧的视觉效果。<relation_1>最终图像完整保留了`image_1`中标志的全部设计元素。这包括其完美的圆形轮廓，以及内部由深蓝、中蓝、浅蓝和白色组成的精确旋涡状图案布局、形状和色彩关系。"
-           "</relation_1><relation_2>最终图像的分辨率、背景和材质均来自`image_2`。背景中灰色织物的纹理和质感被完整保留。冰箱贴的材质被完美重构，精确复刻了`image_2`中冰箱贴所展示的光滑珐琅质感、抛光金属边框的视觉效果，以及整体柔和、均匀的布光环境和由此产生的自然阴影。</relation_2></recaption><answer><boi>"
-           "<img_size_1024><img_ratio_36><timestep>[<img>]{3600}<eoi></answer>")
-# fmt: on
-
 
 def _make_config(enable_kv_reuse: bool, path: Path) -> None:
     config = copy.deepcopy(_BASE_CONFIG)
@@ -131,25 +86,24 @@ def _make_config(enable_kv_reuse: bool, path: Path) -> None:
     path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
 
 
-def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
+def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
     from transformers import AutoTokenizer
 
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
     from vllm_omni.platforms import current_omni_platform
 
-    build_kwargs: dict = {"task": "it2i", "bot_task": "think_recaption", "sys_type": "en_unified", "num_images": 2}
-
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
     result = build_prompt_tokens(
         PROMPT,
         tokenizer,
-        **build_kwargs,
+        task="it2i",
+        bot_task="recaption",
+        sys_type="en_unified",
     )
     token_ids = result.token_ids
     system_prompt_type = result.system_prompt_type
 
-    ar_stop_token_ids = resolve_stop_token_ids(task="it2i", bot_task="think_recaption", tokenizer=tokenizer)
-    with OmniRunner(MODEL_NAME, stage_configs_path=stage_configs_path) as runner:
+    with OmniRunner(MODEL_NAME, stage_configs_path=stage_config_path) as runner:
         params_list = list(runner.omni.default_sampling_params_list)
         for sp in params_list:
             if isinstance(sp, OmniDiffusionSamplingParams):
@@ -157,17 +111,14 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
                 sp.guidance_scale = GUIDANCE_SCALE
                 sp.seed = SEED
                 sp.generator = torch.Generator(device=current_omni_platform.device_type or "cuda").manual_seed(SEED)
-            elif hasattr(sp, "stop_token_ids"):
-                sp.stop_token_ids = ar_stop_token_ids
 
-        images = download_images(TEST_IMAGE_URLS)
         prompts: list[OmniPromptType] = [
             {
                 "prompt_token_ids": token_ids,
                 "prompt": PROMPT,
                 "use_system_prompt": system_prompt_type,
-                "modalities": ["image"],
-                "multi_modal_data": {"image": images},
+                "height": HEIGHT,
+                "width": WIDTH,
             }
         ]
         t0 = time.perf_counter()
@@ -180,13 +131,9 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
     for out in outputs:
         ro = getattr(out, "request_output", None)
         if ro and getattr(ro, "outputs", None):
-            cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
-        if not cot_text:
-            ar_text = getattr(out, "custom_output", {}).get("ar_generated_text")
-            if isinstance(ar_text, list):
-                cot_text = "\n".join(text for text in ar_text if text)
-            else:
-                cot_text = ar_text or ""
+            txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
+            if txt:
+                cot_text = txt
 
         imgs = getattr(out, "images", None)
         if not imgs and ro and hasattr(ro, "images"):
@@ -195,11 +142,9 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
             images = imgs
 
     assert images, "Pipeline output had no images"
-    cot_text = cot_text.lstrip("\n")
 
     image = images[0].convert("RGB")
-    image.save(output_path / "image_offline.png")
-    (output_path / "cot_offline.txt").write_text(cot_text, encoding="utf-8")
+    image.save(output_path)
     gc.collect()
     if torch.accelerator.is_available():
         torch.accelerator.empty_cache()
@@ -207,44 +152,71 @@ def _run_offline(stage_configs_path: str, output_path: Path) -> tuple[Image.Imag
 
 
 @pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Needs 4+ GPUs (2 AR + 2 DiT)")
-def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_root: Path) -> None:
-    if importlib.util.find_spec("FlagEmbedding") is None:
-        raise ImportError("Missing dependency: FlagEmbedding\nInstall with: pip install FlagEmbedding")
-    from tabulate import tabulate  # lazy import
-
-    """KV reuse ON vs OFF: same pipeline, same seed → PSNR >= 10 dB."""
-    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-offline-kv-reuse")
+def test_text_to_image_alignment(accuracy_artifact_root: Path) -> None:
+    """KV reuse ON vs OFF: same pipeline, same seed → PSNR >= 40 dB."""
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-kv-reuse")
 
     with tempfile.TemporaryDirectory() as tmpdir:
         tmp = Path(tmpdir)
+        _make_config(False, tmp / "off.yaml")
+        image_no_reuse_1, cot_no_reuse_1, time_no_reuse_1 = _run(
+            str(tmp / "off.yaml"), output_dir / "without_kv_reuse_1.png"
+        )
+        image_no_reuse_2, cot_no_reuse_2, time_no_reuse_2 = _run(
+            str(tmp / "off.yaml"), output_dir / "without_kv_reuse_2.png"
+        )
         _make_config(True, tmp / "on.yaml")
-        omni_image, omni_cot, time_reuse = _run_offline(str(tmp / "on.yaml"), output_dir)
+        image_reuse, cot_reuse, time_reuse = _run(str(tmp / "on.yaml"), output_dir / "with_kv_reuse.png")
+
+    print("\n--- End-to-end time ---")
+    print(f"  WITHOUT KV reuse 1: {time_no_reuse_1:.2f}s")
+    print(f"  WITHOUT KV reuse 2: {time_no_reuse_2:.2f}s")
+    print(f"  WITH KV reuse:    {time_reuse:.2f}s")
+
+    (output_dir / "cot_with_kv_reuse.txt").write_text(cot_reuse, encoding="utf-8")
+    (output_dir / "cot_without_kv_reuse_1.txt").write_text(cot_no_reuse_1, encoding="utf-8")
+    (output_dir / "cot_without_kv_reuse_2.txt").write_text(cot_no_reuse_2, encoding="utf-8")
+
+    print(f"\n--- CoT WITH KV reuse (len={len(cot_reuse)}) ---\n{cot_reuse}")
+    print(f"\n--- CoT WITHOUT KV reuse 1 (len={len(cot_no_reuse_1)}) ---\n{cot_no_reuse_1}")
+    print(f"\n--- CoT WITHOUT KV reuse 2 (len={len(cot_no_reuse_2)}) ---\n{cot_no_reuse_2}")
 
-    scorer = SemanticSimilarityScorer()
     clip_scorer = CLIPScorer()
-    cot_results = scorer.text_similarity(omni_cot, COT_REF)
-    image_ref = Image.open(str(accuracy_assets_root / "hunyuan_image_ref.png")).convert("RGB")
-    image_clip_score = clip_scorer.image_image_score(omni_image, image_ref)
-    ssim_value, psnr_value = compute_image_ssim_psnr(prediction=omni_image, reference=image_ref, compare_mode="RGB")
-
-    table = [
-        ["COT similarity to reference", f"{cot_results['cot_semantic_sim']:.4f}", 0.9644],
-        ["COT prefix match", f"{cot_results['text_prefix_match_count']:.4f}", 29],
-        ["Image-Image similarity", f"{image_clip_score:.4f}", 94.5538],
-        ["SSIM", f"{ssim_value:.4f}", 0.242],
-        ["PSNR (dB)", f"{psnr_value:.2f}", 14.1],
-    ]
-
-    print(tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
-
-    assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
-        f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {THRESHOLDS['cot_semantic_sim']}"
+    clip_scorer.assert_score(
+        model_name=f"{MODEL_NAME} with-reuse",
+        image=image_reuse,
+        text=PROMPT,
+        threshold=CLIP_SCORE_THRESHOLD,
     )
-    assert cot_results["text_prefix_match_count"] >= THRESHOLDS["text_prefix_match"], (
-        f"COT prefix match count {cot_results['text_prefix_match_count']} is below threshold {THRESHOLDS['text_prefix_match']}"
+    clip_scorer.assert_score(
+        model_name=f"{MODEL_NAME} without-reuse 1",
+        image=image_no_reuse_1,
+        text=PROMPT,
+        threshold=CLIP_SCORE_THRESHOLD,
     )
-    assert image_clip_score >= THRESHOLDS["clip_score"], (
-        f"Image-Image similarity{image_clip_score:.4f} is below threshold {THRESHOLDS['clip_score']}"
+    clip_scorer.assert_score(
+        model_name=f"{MODEL_NAME} without-reuse 2",
+        image=image_no_reuse_2,
+        text=PROMPT,
+        threshold=CLIP_SCORE_THRESHOLD,
+    )
+
+    assert_similarity(
+        model_name=f"{MODEL_NAME} non-reuse run1 vs run2",
+        vllm_image=image_no_reuse_1,
+        diffusers_image=image_no_reuse_2,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+        width=WIDTH,
+        height=HEIGHT,
+    )
+
+    assert_similarity(
+        model_name=f"{MODEL_NAME} KV-reuse vs no-reuse",
+        vllm_image=image_reuse,
+        diffusers_image=image_no_reuse_1,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+        width=WIDTH,
+        height=HEIGHT,
     )
-    assert ssim_value >= THRESHOLDS["ssim"], f"SSIM {ssim_value:.4f} is below threshold {THRESHOLDS['ssim']}"
-    assert psnr_value >= THRESHOLDS["psnr"], f"PSNR {psnr_value:.2f} dB is below threshold {THRESHOLDS['psnr']} dB"

tests/e2e/accuracy/test_qwen_image.py

@@ -12,7 +12,7 @@ import torch
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from PIL import Image
 
-from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
+from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer
@@ -32,38 +32,11 @@ SEED = 42
 SSIM_THRESHOLD = 0.97
 PSNR_THRESHOLD = 30.0
 
-MODEL_2512_ID = "Qwen/Qwen-Image-2512"
-MODEL_2512_ENV_VAR = "QWEN_IMAGE_2512_MODEL"
-PROMPT_2512 = (
-    "A 20-year-old East Asian girl with delicate, charming features and large, bright brown eyes—expressive and "
-    "lively, with a cheerful or subtly smiling expression. Her naturally wavy long hair is either loose or tied in "
-    "twin ponytails. She has fair skin and light makeup accentuating her youthful freshness. She wears a modern, "
-    "cute dress or relaxed outfit in bright, soft colors—lightweight fabric, minimalist cut. She stands indoors at "
-    "an anime convention, surrounded by banners, posters, or stalls. Lighting is typical indoor illumination—no "
-    "staged lighting—and the image resembles a casual iPhone snapshot: unpretentious composition, yet brimming "
-    "with vivid, fresh, youthful charm."
-)
-NEGATIVE_PROMPT_2512 = (
-    "低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感。"
-    "构图混乱。文字模糊，扭曲。"
-)
-WIDTH_2512 = 1664
-HEIGHT_2512 = 928
-NUM_INFERENCE_STEPS_2512 = 50
-TRUE_CFG_SCALE_2512 = 4.0
-SEED_2512 = 42
-MEAN_ABS_DIFF_THRESHOLD_2512 = 3e-2
-P99_ABS_DIFF_THRESHOLD_2512 = 4e-1
-
 
 def _model_name() -> str:
     return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
 
 
-def _model_2512_name() -> str:
-    return os.environ.get(MODEL_2512_ENV_VAR, MODEL_2512_ID)
-
-
 def _local_files_only(model: str) -> bool:
     return Path(model).exists()
 
@@ -129,67 +102,6 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
         run_post_test_cleanup()
 
 
-def _run_vllm_omni_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image:
-    server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
-    with OmniServer(model, server_args, use_omni=True) as omni_server:
-        response = requests.post(
-            f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
-            json={
-                "model": omni_server.model,
-                "prompt": PROMPT_2512,
-                "size": f"{WIDTH_2512}x{HEIGHT_2512}",
-                "n": 1,
-                "response_format": "b64_json",
-                "negative_prompt": NEGATIVE_PROMPT_2512,
-                "num_inference_steps": NUM_INFERENCE_STEPS_2512,
-                "true_cfg_scale": TRUE_CFG_SCALE_2512,
-                "seed": SEED_2512,
-            },
-            timeout=1200,
-        )
-        response.raise_for_status()
-        payload = response.json()
-        assert len(payload["data"]) == 1
-        image_bytes = base64.b64decode(payload["data"][0]["b64_json"])
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        image.load()
-        image.save(output_path)
-        return image
-
-
-def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Image:
-    run_pre_test_cleanup()
-    pipe: DiffusionPipeline | None = None
-    try:
-        pipe = DiffusionPipeline.from_pretrained(
-            model,
-            torch_dtype=torch.bfloat16,
-            trust_remote_code=True,
-            local_files_only=_local_files_only(model),
-        ).to("cuda")
-        generator = torch.Generator(device="cuda").manual_seed(SEED_2512)
-        result = pipe(  # pyright: ignore[reportCallIssue]
-            prompt=PROMPT_2512,
-            negative_prompt=NEGATIVE_PROMPT_2512,
-            width=WIDTH_2512,
-            height=HEIGHT_2512,
-            num_inference_steps=NUM_INFERENCE_STEPS_2512,
-            true_cfg_scale=TRUE_CFG_SCALE_2512,
-            generator=generator,
-        )
-        output_image = result.images[0].convert("RGB")
-        output_image.save(output_path)
-        return output_image
-    finally:
-        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
-            pipe.maybe_free_model_hooks()
-        del pipe
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.accelerator.empty_cache()
-        run_post_test_cleanup()
-
-
 @pytest.mark.benchmark
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None:
@@ -208,27 +120,3 @@ def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None:
         ssim_threshold=SSIM_THRESHOLD,
         psnr_threshold=PSNR_THRESHOLD,
     )
-
-
-@pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_qwen_image_2512_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
-    model = _model_2512_name()
-    output_dir = model_output_dir(accuracy_artifact_root, MODEL_2512_ID)
-    vllm_output_path = output_dir / "vllm_omni_2512.png"
-    diffusers_output_path = output_dir / "diffusers_2512.png"
-
-    vllm_output = _run_vllm_omni_qwen_image_2512(model=model, output_path=vllm_output_path)
-    diffusers_output = _run_diffusers_qwen_image_2512(model=model, output_path=diffusers_output_path)
-
-    print(f"{MODEL_2512_ID} generated images:")
-    print(f"  vllm_omni: {vllm_output_path}")
-    print(f"  diffusers: {diffusers_output_path}")
-
-    assert_images_pixel_close(
-        model_name=MODEL_2512_ID,
-        vllm_image=vllm_output,
-        diffusers_image=diffusers_output,
-        mean_threshold=MEAN_ABS_DIFF_THRESHOLD_2512,
-        p99_threshold=P99_ABS_DIFF_THRESHOLD_2512,
-    )

tests/e2e/accuracy/test_qwen_image_edit.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import gc
-import os
 from pathlib import Path
 
 import pytest
@@ -11,7 +10,7 @@ from diffusers import QwenImageEditPipeline, QwenImageEditPlusPipeline
 from PIL import Image
 
 from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes
-from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
+from tests.e2e.accuracy.helpers import assert_similarity, model_output_dir
 from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServer
@@ -29,43 +28,12 @@ SEED = 42
 SSIM_THRESHOLD = 0.94
 PSNR_THRESHOLD = 28.0
 
-EDIT_2511_MODEL = "Qwen/Qwen-Image-Edit-2511"
-EDIT_2511_MODEL_ENV_VAR = "QWEN_IMAGE_EDIT_2511_MODEL"
-EDIT_2511_WIDTH = 1190
-EDIT_2511_HEIGHT = 2032
-EDIT_2511_NUM_INFERENCE_STEPS = 40
-EDIT_2511_TRUE_CFG_SCALE = 4.0
-EDIT_2511_GUIDANCE_SCALE = 1.0
-EDIT_2511_SEED = 42
-EDIT_2511_MEAN_ABS_DIFF_THRESHOLD = 3e-2
-EDIT_2511_P99_ABS_DIFF_THRESHOLD = 4e-1
-EDIT_2511_NEGATIVE_PROMPT = " "
-EDIT_2511_PROMPT = (
-    "将第二张图中人脸/猫狗脸转换为3D卡通形象，质量要求：毛发纹理自然电影级色彩校准，注意保留原图的外貌、"
-    "肤色、发型特征，仅替换第一张财神爷/猫狗财神的用绿框框住的头部区域，耳朵样式和原图的3D卡通形象一致， "
-    "（帽子紧贴头顶，未贴合部分可调整帽子摆放角度达到完全贴合）， 头部与身体过渡必须自然，无拼接痕迹，"
-    "无额外元素出现在头部四周，纯白背景，不要出现绿色圆圈"
-)
-EDIT_2511_ASSET_DIR = Path(__file__).resolve().parent / "assets"
-EDIT_2511_IMAGE_PATHS = [
-    EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test1.png",
-    EDIT_2511_ASSET_DIR / "qwen_image_edit_2511_test2.png",
-]
-
 PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure."
 PROMPT_MULTIPLE_IMAGE = "Put the cartoon bear mascot and the furry rabbit into one coherent scene with a painterly oil artwork style and consistent lighting."
 NEGATIVE_PROMPT = "low quality, blurry, artifacts, distortion"
 SERVER_ARGS = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"]
 
 
-def _edit_2511_model_name() -> str:
-    return os.environ.get(EDIT_2511_MODEL_ENV_VAR, EDIT_2511_MODEL)
-
-
-def _local_files_only(model: str) -> bool:
-    return Path(model).exists()
-
-
 def _run_vllm_omni_image_edit(
     *,
     omni_server: OmniServer,
@@ -111,15 +79,17 @@ def _run_diffusers_image_edit(
 ) -> Image.Image:
     run_pre_test_cleanup()
     pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None
+    device = torch.device("cuda:0")
+    torch.cuda.set_device(device)
     try:
         images = input_images[0] if len(input_images) == 1 else input_images
         pipe = pipeline_class.from_pretrained(
             model,
             torch_dtype=torch.bfloat16,
             trust_remote_code=True,
-        ).to("cuda")
+        ).to(device)
         pipe.set_progress_bar_config(disable=False)
-        generator = torch.Generator(device="cuda").manual_seed(SEED)
+        generator = torch.Generator(device=device).manual_seed(SEED)
         result = pipe(  # pyright: ignore[reportCallIssue]
             prompt=prompt,
             image=images,
@@ -143,88 +113,6 @@ def _run_diffusers_image_edit(
         run_post_test_cleanup()
 
 
-def _load_edit_2511_input_images() -> list[Image.Image]:
-    return [Image.open(path).convert("RGB") for path in EDIT_2511_IMAGE_PATHS]
-
-
-def _run_vllm_omni_image_edit_2511(
-    *,
-    model: str,
-    input_images: list[Image.Image],
-    output_path: Path,
-) -> Image.Image:
-    with OmniServer(model=model, serve_args=SERVER_ARGS) as server:
-        response = requests.post(
-            f"http://{server.host}:{server.port}/v1/images/edits",
-            data={
-                "model": server.model,
-                "prompt": EDIT_2511_PROMPT,
-                "size": f"{EDIT_2511_WIDTH}x{EDIT_2511_HEIGHT}",
-                "n": 1,
-                "response_format": "b64_json",
-                "negative_prompt": EDIT_2511_NEGATIVE_PROMPT,
-                "num_inference_steps": EDIT_2511_NUM_INFERENCE_STEPS,
-                "true_cfg_scale": EDIT_2511_TRUE_CFG_SCALE,
-                "guidance_scale": EDIT_2511_GUIDANCE_SCALE,
-                "seed": EDIT_2511_SEED,
-            },
-            files=[
-                ("image", (f"image_{index}.png", pil_to_png_bytes(image), "image/png"))
-                for index, image in enumerate(input_images)
-            ],
-            timeout=1200,
-        )
-        response.raise_for_status()
-        payload = response.json()
-        assert len(payload["data"]) == 1
-        image = decode_base64_image(payload["data"][0]["b64_json"])
-        image.load()
-        image.save(output_path)
-        return image
-
-
-def _run_diffusers_image_edit_2511(
-    *,
-    model: str,
-    input_images: list[Image.Image],
-    output_path: Path,
-) -> Image.Image:
-    run_pre_test_cleanup()
-    pipe: QwenImageEditPlusPipeline | None = None
-    try:
-        pipe = QwenImageEditPlusPipeline.from_pretrained(
-            model,
-            torch_dtype=torch.bfloat16,
-            trust_remote_code=True,
-            local_files_only=_local_files_only(model),
-        ).to("cuda")
-        pipe.set_progress_bar_config(disable=False)
-        generator = torch.Generator(device="cuda").manual_seed(EDIT_2511_SEED)
-        result = pipe(  # pyright: ignore[reportCallIssue]
-            image=input_images,
-            prompt=EDIT_2511_PROMPT,
-            generator=generator,
-            true_cfg_scale=EDIT_2511_TRUE_CFG_SCALE,
-            negative_prompt=EDIT_2511_NEGATIVE_PROMPT,
-            num_inference_steps=EDIT_2511_NUM_INFERENCE_STEPS,
-            guidance_scale=EDIT_2511_GUIDANCE_SCALE,
-            num_images_per_prompt=1,
-            width=EDIT_2511_WIDTH,
-            height=EDIT_2511_HEIGHT,
-        )
-        output_image = result.images[0].convert("RGB")  # pyright: ignore[reportAttributeAccessIssue]
-        output_image.save(output_path)
-        return output_image
-    finally:
-        if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"):
-            pipe.maybe_free_model_hooks()
-        del pipe
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.accelerator.empty_cache()
-        run_post_test_cleanup()
-
-
 def _vllm_omni_output_single_image(
     accuracy_artifact_root: Path,
     qwen_bear_image: Image.Image,
@@ -338,43 +226,3 @@ def test_qwen_image_edit_multiple_matches_diffusers(
         ssim_threshold=SSIM_THRESHOLD,
         psnr_threshold=PSNR_THRESHOLD,
     )
-
-
-@pytest.mark.benchmark
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
-    model = _edit_2511_model_name()
-    output_dir = model_output_dir(accuracy_artifact_root, EDIT_2511_MODEL)
-    input_images = _load_edit_2511_input_images()
-    input_paths: list[Path] = []
-    for index, image in enumerate(input_images, start=1):
-        input_path = output_dir / f"edit_2511_input_{index}.png"
-        image.save(input_path)
-        input_paths.append(input_path)
-    vllm_output_path = output_dir / "vllm_omni_edit_2511.png"
-    diffusers_output_path = output_dir / "diffusers_edit_2511.png"
-
-    vllm_image = _run_vllm_omni_image_edit_2511(
-        model=model,
-        input_images=input_images,
-        output_path=vllm_output_path,
-    )
-    diffusers_image = _run_diffusers_image_edit_2511(
-        model=model,
-        input_images=input_images,
-        output_path=diffusers_output_path,
-    )
-
-    print(f"{EDIT_2511_MODEL} generated images:")
-    for index, input_path in enumerate(input_paths, start=1):
-        print(f"  input_{index}: {input_path}")
-    print(f"  vllm_omni: {vllm_output_path}")
-    print(f"  diffusers: {diffusers_output_path}")
-
-    assert_images_pixel_close(
-        model_name=EDIT_2511_MODEL,
-        vllm_image=vllm_image,
-        diffusers_image=diffusers_image,
-        mean_threshold=EDIT_2511_MEAN_ABS_DIFF_THRESHOLD,
-        p99_threshold=EDIT_2511_P99_ABS_DIFF_THRESHOLD,
-    )

tests/e2e/offline_inference/test_cosyvoice3_expansion.py

@@ -27,9 +27,9 @@ from tests.helpers.mark import hardware_test
 from tests.helpers.media import get_asset_path
 from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 
 MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
 MODEL_DIR_ENV = "VLLM_OMNI_COSYVOICE3_MODEL_DIR"

tests/entrypoints/test_omni_entrypoints.py

@@ -943,20 +943,6 @@ def test_omni_base_errored_false_when_alive():
     assert base.errored is False
 
 
-def test_omni_base_is_running_false_when_stage_engine_dead():
-    base = _make_base()
-    base.engine.is_alive.return_value = True
-    base.engine.stage_clients = [SimpleNamespace(_engine_dead=True)]
-    assert base.is_running is False
-
-
-def test_omni_base_is_running_false_when_stage_resources_engine_dead():
-    base = _make_base()
-    base.engine.is_alive.return_value = True
-    base.engine.stage_clients = [SimpleNamespace(resources=SimpleNamespace(engine_dead=True))]
-    assert base.is_running is False
-
-
 def test_omni_base_errored_true_when_orchestrator_dead():
     base = _make_base()
     base.engine.is_alive.return_value = False

tests/test_config_factory.py

@@ -4,40 +4,27 @@
 Unit tests for StageConfigFactory and related classes.
 """
 
-import importlib
-import inspect
 import warnings
-from dataclasses import dataclass, fields
+from dataclasses import dataclass
 from pathlib import Path
-from unittest.mock import patch
 
 import pytest
 
-from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.config.stage_config import (
     _PIPELINE_REGISTRY,
-    DeployConfig,
     ModelPipeline,
     PipelineConfig,
     StageConfig,
     StageConfigFactory,
-    StageDeployConfig,
     StageExecutionType,
     StagePipelineConfig,
     StageType,
-    _apply_platform_overrides,
-    _deep_merge_stage,
     _resolve_scheduler,
     build_stage_runtime_overrides,
-    deploy_override_field_names,
-    load_deploy_config,
-    merge_pipeline_deploy,
     register_pipeline,
     strip_parent_engine_args,
 )
-from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
 
 
 class TestStageType:
@@ -112,6 +99,8 @@ class TestStageConfig:
 
     def test_to_omegaconf_with_runtime_overrides(self):
         """Test that runtime overrides are applied to OmegaConf output."""
+        import warnings
+
         config = StageConfig(
             stage_id=0,
             model_stage="thinker",
@@ -134,6 +123,8 @@ class TestStageConfig:
 
     def test_to_omegaconf_max_batch_size_deprecation(self):
         """Test that runtime.max_batch_size emits a FutureWarning."""
+        import warnings
+
         config = StageConfig(
             stage_id=0,
             model_stage="thinker",
@@ -158,6 +149,7 @@ class TestStageConfig:
 
     def test_to_omegaconf_omits_none_deploy_overrides_for_engine_args(self):
         """None deploy overrides must fall through to EngineArgs defaults."""
+        from vllm_omni.config.stage_config import deploy_override_field_names
 
         config = StageConfig(
             stage_id=0,
@@ -378,7 +370,11 @@ class TestStageResolutionHelpers:
         assert "parallel_config" not in overrides
 
     def test_strip_parent_engine_args_reports_only_surprising_parent_overrides(self):
-        parent_fields = {f.name: f for f in fields(EngineArgs)}
+        from dataclasses import fields as dc_fields
+
+        from vllm.engine.arg_utils import EngineArgs
+
+        parent_fields = {f.name: f for f in dc_fields(EngineArgs)}
         filtered, overridden = strip_parent_engine_args(
             {
                 "model": "some/model",
@@ -400,7 +396,11 @@ class TestStageResolutionHelpers:
         assert overridden == ["tensor_parallel_size"]
 
     def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self):
-        parent_fields = {f.name: f for f in fields(EngineArgs)}
+        from dataclasses import fields as dc_fields
+
+        from vllm.engine.arg_utils import EngineArgs
+
+        parent_fields = {f.name: f for f in dc_fields(EngineArgs)}
         filtered, overridden = strip_parent_engine_args(
             {
                 "model": "some/model",
@@ -826,6 +826,8 @@ class TestPipelineRegistry:
 
 class TestDeployConfigLoading:
     def test_deploy_override_fields_include_deploy_schema_fields(self):
+        from vllm_omni.config.stage_config import deploy_override_field_names
+
         expected_fields = {
             "async_chunk",
             "async_scheduling",
@@ -862,6 +864,10 @@ class TestDeployConfigLoading:
         )
 
     def test_load_qwen3_omni_moe_deploy_config(self):
+        from pathlib import Path
+
+        from vllm_omni.config.stage_config import load_deploy_config
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         deploy = load_deploy_config(deploy_path)
         assert len(deploy.stages) == 3
@@ -870,6 +876,10 @@ class TestDeployConfigLoading:
         assert deploy.platforms is not None
 
     def test_load_voxtral_tts_deploy_config_schema_fields(self):
+        from pathlib import Path
+
+        from vllm_omni.config.stage_config import load_deploy_config
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "voxtral_tts.yaml"
         deploy = load_deploy_config(deploy_path)
         assert deploy.stages[0].config_format == "mistral"
@@ -880,12 +890,21 @@ class TestDeployConfigLoading:
         )
 
     def test_load_ming_flash_omni_deploy_config_schema_fields(self):
+        from pathlib import Path
+
+        from vllm_omni.config.stage_config import load_deploy_config
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
         deploy = load_deploy_config(deploy_path)
         assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}}
         assert "compilation_config" not in deploy.stages[0].engine_extras
 
     def test_merge_pipeline_deploy(self):
+        from pathlib import Path
+
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
@@ -902,6 +921,11 @@ class TestDeployConfigLoading:
         assert s0.yaml_extras["default_sampling_params"]["detokenize"] is True
 
     def test_merge_pipeline_deploy_preserves_num_replicas(self, tmp_path):
+        from pathlib import Path
+
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not base.exists():
@@ -918,6 +942,15 @@ class TestDeployConfigLoading:
         assert stages[1].yaml_runtime["num_replicas"] == 2
 
     def test_merge_pipeline_deploy_preserves_requires_multimodal_data(self):
+        from vllm_omni.config.stage_config import (
+            DeployConfig,
+            PipelineConfig,
+            StageDeployConfig,
+            StageExecutionType,
+            StagePipelineConfig,
+            merge_pipeline_deploy,
+        )
+
         pipeline = PipelineConfig(
             model_type="test_mm",
             model_arch="TestModel",
@@ -936,186 +969,11 @@ class TestDeployConfigLoading:
 
         assert stages[0].yaml_runtime["requires_multimodal_data"] is True
 
-    def test_mixed_schema_preserves_flat_fields(self):
-        """Ensure flat fields are not dropped when engine_args are present."""
-        fake_config = {
-            "stages": [
-                {
-                    "stage_id": 0,
-                    "gpu_memory_utilization": 0.7,
-                    "max_num_seqs": 8,
-                    "engine_args": {
-                        "tensor_parallel_size": 4,
-                        "enforce_eager": True,
-                    },
-                },
-            ]
-        }
-
-        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
-            deploy = load_deploy_config("dummy.yaml")
-
-        assert len(deploy.stages) == 1
-        stage = deploy.stages[0]
-        # Check that the engine args are set
-        assert stage.tensor_parallel_size == 4
-        assert stage.enforce_eager is True
-        # Check that the other settings are also preserved
-        assert stage.gpu_memory_utilization == 0.7
-        assert stage.max_num_seqs == 8
-
-    def test_engine_parse_engine_fields(self):
-        """Test that we correctly parse & recursively merge stage deploy fields."""
-        fake_config = {
-            "stages": [
-                {
-                    "stage_id": 0,
-                    "compilation_config": {
-                        "encoder_cudagraph_token_budgets": [1024, 2048],
-                        "pass_config": {"fuse_norm_quant": True},
-                    },
-                    "engine_args": {
-                        "compilation_config": {
-                            "cudagraph_mm_encoder": True,
-                            "pass_config": {"fuse_allreduce_rms": False},
-                        },
-                    },
-                },
-            ]
-        }
-
-        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
-            deploy = load_deploy_config("dummy.yaml")
-
-        assert len(deploy.stages) == 1
-        stage = deploy.stages[0]
-        assert stage.compilation_config == {
-            "encoder_cudagraph_token_budgets": [1024, 2048],
-            "pass_config": {
-                "fuse_norm_quant": True,
-                "fuse_allreduce_rms": False,
-            },
-            "cudagraph_mm_encoder": True,
-        }
-
-    def test_engine_extras_deep_merges_dicts_simple(self):
-        """Ensure dictionary valued keys merge properly for top level dicts."""
-        fake_config = {
-            "stages": [
-                {
-                    "stage_id": 0,
-                    "foo": {"a": 111, "b": 1},
-                    "engine_args": {
-                        "foo": {"b": 2, "c": 3},
-                    },
-                },
-            ]
-        }
-
-        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
-            deploy = load_deploy_config("dummy.yaml")
-
-        assert len(deploy.stages) == 1
-        stage = deploy.stages[0]
-        assert "foo" in stage.engine_extras
-        assert stage.engine_extras["foo"] == {"a": 111, "b": 2, "c": 3}
-
-    def test_engine_extras_dict_type_mismatch(self):
-        """Ensure that we handle type mismatches with nested dicts correctly."""
-        fake_config = {
-            "stages": [
-                {
-                    "stage_id": 0,
-                    "foo": {"b": {1: 1}},
-                    "engine_args": {
-                        "foo": {"b": 2},
-                    },
-                },
-            ]
-        }
-
-        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
-            deploy = load_deploy_config("dummy.yaml")
-
-        assert len(deploy.stages) == 1
-        stage = deploy.stages[0]
-        assert "foo" in stage.engine_extras
-        assert stage.engine_extras["foo"] == {"b": 2}
-
-    def test_mixed_engine_extras_deep_merges_dicts(self):
-        """Ensure dictionary valued keys merge properly for nested dicts."""
-        fake_config = {
-            "stages": [
-                {
-                    "stage_id": 0,
-                    "foo": {"a": 111, "b": {"e": 199}},
-                    "engine_args": {
-                        "foo": {"b": {"d": 9}, "c": 3},
-                    },
-                },
-            ]
-        }
-
-        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
-            deploy = load_deploy_config("dummy.yaml")
-
-        assert len(deploy.stages) == 1
-        stage = deploy.stages[0]
-        assert "foo" in stage.engine_extras
-        assert stage.engine_extras["foo"] == {"a": 111, "b": {"d": 9, "e": 199}, "c": 3}
-
-    def test_deep_merge_does_not_mutate_inputs(self):
-        """Merging engine_args must not mutate the base stage dict."""
-        fake_config = {
-            "stages": [
-                {
-                    "stage_id": 0,
-                    "foo": {"a": 1, "b": {"x": 10, "y": 20}},
-                    "engine_args": {
-                        "foo": {"b": {"y": 99, "z": 30}, "c": 3},
-                    },
-                },
-            ]
-        }
-        original_foo = fake_config["stages"][0]["foo"]
-        original_b = original_foo["b"]
-
-        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
-            deploy = load_deploy_config("dummy.yaml")
-
-        # Values should merge into the new dict correctly
-        assert deploy.stages[0].engine_extras["foo"] == {
-            "a": 1,
-            "b": {"x": 10, "y": 99, "z": 30},
-            "c": 3,
-        }
-        # But the original nested dicts foo / b should be unchanged, since
-        # we recursively shallow copy to avoid mutating in place.
-        assert original_foo == {"a": 1, "b": {"x": 10, "y": 20}}
-        assert original_b == {"x": 10, "y": 20}
-
-    def test_mixed_schema_engine_args_wins_scalars(self):
-        """engine_args takes precedence over flat fields for scalar conflicts."""
-        fake_config = {
-            "stages": [
-                {
-                    "stage_id": 0,
-                    "gpu_memory_utilization": 0.7,
-                    "engine_args": {
-                        "gpu_memory_utilization": 0.5,
-                    },
-                },
-            ]
-        }
-
-        with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
-            deploy = load_deploy_config("dummy.yaml")
-
-        assert deploy.stages[0].gpu_memory_utilization == 0.5
-
 
 class TestQwen3OmniPipeline:
     def test_registered(self):
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+
         p = _PIPELINE_REGISTRY.get("qwen3_omni_moe")
         assert p is not None
         assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration"
@@ -1123,6 +981,8 @@ class TestQwen3OmniPipeline:
         assert p.validate() == []
 
     def test_thinker(self):
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(0)
         assert s.model_stage == "thinker"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -1131,6 +991,8 @@ class TestQwen3OmniPipeline:
         assert s.sampling_constraints["detokenize"] is True
 
     def test_talker(self):
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(1)
         assert s.input_sources == (0,)
         assert s.sampling_constraints["stop_token_ids"] == [2150]
@@ -1138,6 +1000,8 @@ class TestQwen3OmniPipeline:
         assert s.custom_process_next_stage_input_func is not None
 
     def test_code2wav(self):
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(2)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
@@ -1146,6 +1010,8 @@ class TestQwen3OmniPipeline:
 
 class TestQwen2_5OmniPipeline:
     def test_registered(self):
+        import vllm_omni.model_executor.models.qwen2_5_omni.pipeline  # noqa: F401
+
         p = _PIPELINE_REGISTRY.get("qwen2_5_omni")
         assert p is not None
         assert p.model_arch == "Qwen2_5OmniForConditionalGeneration"
@@ -1153,6 +1019,8 @@ class TestQwen2_5OmniPipeline:
         assert p.validate() == []
 
     def test_thinker(self):
+        import vllm_omni.model_executor.models.qwen2_5_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(0)
         assert s.model_stage == "thinker"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -1161,12 +1029,16 @@ class TestQwen2_5OmniPipeline:
         assert s.requires_multimodal_data is True
 
     def test_talker(self):
+        import vllm_omni.model_executor.models.qwen2_5_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(1)
         assert s.input_sources == (0,)
         assert s.sampling_constraints["stop_token_ids"] == [8294]
         assert s.custom_process_input_func is not None
 
     def test_code2wav(self):
+        import vllm_omni.model_executor.models.qwen2_5_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(2)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
@@ -1175,6 +1047,8 @@ class TestQwen2_5OmniPipeline:
 
 class TestQwen3TTSPipeline:
     def test_registered(self):
+        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
+
         p = _PIPELINE_REGISTRY.get("qwen3_tts")
         assert p is not None
         assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration"
@@ -1182,6 +1056,8 @@ class TestQwen3TTSPipeline:
         assert p.validate() == []
 
     def test_talker_stage(self):
+        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(0)
         assert s.model_stage == "qwen3_tts"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -1192,6 +1068,8 @@ class TestQwen3TTSPipeline:
         assert s.model_arch is None
 
     def test_code2wav_stage_has_per_stage_model_arch(self):
+        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(1)
         assert s.execution_type == StageExecutionType.LLM_GENERATION
         assert s.final_output_type == "audio"
@@ -1203,6 +1081,9 @@ class TestQwen3TTSPipeline:
 
     def test_per_stage_model_arch_flows_through_merge(self, tmp_path):
         """Verify the new ps.model_arch override survives merge_pipeline_deploy."""
+        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_tts.yaml"
         if not deploy_path.exists():
             pytest.skip("qwen3_tts deploy yaml not found")
@@ -1218,6 +1099,8 @@ class TestQwen3TTSPipeline:
 
     def test_subtalker_sampling_params_deep_merge_preserves_base_keys(self):
         """Verify subtalker sampling params participate in stage deep-merge."""
+        from vllm_omni.config.stage_config import _deep_merge_stage
+
         base = {
             "stage_id": 0,
             "subtalker_sampling_params": {
@@ -1247,6 +1130,8 @@ class TestQwen3TTSPipeline:
 
 class TestMingFlashOmniPipeline:
     def test_registered(self):
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+
         p = _PIPELINE_REGISTRY.get("ming_flash_omni")
         assert p is not None
         assert p.model_arch == "MingFlashOmniForConditionalGeneration"
@@ -1254,6 +1139,8 @@ class TestMingFlashOmniPipeline:
         assert p.validate() == []
 
     def test_thinker_stage(self):
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(0)
         assert s.model_stage == "thinker"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -1264,6 +1151,8 @@ class TestMingFlashOmniPipeline:
         assert s.sampling_constraints["detokenize"] is True
 
     def test_talker_stage(self):
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1)
         assert s.model_stage == "ming_tts"
         assert s.execution_type == StageExecutionType.LLM_GENERATION
@@ -1282,12 +1171,18 @@ class TestMingFlashOmniPipeline:
         Lazy string references only fail at first inference otherwise — this
         catches typos in the pipeline declaration at import / registration time.
         """
+        import importlib
+
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1)
         module_path, _, attr = s.custom_process_input_func.rpartition(".")
         module = importlib.import_module(module_path)
         assert callable(getattr(module, attr))
 
     def test_tts_pipeline_registered(self):
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+
         p = _PIPELINE_REGISTRY.get("ming_flash_omni_tts")
         assert p is not None
         assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration"
@@ -1295,6 +1190,8 @@ class TestMingFlashOmniPipeline:
         assert p.validate() == []
 
     def test_tts_stage(self):
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["ming_flash_omni_tts"].get_stage(0)
         assert s.model_stage == "ming_tts"
         assert s.execution_type == StageExecutionType.LLM_GENERATION
@@ -1307,6 +1204,9 @@ class TestMingFlashOmniPipeline:
 
     def test_full_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni.yaml parses and merges with the registered pipeline."""
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni.yaml"
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni deploy yaml not found")
@@ -1326,6 +1226,9 @@ class TestMingFlashOmniPipeline:
 
     def test_tts_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_tts.yaml parses and routes to the TTS-only pipeline."""
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_tts.yaml"
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_tts deploy yaml not found")
@@ -1340,6 +1243,8 @@ class TestMingFlashOmniPipeline:
         assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniTalkerForConditionalGeneration"
 
     def test_thinker_only_pipeline_registered(self):
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+
         p = _PIPELINE_REGISTRY.get("ming_flash_omni_thinker_only")
         assert p is not None
         assert p.model_arch == "MingFlashOmniForConditionalGeneration"
@@ -1347,6 +1252,8 @@ class TestMingFlashOmniPipeline:
         assert p.validate() == []
 
     def test_thinker_only_stage(self):
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+
         s = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"].get_stage(0)
         assert s.model_stage == "thinker"
         assert s.execution_type == StageExecutionType.LLM_AR
@@ -1360,6 +1267,9 @@ class TestMingFlashOmniPipeline:
 
     def test_thinker_only_yaml_loads_and_merges(self):
         """deploy/ming_flash_omni_thinker_only.yaml parses and routes to the thinker-only pipeline."""
+        import vllm_omni.model_executor.models.ming_flash_omni.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "ming_flash_omni_thinker_only.yaml"
         if not deploy_path.exists():
             pytest.skip("ming_flash_omni_thinker_only deploy yaml not found")
@@ -1378,6 +1288,9 @@ class TestBaseConfigInheritance:
     """Test deploy YAML base_config inheritance."""
 
     def test_ci_inherits_from_main(self):
+        from tests.helpers.stage_config import get_deploy_config_path
+        from vllm_omni.config.stage_config import load_deploy_config
+
         ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml"))
         if not ci_path.exists():
             pytest.skip("CI deploy config not found")
@@ -1385,7 +1298,7 @@ class TestBaseConfigInheritance:
         deploy = load_deploy_config(ci_path)
         assert len(deploy.stages) == 3
         # CI overrides
-        assert deploy.stages[0].load_format is None
+        assert deploy.stages[0].load_format == "dummy"
         assert "load_format" not in deploy.stages[0].engine_extras
         assert deploy.stages[0].max_num_seqs == 5
         # Inherited from base
@@ -1398,6 +1311,9 @@ class TestBaseConfigInheritance:
         assert deploy.async_chunk is False
 
     def test_ci_sampling_merge(self):
+        from tests.helpers.stage_config import get_deploy_config_path
+        from vllm_omni.config.stage_config import load_deploy_config
+
         ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml"))
         if not ci_path.exists():
             pytest.skip("CI deploy config not found")
@@ -1412,6 +1328,8 @@ class TestBaseConfigInheritance:
 
     def test_pure_inheritance_overlay(self, tmp_path):
         """An overlay with only ``base_config`` inherits everything."""
+        from vllm_omni.config.stage_config import load_deploy_config
+
         base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not base.exists():
             pytest.skip("Base deploy config not found")
@@ -1425,6 +1343,8 @@ class TestBaseConfigInheritance:
 
     def test_single_field_overlay(self, tmp_path):
         """An overlay overriding one stage field merges with the base."""
+        from vllm_omni.config.stage_config import load_deploy_config
+
         base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not base.exists():
             pytest.skip("Base deploy config not found")
@@ -1442,6 +1362,10 @@ class TestPlatformOverrides:
     """Test platform-specific deploy config overrides."""
 
     def test_npu_overrides(self):
+        from pathlib import Path
+
+        from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
@@ -1456,6 +1380,10 @@ class TestPlatformOverrides:
         assert deploy.stages[2].enforce_eager is True
 
     def test_xpu_overrides(self):
+        from pathlib import Path
+
+        from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
@@ -1468,6 +1396,10 @@ class TestPlatformOverrides:
         assert deploy.stages[0].engine_extras.get("max_cudagraph_capture_size") == 0
 
     def test_unknown_platform_noop(self):
+        from pathlib import Path
+
+        from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
+
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
             pytest.skip("Deploy config not found")
@@ -1479,6 +1411,8 @@ class TestPlatformOverrides:
 
     def test_platforms_deep_merge_inheritance(self, tmp_path):
         """Overlay's platforms: block layers onto base's, per-stage."""
+        from vllm_omni.config.stage_config import _apply_platform_overrides, load_deploy_config
+
         base = tmp_path / "base.yaml"
         base.write_text(
             "stages:\n"
@@ -1513,6 +1447,11 @@ class TestCLIOverrideFlow:
     """Test --stage-overrides JSON merge into StageConfig."""
 
     def test_stage_overrides_merge(self):
+        from pathlib import Path
+
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
@@ -1527,6 +1466,11 @@ class TestCLIOverrideFlow:
         assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5
 
     def test_global_override_applies_to_all(self):
+        from pathlib import Path
+
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():
@@ -1546,6 +1490,8 @@ class TestSentinelDefaultPrecedence:
     to YAML / dataclass defaults (#3035)."""
 
     def _stages(self, cli_overrides):
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+
         return StageConfigFactory._create_from_registry(
             "qwen3_omni_moe",
             cli_overrides=cli_overrides,
@@ -1558,7 +1504,7 @@ class TestSentinelDefaultPrecedence:
     def test_none_value_skipped_yaml_wins(self):
         stages = self._stages({"max_num_seqs": None})
         assert stages[2].runtime_overrides.get("max_num_seqs") is None
-        assert stages[2].yaml_engine_args.get("max_num_seqs") == 64
+        assert stages[2].yaml_engine_args["max_num_seqs"] == 64
 
     def test_empty_kwargs_yaml_only(self):
         stages = self._stages({})
@@ -1604,11 +1550,15 @@ class TestSentinelDefaultPrecedence:
             assert stage.runtime_overrides == {}
 
     def test_create_from_registry_no_cli_explicit_keys_param(self):
+        import inspect
+
         sig = inspect.signature(StageConfigFactory._create_from_registry)
         named = [p for p in sig.parameters.values() if p.kind != p.VAR_KEYWORD]
         assert "cli_explicit_keys" not in {p.name for p in named}
 
     def test_cli_explicit_keys_kwarg_emits_deprecation(self):
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+
         with warnings.catch_warnings(record=True) as w:
             warnings.simplefilter("always")
             StageConfigFactory._create_from_registry(
@@ -1622,6 +1572,13 @@ class TestSentinelDefaultPrecedence:
         """A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end
         processors based on ``deploy.async_chunk``, without needing a
         separate variant pipeline registration."""
+        import vllm_omni.model_executor.models.qwen3_tts.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import (
+            _PIPELINE_REGISTRY,
+            DeployConfig,
+            merge_pipeline_deploy,
+        )
+
         pipeline = _PIPELINE_REGISTRY["qwen3_tts"]
 
         # async_chunk=True → stage 0's per-chunk processor wires up, stage 1
@@ -1681,6 +1638,11 @@ class TestSamplingConstraintsPrecedence:
     """Test that pipeline sampling_constraints override deploy defaults."""
 
     def test_constraints_win(self):
+        from pathlib import Path
+
+        import vllm_omni.model_executor.models.qwen3_omni.pipeline  # noqa: F401
+        from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
         pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"]
         deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml"
         if not deploy_path.exists():

vllm_omni/config/stage_config.py

@@ -473,7 +473,7 @@ class DeployConfig:
     pipeline_parallel_size: int | None = None
 
 
-_STAGE_RESERVED_KEYS = frozenset(
+_STAGE_NON_ENGINE_KEYS = frozenset(
     {
         "stage_id",
         "devices",
@@ -482,31 +482,24 @@ _STAGE_RESERVED_KEYS = frozenset(
         "input_connectors",
         "default_sampling_params",
         "engine_extras",
-        "engine_args",
-        "runtime",
     }
 )
 
 # Fields on StageDeployConfig that are populated from engine_args dict
-_STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_RESERVED_KEYS}
+_STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_NON_ENGINE_KEYS}
 
 
 def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
     """Parse a single stage entry from deploy YAML into StageDeployConfig."""
-    # Get the non-reserved keys for this stage
-    flat_args = {k: v for k, v in stage_data.items() if k not in _STAGE_RESERVED_KEYS}
-    runtime_cfg = dict(stage_data.get("runtime", {}))
-    devices = runtime_cfg.get("devices", stage_data.get("devices"))
-    num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1))
-
     if "engine_args" in stage_data:
-        for k, v in stage_data["engine_args"].items():
-            existing = flat_args.get(k)
-            # If we have multiple dictionaries, merge recursively.
-            if isinstance(v, dict) and isinstance(existing, dict):
-                flat_args[k] = _get_recursively_merged_dict(existing, v)
-            else:
-                flat_args[k] = v
+        runtime_cfg = dict(stage_data.get("runtime", {}))
+        engine_args = dict(stage_data["engine_args"])
+        devices = stage_data.get("runtime", {}).get("devices", stage_data.get("devices"))
+        num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1))
+    else:
+        engine_args = {k: v for k, v in stage_data.items() if k not in _STAGE_NON_ENGINE_KEYS and k != "stage_id"}
+        devices = stage_data.get("devices")
+        num_replicas = stage_data.get("num_replicas", stage_data.get("runtime", {}).get("num_replicas", 1))
 
     kwargs: dict[str, Any] = {
         "stage_id": stage_data["stage_id"],
@@ -514,13 +507,13 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
         "num_replicas": int(num_replicas),
     }
     for name, f in _STAGE_DEPLOY_FIELDS.items():
-        if name in flat_args:
-            kwargs[name] = flat_args.pop(name)
+        if name in engine_args:
+            kwargs[name] = engine_args.pop(name)
 
     kwargs["output_connectors"] = stage_data.get("output_connectors")
     kwargs["input_connectors"] = stage_data.get("input_connectors")
     kwargs["default_sampling_params"] = stage_data.get("default_sampling_params")
-    kwargs["engine_extras"] = flat_args
+    kwargs["engine_extras"] = engine_args
     return StageDeployConfig(**kwargs)
 
 
@@ -529,34 +522,24 @@ _DEEP_MERGE_KEYS = frozenset({"default_sampling_params", "subtalker_sampling_par
 
 def _deep_merge_stage(base: dict, overlay: dict) -> dict:
     """Deep-merge ``_DEEP_MERGE_KEYS`` so thin overlays don't drop base keys."""
-    # Deep merge _DEEP_MERGE_KEYS recursively
-    base_merge_dict = {k: v for k, v in base.items() if k in _DEEP_MERGE_KEYS}
-    overlay_merge_dict = {k: v for k, v in overlay.items() if k in _DEEP_MERGE_KEYS}
-
-    # Get the merge dict; priority is base < overlay < merged sub
-    merged_subdict = _get_recursively_merged_dict(original=base_merge_dict, update=overlay_merge_dict)
-    merged_dict = {**base, **overlay, **merged_subdict}
-    return merged_dict
-
-
-def _get_recursively_merged_dict(original: dict, update: dict) -> dict:
-    """Recursively merge two dicts, returning a new dict."""
-    merged = original.copy()
-    for k, update_v in update.items():
-        orig_v = merged.get(k)
-        if isinstance(orig_v, dict) and isinstance(update_v, dict):
-            merged[k] = _get_recursively_merged_dict(orig_v, update_v)
-        else:
-            if orig_v is not None and (isinstance(orig_v, dict) != isinstance(update_v, dict)):
+    merged = dict(base)
+    for k, v in overlay.items():
+        if k in _DEEP_MERGE_KEYS:
+            base_val = merged.get(k)
+            if isinstance(v, dict) and isinstance(base_val, dict):
+                merged[k] = {**base_val, **v}
+                continue
+            # Deep-merge key but at least one side isn't a dict: surface the
+            # silent clobber so mismatched YAML types don't get past review.
+            if base_val is not None:
                 logger.warning(
                     "Deep-merge key %r has non-dict value (base=%s, overlay=%s); "
                     "overlay will fully replace base instead of merging.",
                     k,
-                    type(orig_v).__name__,
-                    type(update_v).__name__,
+                    type(base_val).__name__,
+                    type(v).__name__,
                 )
-
-            merged[k] = update_v
+        merged[k] = v
     return merged
 
 
@@ -796,7 +779,7 @@ def _build_engine_args(
     # Per-stage StageDeployConfig values override pipeline-wide settings.
     if ds is not None:
         for k, v in asdict(ds).items():
-            if k in _STAGE_RESERVED_KEYS or v is None:
+            if k in _STAGE_NON_ENGINE_KEYS or v is None:
                 continue
             engine_args[k] = v
         engine_args.update(ds.engine_extras)

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -44,7 +44,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self.chunk_transfer_adapter = None
         if getattr(model_config, "async_chunk", False):
             self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
-        self._pending_finish_reqs: list[Request] = []
         self.input_coordinator: OmniSchedulingCoordinator | None = None
         if uses_qwen3_omni_full_payload_input_coordinator(model_config):
             self.input_coordinator = OmniSchedulingCoordinator(
@@ -112,14 +111,23 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             num_computed_tokens = request.num_computed_tokens
             required_tokens = len(request.prompt_token_ids) - num_computed_tokens
+            # async_chunk: don't schedule placeholder tokens when no new chunk is available.
             if required_tokens <= 0:
                 if (
                     self.chunk_transfer_adapter is not None
                     and request.request_id in self.chunk_transfer_adapter.finished_requests
                 ):
-                    self._pending_finish_reqs.append(request)
-                req_index += 1
-                continue
+                    # Upstream may finish with no terminal tokens; append one pad token so we can emit FINISHED.
+                    if len(request.prompt_token_ids) <= num_computed_tokens:
+                        request.prompt_token_ids.append(0)
+                        try:
+                            request._all_token_ids.append(0)  # type: ignore[attr-defined]
+                        except Exception:
+                            pass
+                    required_tokens = len(request.prompt_token_ids) - num_computed_tokens
+                else:
+                    req_index += 1
+                    continue
             num_new_tokens = min(required_tokens, token_budget)
             new_blocks = self.kv_cache_manager.allocate_slots(
                 request,
@@ -166,9 +174,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # async_chunk: wait for the first upstream chunk (don't start with placeholders).
             if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
                 if request.request_id in self.chunk_transfer_adapter.finished_requests:
-                    self.waiting.pop_request()
-                    self._pending_finish_reqs.append(request)
-                    continue
+                    request.prompt_token_ids.append(0)
+                    try:
+                        request._all_token_ids.append(0)  # type: ignore[attr-defined]
+                    except Exception:
+                        pass
                 else:
                     self.waiting.pop_request()
                     skipped_waiting_requests.prepend_request(request)
@@ -571,36 +581,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # Invariant: EngineCore returns no partial prefill outputs.
                 assert not prompt_logprobs_tensors
 
-        # Finish async_chunk requests that schedule() collected because their
-        # upstream completed with no remaining codec tokens.
-        for request in self._pending_finish_reqs:
-            if request.is_finished():
-                continue
-            request.status = RequestStatus.FINISHED_STOPPED
-            finish_reason = request.get_finished_reason()
-            finished = self._handle_stopped_request(request)
-            kv_transfer_params = None
-            if finished:
-                kv_transfer_params = self._free_request(request)
-                if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.cleanup(
-                        request.request_id,
-                        getattr(request, "external_req_id", None),
-                    )
-            outputs[request.client_index].append(
-                EngineCoreOutput(
-                    request_id=request.request_id,
-                    new_token_ids=[],
-                    finish_reason=finish_reason,
-                    stop_reason=request.stop_reason,
-                    events=request.take_events(),
-                    kv_transfer_params=kv_transfer_params,
-                    trace_headers=request.trace_headers,
-                )
-            )
-            stopped_running_reqs.add(request)
-        self._pending_finish_reqs.clear()
-
         # Remove the stopped requests from the running and waiting queues.
         if stopped_running_reqs:
             self.running = remove_all(self.running, stopped_running_reqs)

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -1383,68 +1383,6 @@ def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[
     return refresh_cache_context
 
 
-def enable_cache_for_helios(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
-    """Enable cache-dit for Helios pipeline.
-
-    Helios extends Wan2.2 with multi-term memory patches and guidance cross-attention.
-    Its transformer blocks have the same single-output signature as Wan2.2 blocks
-    (returns hidden_states only), so we use ForwardPattern.Pattern_2.
-
-    Args:
-        pipeline: The HeliosPipeline instance.
-        cache_config: DiffusionCacheConfig instance with cache configuration.
-
-    Returns:
-        A refresh function that can be called to update cache context with new num_inference_steps.
-    """
-    db_cache_config = _build_db_cache_config(cache_config)
-
-    calibrator_config = None
-    if cache_config.enable_taylorseer:
-        taylorseer_order = cache_config.taylorseer_order
-        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
-        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
-
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=pipeline.transformer,
-            blocks=[pipeline.transformer.blocks],
-            forward_pattern=[ForwardPattern.Pattern_2],
-            params_modifiers=[
-                ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config),
-            ],
-            has_separate_cfg=False,
-        ),
-        cache_config=db_cache_config,
-        calibrator_config=calibrator_config,
-    )
-
-    logger.info(
-        f"Enabling cache-dit on Helios transformer: "
-        f"Fn={db_cache_config.Fn_compute_blocks}, "
-        f"Bn={db_cache_config.Bn_compute_blocks}, "
-        f"W={db_cache_config.max_warmup_steps}, "
-    )
-
-    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
-        if cache_config.scm_steps_mask_policy is None:
-            cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
-        else:
-            cache_dit.refresh_context(
-                pipeline.transformer,
-                cache_config=DBCacheConfig().reset(
-                    num_inference_steps=num_inference_steps,
-                    steps_computation_mask=cache_dit.steps_mask(
-                        mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_inference_steps
-                    ),
-                    steps_computation_policy=cache_config.scm_steps_policy,
-                ),
-                verbose=verbose,
-            )
-
-    return refresh_cache_context
-
-
 def enable_cache_for_hunyuan_video_15(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for HunyuanVideo 1.5 pipeline.
 
@@ -1525,7 +1463,6 @@ CUSTOM_DIT_ENABLERS.update(
         "ErnieImagePipeline": enable_cache_for_ernie_image,
         "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15,
         "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15,
-        "HeliosPipeline": enable_cache_for_helios,
     }
 )
 

vllm_omni/diffusion/executor/abstract.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 from vllm.utils.import_utils import resolve_obj_by_qualname
@@ -102,14 +101,6 @@ class DiffusionExecutor(ABC):
         """Check if the executor and workers are healthy."""
         pass
 
-    def register_failure_callback(self, callback: Callable[[], None]) -> None:
-        """Register a callback invoked when the executor fatally fails.
-
-        Executors without a background failure monitor can keep the default
-        no-op implementation.
-        """
-        return None
-
     @abstractmethod
     def shutdown(self) -> None:
         """Shutdown the executor and release resources."""

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -64,11 +64,8 @@ class InlineStageDiffusionClient(StageClientBase):
 
         self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
         self._tasks: dict[str, asyncio.Task] = {}
-        self._engine_dead = False
         self._shutting_down = False
 
-        self._engine.executor.register_failure_callback(self._mark_engine_dead)
-
         logger.info(
             "[InlineStageDiffusionClient] stage-%s [rep-%s] initialized inline (batch_size=%d)",
             self.stage_id,
@@ -80,16 +77,6 @@ class InlineStageDiffusionClient(StageClientBase):
         """Load model metadata from HuggingFace and populate od_config fields."""
         self.od_config.enrich_config()
 
-    def _mark_engine_dead(self) -> None:
-        if self._engine_dead:
-            return
-        self._engine_dead = True
-        logger.error(
-            "[InlineStageDiffusionClient] stage-%s [rep-%s] diffusion executor died unexpectedly.",
-            self.stage_id,
-            self.replica_id,
-        )
-
     # ------------------------------------------------------------------
     # Request processing
     # ------------------------------------------------------------------
@@ -262,8 +249,6 @@ class InlineStageDiffusionClient(StageClientBase):
         try:
             return self._output_queue.get_nowait()
         except asyncio.QueueEmpty:
-            if self._engine_dead:
-                raise EngineDeadError(f"Stage-{self.stage_id} inline diffusion engine is dead")
             return None
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:
@@ -366,11 +351,7 @@ class InlineStageDiffusionClient(StageClientBase):
         """Check if the inline diffusion engine and its workers are healthy."""
         if self._shutting_down:
             raise EngineDeadError("InlineStageDiffusionClient is shutting down")
-        try:
-            self._engine.executor.check_health()
-        except EngineDeadError:
-            self._mark_engine_dead()
-            raise
+        self._engine.executor.check_health()
 
     def shutdown(self) -> None:
         self._shutting_down = True

vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py

@@ -26,10 +26,10 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.models.interface import SupportAudioOutput
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
 from vllm_omni.model_executor.models.omnivoice.duration import RuleDurationEstimator
 from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import OmniVoiceDecoder
 from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator
-from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
 from vllm_omni.utils.speaker_cache import get_speaker_cache
 
 try:

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -796,14 +796,6 @@ class WorkerProc:
         """Worker initialization and execution loops."""
         from vllm_omni.plugins import load_omni_general_plugins
 
-        # Set process title for visibility in nvidia-smi / htop (optional, non-fatal)
-        try:
-            import setproctitle
-
-            setproctitle.setproctitle(f"vLLM-Omni::StageDiffusionProc-{rank}")
-        except ImportError:
-            pass  # setproctitle not installed, skip process title setting
-
         load_omni_general_plugins()
         worker_proc = WorkerProc(
             od_config,

vllm_omni/engine/arg_utils.py

@@ -35,9 +35,13 @@ def _register_omni_hf_configs() -> None:
     try:
         from transformers import AutoConfig
 
+        from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
+        from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
         from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
             Qwen3TTSConfig,
         )
+        from vllm_omni.transformers_utils.configs.voxcpm import VoxCPMConfig
+        from vllm_omni.transformers_utils.configs.voxcpm2 import VoxCPM2Config
     except Exception as exc:  # pragma: no cover - best-effort optional registration
         logger.warning("Skipping omni HF config registration due to import error: %s", exc)
         return
@@ -52,6 +56,10 @@ def _register_omni_hf_configs() -> None:
 
     for model_type, config_cls in [
         ("qwen3_tts", Qwen3TTSConfig),
+        ("cosyvoice3", CosyVoice3Config),
+        ("omnivoice", OmniVoiceConfig),
+        ("voxcpm", VoxCPMConfig),
+        ("voxcpm2", VoxCPM2Config),
     ]:
         try:
             AutoConfig.register(model_type, config_cls)

vllm_omni/entrypoints/omni_base.py

@@ -218,18 +218,9 @@ class OmniBase(PDDisaggregationMixin):
         """Expose engine stage configs for PD disaggregation detection and validation."""
         return self.engine.stage_configs
 
-    def _has_dead_stage(self) -> bool:
-        for stage_client in self.engine.stage_clients:
-            if getattr(stage_client, "_engine_dead", False):
-                return True
-            resources = getattr(stage_client, "resources", None)
-            if resources is not None and getattr(resources, "engine_dead", False):
-                return True
-        return False
-
     @property
     def is_running(self) -> bool:
-        return self.engine.is_alive() and not self._has_dead_stage()
+        return self.engine.is_alive()
 
     @property
     def errored(self) -> bool:
@@ -242,7 +233,15 @@ class OmniBase(PDDisaggregationMixin):
         ``resources.engine_dead`` (StageEngineCoreClient / AsyncMPClient)
         since the two client types store the flag differently.
         """
-        return not self.engine.is_alive() or self._has_dead_stage()
+        if not self.engine.is_alive():
+            return True
+        for stage_client in self.engine.stage_clients:
+            if getattr(stage_client, "_engine_dead", False):
+                return True
+            resources = getattr(stage_client, "resources", None)
+            if resources is not None and getattr(resources, "engine_dead", False):
+                return True
+        return False
 
     def check_health(self) -> None:
         if not self.engine.is_alive():

vllm_omni/model_executor/models/cosyvoice3/config.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-from transformers import AutoConfig
 from transformers.configuration_utils import PretrainedConfig
 
 
@@ -132,8 +131,3 @@ class CosyVoice3Config(PretrainedConfig):
                 "cond_channels": 512,
             },
         }
-
-
-AutoConfig.register("cosyvoice3", CosyVoice3Config)
-
-__all__ = ["CosyVoice3Config"]

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -33,6 +33,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct
+from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.utils import (
     concat_text_with_prompt_ids,
     extract_speech_feat,
@@ -41,7 +42,6 @@ from vllm_omni.model_executor.models.cosyvoice3.utils import (
     extract_text_token,
 )
 from vllm_omni.model_executor.models.output_templates import OmniOutput
-from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 from vllm_omni.utils.speaker_cache import get_speaker_cache
 
 logger = init_logger(__name__)

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py

@@ -29,7 +29,7 @@ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
     CausalHiFTGenerator,
 )
 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer
-from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
+from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 
 logger = init_logger(__name__)
 

vllm_omni/model_executor/models/omnivoice/config.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """OmniVoice configuration for vLLM-Omni two-stage pipeline."""
 
-from transformers import AutoConfig
 from transformers.configuration_utils import PretrainedConfig
 
 
@@ -80,8 +79,3 @@ class OmniVoiceConfig(PretrainedConfig):
 
         # Serving
         self.speculative_config = None
-
-
-AutoConfig.register("omnivoice", OmniVoiceConfig)
-
-__all__ = ["OmniVoiceConfig"]

vllm_omni/model_executor/models/omnivoice/omnivoice.py

@@ -37,8 +37,8 @@ from vllm.multimodal.processing import (
 )
 from vllm.sequence import IntermediateTensors
 
+from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
 from vllm_omni.model_executor.models.output_templates import OmniOutput
-from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
 
 logger = init_logger(__name__)
 

vllm_omni/model_executor/models/omnivoice/omnivoice_decoder.py

@@ -23,7 +23,7 @@ import torch
 import torch.nn as nn
 from vllm.logger import init_logger
 
-from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
 
 logger = init_logger(__name__)
 

vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py

@@ -19,7 +19,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from vllm.logger import init_logger
 
-from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig
+from vllm_omni.model_executor.models.omnivoice.config import OmniVoiceConfig
 
 logger = init_logger(__name__)
 

vllm_omni/model_executor/models/voxcpm/__init__.py

@@ -1,5 +1,7 @@
+from .configuration_voxcpm import VoxCPMConfig
 from .voxcpm import VoxCPMForConditionalGeneration
 
 __all__ = [
+    "VoxCPMConfig",
     "VoxCPMForConditionalGeneration",
 ]

vllm_omni/model_executor/models/voxcpm/configuration_voxcpm.py

@@ -0,0 +1,3 @@
+from vllm_omni.transformers_utils.configs.voxcpm import VoxCPMConfig
+
+__all__ = ["VoxCPMConfig"]

vllm_omni/transformers_utils/configs/__init__.py

@@ -20,8 +20,6 @@ _CLASS_TO_MODULE: dict[str, str] = {
     "VoxCPMConfig": "vllm_omni.transformers_utils.configs.voxcpm",
     "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2",
     "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts",
-    "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3",
-    "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice",
     "BailingMoeV2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "BailingMM2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni",
     "MingFlashOmniConfig": "vllm_omni.transformers_utils.configs.ming_flash_omni",
@@ -40,8 +38,6 @@ __all__ = [
     "VoxCPMConfig",
     "VoxCPM2Config",
     "VoxtralTTSConfig",
-    "CosyVoice3Config",
-    "OmniVoiceConfig",
     "BailingMoeV2Config",
     "BailingMM2Config",
     "MingFlashOmniConfig",
@@ -65,11 +61,9 @@ def __dir__():
 
 # Eagerly import all config modules so their AutoConfig.register() side-effects
 # run as soon as `vllm_omni.transformers_utils.configs` is imported.
-from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni  # noqa: F401, E402
-from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import voxcpm as _voxcpm  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import voxcpm2 as _voxcpm2  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import voxtral_tts as _voxtral_tts  # noqa: F401, E402