[Model] Add Dynin-omni model in vllm-omni (#1759)

Signed-off-by: kdg6245@snu.ac.kr <kdg6245@snu.ac.kr>
Signed-off-by: Yejoon Lee (IPAI) <leeyejoon@snu.ac.kr>
Signed-off-by: aidas (arpa-kt) <aidaslab@gmail.com>
Signed-off-by: Dogeun Kim <82812668+DOGEUNNKIM@users.noreply.github.com>
Co-authored-by: Yejoon Lee (IPAI) <leeyejoon@snu.ac.kr>
Co-authored-by: aidas (arpa-kt) <aidaslab@gmail.com>

docs/models/supported_models.md

@@ -60,5 +60,6 @@ th {
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
+|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
 
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.

examples/offline_inference/dynin_omni/README.md

@@ -0,0 +1,110 @@
+# Dynin-Omni Offline End2End Example
+
+This folder contains a unified offline inference entrypoint:
+
+- `end2end.py`
+
+## 1. Environment Setup
+
+Run from repository root:
+
+```bash
+cd <REPO_ROOT>
+```
+
+If needed, install this repo in editable mode:
+
+```bash
+pip install -e .
+```
+
+## 2. Extra Dependencies (EMOVA)
+
+Install the following packages for EMOVA-related components:
+
+```bash
+pip install \
+  "phonemizer==3.3.0" \
+  "Unidecode==1.4.0" \
+  "hydra-core==1.3.2" \
+  "pytorch-lightning==1.1.0" \
+  "wget==3.2" \
+  "wrapt==2.1.1" \
+  "onnx==1.20.1" \
+  "frozendict==2.4.7" \
+  "inflect==7.5.0" \
+  "braceexpand==0.1.7" \
+  "webdataset==1.0.2" \
+  "torch-stft==0.1.4" \
+  "editdistance==0.8.1"
+```
+
+## 3. Hardware and VRAM Requirements
+
+This example uses a 3-stage pipeline on one GPU by default
+([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
+
+- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
+- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
+- Stage-2 (`token2audio`): `gpu_memory_utilization: 0.1`
+
+### Requested GPU Memory Budget from `gpu_memory_utilization`
+
+| Stage | Utilization | A100 80GB | H200 141GB |
+| :-- | :-- | :-- | :-- |
+| Stage-0 (token2text) | 0.5 | ~40.0 GB | ~70.5 GB |
+| Stage-1 (token2image) | 0.1 | ~8.0 GB | ~14.1 GB |
+| Stage-2 (token2audio) | 0.1 | ~8.0 GB | ~14.1 GB |
+| Total requested budget | 0.7 | ~56.0 GB | ~98.7 GB |
+
+### Observed Runtime Signal (from your log)
+
+- Stage-0 reported: `Model loading took 15.12 GiB memory` (weights footprint signal).
+- Stages 1/2 can still add runtime memory depending on task path and backend allocations.
+- Keep extra headroom for CUDA/PyTorch overhead and temporary allocations.
+
+### GPU Compatibility
+
+- Confirmed target GPUs for this setup: **NVIDIA H200**, **NVIDIA A100**.
+- CI/e2e coverage in this repo also includes CUDA **L4** markers for Dynin tests.
+
+## 4. End2End Run Examples
+
+```bash
+# t2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task t2t --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
+
+# i2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task i2t --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text "Please describe this image in detail."
+
+# s2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task s2t --model snu-aidas/Dynin-Omni --audio <AUDIO_PATH> --text "Transcribe the given audio."
+
+# t2i
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task t2i --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
+
+# v2t
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task v2t --model snu-aidas/Dynin-Omni --video <VIDEO_PATH> --text "Describe this video in detail."
+
+# i2i
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task i2i --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text <INSTRUCTION_TEXT>
+
+# t2s
+python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
+  --task t2s --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
+```
+
+## 5. Notes
+
+- Outputs are saved under task-specific directories in `/tmp` by default.
+- You can override output path with `--output-dir`.
+- If you want to force local config resolution, pass `--dynin-config-path <PATH_TO_DYNIN_OMNI_YAML>`.
+- If you see the warning
+  `max_num_batched_tokens (32768) exceeds max_num_seqs * max_model_len (4096)`,
+  reduce `max_num_batched_tokens` in stage config (for example, `4096` in CI config).

examples/online_serving/dynin_omni/README.md

@@ -0,0 +1,97 @@
+# Dynin-Omni Online Serving Example
+
+## Installation
+
+Please refer to [README.md](../../../README.md).
+
+## Launch the Server
+
+First, find the `transformers_modules` path:
+
+```bash
+python - <<'PY'
+from transformers.utils.hub import HF_MODULES_CACHE
+print(HF_MODULES_CACHE)
+PY
+```
+
+Then export it for both `PYTHONPATH` and `HF_MODULES_CACHE`:
+
+```bash
+export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
+export HF_MODULES_CACHE=<transformers_modules_path>
+```
+
+Run from repository root:
+
+```bash
+vllm-omni serve snu-aidas/Dynin-Omni \
+  --omni \
+  --port 8091 \
+  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+```
+
+If `vllm-omni` is not in PATH, run:
+
+```bash
+PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
+  --omni \
+  --port 8091 \
+  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
+```
+
+Wait until the server logs show both `All stages initialized successfully` and
+`Application startup complete.` before sending requests.
+
+## Send Requests via Python Client
+
+Move to the example directory:
+
+```bash
+cd examples/online_serving/dynin_omni
+```
+
+### Text -> Image
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+  --query-type t2i \
+  --prompt "A realistic indoor living room with natural daylight."
+```
+
+### Image -> Image
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+  --query-type i2i \
+  --image-path ../../offline_inference/dynin_omni/data/image/sofa_under_water.jpg \
+  --prompt "Transform this surreal underwater setting into a realistic indoor living room while preserving the sofa layout."
+```
+
+### Text -> Speech
+
+```bash
+python openai_chat_completion_client_for_multimodal_generation.py \
+  --query-type t2s \
+  --prompt "Hello. This is Dynin-omni."
+```
+
+## CLI Arguments
+
+- `--query-type` (`t2i|t2s|i2i`)
+- `--model` (default: `snu-aidas/Dynin-Omni`)
+- `--host` / `--port` (OpenAI-compatible vLLM endpoint)
+- `--prompt` (custom text)
+- `--image-path` (required for `i2i`)
+- `--modalities` (optional output modalities override)
+- `--output-dir` (default: `/tmp/dynin_online_outputs`)
+
+## Notes
+
+- This client currently supports only `t2i`, `t2s`, and `i2i`.
+- `t2t` is intentionally not exposed in this online example.
+- This example intentionally uses the OpenAI-compatible chat completion endpoint.
+- Task routing for non-text outputs relies on Dynin task trigger tokens (`<|t2i|>`, `<|i2i|>`, `<|t2s|>`) injected by the client.
+- Outputs are saved under `/tmp/dynin_online_outputs` by default.
+- Dynin stage-0 warmup can take a while on first startup; do not send requests before startup completes.
+- Dynin itself can execute text-returning tasks such as `t2t`, `s2t`, `i2t`, and `v2t`, but this online serving example currently runs stage-0 in `generation` mode. In that path, the generation worker does not surface the final text as `output.text`, so OpenAI chat responses for those text-output tasks may complete internally but still return empty text.

examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -0,0 +1,342 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import argparse
+import base64
+import json
+import mimetypes
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+DEFAULT_MODEL = "snu-aidas/Dynin-Omni"
+DEFAULT_OUTPUT_DIR = "/tmp/dynin_online_outputs"
+
+QUERY_CHOICES = ("t2i", "t2s", "i2i")
+DEFAULT_PROMPT_BY_QUERY = {
+    "t2i": "A high quality detailed living room interior photo.",
+    "t2s": "Please read this sentence naturally: Hello from Dynin-Omni online serving.",
+    "i2i": "Transform this image into a realistic indoor living room while preserving layout.",
+}
+DEFAULT_MODALITIES_BY_QUERY = {
+    "t2i": ["image"],
+    "t2s": ["audio"],
+    "i2i": ["image"],
+}
+OFFLINE_PARITY_STAGE_COUNT = 3
+OFFLINE_PARITY_STAGE_SAMPLING = {
+    "max_tokens": 1,
+    "temperature": 0.0,
+    "top_p": 1.0,
+    "detokenize": False,
+}
+
+
+def _infer_mime_type(path: Path) -> str:
+    mime_type, _ = mimetypes.guess_type(str(path))
+    return mime_type or "application/octet-stream"
+
+
+def _encode_file_as_data_url(path: Path) -> str:
+    mime_type = _infer_mime_type(path)
+    raw = path.read_bytes()
+    encoded = base64.b64encode(raw).decode("utf-8")
+    return f"data:{mime_type};base64,{encoded}"
+
+
+def _to_image_url(path_or_url: str) -> str:
+    value = str(path_or_url)
+    if value.startswith(("http://", "https://", "data:image/")):
+        return value
+    path = Path(value).expanduser().resolve()
+    if not path.exists():
+        raise FileNotFoundError(f"Image file not found: {path}")
+    return _encode_file_as_data_url(path)
+
+
+def _build_user_content(query_type: str, prompt: str, image_path: str | None) -> list[dict[str, Any]]:
+    if query_type == "t2i":
+        return [{"type": "text", "text": f"<|t2i|> {prompt}"}]
+
+    if query_type == "t2s":
+        return [{"type": "text", "text": f"<|t2s|> {prompt}"}]
+
+    if query_type == "i2i":
+        if not image_path:
+            raise ValueError("--image-path is required for query type i2i")
+        return [
+            {"type": "text", "text": f"<|i2i|> {prompt}"},
+            {"type": "image_url", "image_url": {"url": _to_image_url(image_path)}},
+        ]
+
+    raise ValueError(f"Unsupported query_type: {query_type}")
+
+
+def _collect_text_from_content(content: Any) -> list[str]:
+    texts: list[str] = []
+    if isinstance(content, str):
+        stripped = content.strip()
+        if stripped:
+            texts.append(stripped)
+        return texts
+
+    if isinstance(content, dict):
+        for key in ("text", "content", "value", "output_text"):
+            text_value = content.get(key)
+            if isinstance(text_value, str) and text_value.strip():
+                texts.append(text_value.strip())
+        return texts
+
+    if isinstance(content, list):
+        for item in content:
+            texts.extend(_collect_text_from_content(item))
+        return texts
+
+    content_text = getattr(content, "text", None)
+    if isinstance(content_text, str) and content_text.strip():
+        texts.append(content_text.strip())
+    content_value = getattr(content, "content", None)
+    if isinstance(content_value, str) and content_value.strip():
+        texts.append(content_value.strip())
+    output_text = getattr(content, "output_text", None)
+    if isinstance(output_text, str) and output_text.strip():
+        texts.append(output_text.strip())
+    return texts
+
+
+def _extract_text_outputs(chat_completion: Any) -> list[str]:
+    texts: list[str] = []
+    for choice in getattr(chat_completion, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        if message is None:
+            continue
+        content = getattr(message, "content", None)
+        texts.extend(_collect_text_from_content(content))
+        reasoning_content = getattr(message, "reasoning_content", None)
+        if isinstance(reasoning_content, str) and reasoning_content.strip():
+            texts.append(reasoning_content.strip())
+        choice_text = getattr(choice, "text", None)
+        if isinstance(choice_text, str) and choice_text.strip():
+            texts.append(choice_text.strip())
+    top_level_output_text = getattr(chat_completion, "output_text", None)
+    if isinstance(top_level_output_text, str) and top_level_output_text.strip():
+        texts.append(top_level_output_text.strip())
+    return texts
+
+
+def _extract_image_data_urls(chat_completion: Any) -> list[str]:
+    urls: list[str] = []
+    for choice in getattr(chat_completion, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        if message is None:
+            continue
+        content = getattr(message, "content", None)
+        if not isinstance(content, list):
+            continue
+        for item in content:
+            if not isinstance(item, dict):
+                continue
+            if item.get("type") != "image_url":
+                continue
+            image_url = (item.get("image_url") or {}).get("url")
+            if isinstance(image_url, str) and image_url.startswith("data:image"):
+                urls.append(image_url)
+    return urls
+
+
+def _extract_audio_payloads(chat_completion: Any) -> list[bytes]:
+    payloads: list[bytes] = []
+    for choice in getattr(chat_completion, "choices", []) or []:
+        message = getattr(choice, "message", None)
+        if message is None:
+            continue
+        message_audio = getattr(message, "audio", None)
+        if message_audio is None:
+            continue
+        data_b64 = getattr(message_audio, "data", None)
+        if isinstance(data_b64, str) and data_b64:
+            try:
+                payloads.append(base64.b64decode(data_b64))
+            except Exception:
+                continue
+    return payloads
+
+
+def _decode_data_url(data_url: str) -> tuple[bytes, str]:
+    header, data = data_url.split(",", 1)
+    mime_type = "image/png"
+    if ";" in header and ":" in header:
+        mime_type = header.split(":", 1)[1].split(";", 1)[0]
+    return base64.b64decode(data), mime_type
+
+
+def _image_extension_from_mime(mime_type: str) -> str:
+    if mime_type == "image/jpeg":
+        return ".jpg"
+    if mime_type == "image/webp":
+        return ".webp"
+    if mime_type == "image/gif":
+        return ".gif"
+    return ".png"
+
+
+def _save_outputs(
+    *,
+    query_type: str,
+    chat_completion: Any,
+    output_dir: Path,
+) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    stamp = time.strftime("%Y%m%d_%H%M%S")
+
+    text_outputs = _extract_text_outputs(chat_completion)
+    image_data_urls = _extract_image_data_urls(chat_completion)
+    audio_payloads = _extract_audio_payloads(chat_completion)
+
+    if text_outputs:
+        text_path = output_dir / f"{query_type}_{stamp}.txt"
+        text_path.write_text("\n\n".join(text_outputs) + "\n", encoding="utf-8")
+        print(f"[dynin-online] text saved: {text_path}")
+        print(text_outputs[0])
+
+    for idx, image_url in enumerate(image_data_urls):
+        image_bytes, mime_type = _decode_data_url(image_url)
+        ext = _image_extension_from_mime(mime_type)
+        image_path = output_dir / f"{query_type}_{stamp}_{idx}{ext}"
+        image_path.write_bytes(image_bytes)
+        print(f"[dynin-online] image saved: {image_path}")
+
+    for idx, audio_bytes in enumerate(audio_payloads):
+        audio_path = output_dir / f"{query_type}_{stamp}_{idx}.wav"
+        audio_path.write_bytes(audio_bytes)
+        print(f"[dynin-online] audio saved: {audio_path}")
+
+    if not text_outputs and not image_data_urls and not audio_payloads:
+        print("[dynin-online] no output extracted from response")
+        raw_path = output_dir / f"{query_type}_{stamp}_raw_response.json"
+        try:
+            if hasattr(chat_completion, "model_dump_json"):
+                serialized = chat_completion.model_dump_json(indent=2)
+            else:
+                if hasattr(chat_completion, "model_dump"):
+                    raw_payload: Any = chat_completion.model_dump(mode="json")
+                else:
+                    raw_payload = chat_completion
+                try:
+                    serialized = json.dumps(raw_payload, ensure_ascii=False, indent=2)
+                except Exception:
+                    serialized = json.dumps({"repr": repr(raw_payload)}, ensure_ascii=False, indent=2)
+            raw_path.write_text(serialized + "\n", encoding="utf-8")
+            print(f"[dynin-online] raw response saved: {raw_path}")
+        except Exception:
+            pass
+
+
+def _build_offline_parity_sampling_params_list() -> list[dict[str, Any]]:
+    return [dict(OFFLINE_PARITY_STAGE_SAMPLING) for _ in range(OFFLINE_PARITY_STAGE_COUNT)]
+
+
+def run_request(args: argparse.Namespace) -> None:
+    from openai import OpenAI
+
+    client = OpenAI(
+        api_key="EMPTY",
+        base_url=f"http://{args.host}:{args.port}/v1",
+    )
+    prompt = args.prompt.strip() if args.prompt else DEFAULT_PROMPT_BY_QUERY[args.query_type]
+    user_content = _build_user_content(
+        query_type=args.query_type,
+        prompt=prompt,
+        image_path=args.image_path,
+    )
+    if args.modalities:
+        modalities = [item.strip() for item in args.modalities.split(",") if item.strip()]
+    else:
+        modalities = DEFAULT_MODALITIES_BY_QUERY[args.query_type]
+
+    extra_body = {
+        "sampling_params_list": _build_offline_parity_sampling_params_list(),
+    }
+    chat_completion = client.chat.completions.create(
+        model=args.model,
+        messages=[{"role": "user", "content": user_content}],
+        modalities=modalities,
+        extra_body=extra_body,
+    )
+    _save_outputs(
+        query_type=args.query_type,
+        chat_completion=chat_completion,
+        output_dir=Path(args.output_dir).expanduser(),
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Dynin-Omni online chat completion client")
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="t2i",
+        choices=QUERY_CHOICES,
+        help="Dynin query type",
+    )
+    parser.add_argument(
+        "--model",
+        "-m",
+        type=str,
+        default=DEFAULT_MODEL,
+        help="Model name/path",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="localhost",
+        help="Host/IP of the vLLM Omni API server",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8091,
+        help="Port of the vLLM Omni API server",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        type=str,
+        default="",
+        help="Custom prompt text",
+    )
+    parser.add_argument(
+        "--image-path",
+        "-i",
+        type=str,
+        default=None,
+        help="Image path/URL for i2i",
+    )
+    parser.add_argument(
+        "--modalities",
+        type=str,
+        default="",
+        help="Comma-separated output modalities override (e.g., text,image,audio)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        "-o",
+        type=str,
+        default=DEFAULT_OUTPUT_DIR,
+        help="Directory to save outputs",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
+    run_request(args)
+
+
+if __name__ == "__main__":
+    main()

tests/e2e/offline_inference/test_dynin_omni.py

@@ -0,0 +1,419 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E offline smoke tests for Dynin-Omni.
+
+- model: "snu-aidas/Dynin-Omni"
+- stage config: tests/e2e/stage_configs/dynin_omni_ci.yaml
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from transformers import AutoTokenizer
+
+from tests.conftest import OmniRunner
+from tests.utils import hardware_test
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_DEFAULT_DYNIN_CONFIG_PATH: Path | None = None
+_DEFAULT_STAGE_CONFIG_PATH = _REPO_ROOT / "tests" / "e2e" / "stage_configs" / "dynin_omni_ci.yaml"
+
+models = ["snu-aidas/Dynin-Omni"]
+stage_configs = [str(_DEFAULT_STAGE_CONFIG_PATH)]
+test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
+
+DYNIN_CONFIG_PATH = str(_DEFAULT_DYNIN_CONFIG_PATH) if _DEFAULT_DYNIN_CONFIG_PATH is not None else None
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.omni,
+]
+
+
+# prompting util
+def _build_mmu_prompt(tokenizer: Any, question: str, dynin_config_path: str | None) -> dict[str, Any]:
+    encoded = tokenizer(question, return_tensors="pt", add_special_tokens=True)
+    token_ids = [int(v) for v in encoded["input_ids"][0].tolist()]
+    attention_mask = [int(v) for v in encoded["attention_mask"][0].tolist()]
+    additional_information: dict[str, Any] = {
+        "task": ["mmu"],
+        "detok_id": [0],
+        "prompt_length": [len(token_ids)],
+        "attention_mask": [attention_mask],
+        "max_new_tokens": [64],
+        "steps": [64],
+        "block_length": [16],
+        "temperature": [0.0],
+    }
+    if dynin_config_path:
+        additional_information["dynin_config_path"] = [str(dynin_config_path)]
+    return {
+        "prompt_token_ids": token_ids,
+        "additional_information": additional_information,
+        "modalities": ["text"],
+    }
+
+
+def _build_mmu_multimodal_prompt(
+    tokenizer: Any,
+    question: str,
+    dynin_config_path: str | None,
+    *,
+    image: Any | None = None,
+    audio: tuple[np.ndarray, int] | None = None,
+) -> dict[str, Any]:
+    if image is None and audio is None:
+        raise ValueError("At least one multimodal input (image or audio) must be provided.")
+
+    prefix_chunks: list[str] = []
+    mm_data: dict[str, Any] = {}
+    if image is not None:
+        prefix_chunks.append("<|soi|><|image|><|eoi|>")
+        mm_data["image"] = image
+    if audio is not None:
+        prefix_chunks.append("<|soa|><|audio|><|eoa|>")
+        mm_data["audio"] = audio
+
+    prefixed_question = " ".join(prefix_chunks + [question]).strip()
+    prompt = _build_mmu_prompt(
+        tokenizer=tokenizer,
+        question=prefixed_question,
+        dynin_config_path=dynin_config_path,
+    )
+    prompt["multi_modal_data"] = mm_data
+    prompt["modalities"] = ["text"]
+    return prompt
+
+
+def _generate_synthetic_image(width: int = 224, height: int = 224) -> np.ndarray:
+    x = np.linspace(0, 255, width, dtype=np.uint8)
+    y = np.linspace(0, 255, height, dtype=np.uint8)[:, None]
+    red = np.tile(x, (height, 1))
+    green = np.tile(y, (1, width))
+    blue = ((red.astype(np.uint16) + green.astype(np.uint16)) // 2).astype(np.uint8)
+    return np.stack([red, green, blue], axis=-1)
+
+
+def _generate_synthetic_audio(duration_s: int = 5, sample_rate: int = 48_000) -> tuple[np.ndarray, int]:
+    t = np.linspace(0, duration_s, int(sample_rate * duration_s), endpoint=False, dtype=np.float32)
+    waveform = 0.1 * np.sin(2.0 * np.pi * 440.0 * t)
+    return waveform.astype(np.float32), sample_rate
+
+
+# prompting util
+def _build_t2s_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
+    # Bypass stage-0 generation and directly validate token->audio decode path.
+    generated_audio_token_ids = [int(v) for v in ([10, 11, 12, 13, 14] * 32)]
+    additional_information: dict[str, Any] = {
+        "task": ["t2s"],
+        "detok_id": [1],
+        "generated_token_ids": [generated_audio_token_ids],
+        "audio_codebook_size": [4096],
+    }
+    if dynin_config_path:
+        additional_information["dynin_config_path"] = [str(dynin_config_path)]
+    return {
+        "prompt_token_ids": [0],
+        "additional_information": additional_information,
+        "modalities": ["audio"],
+    }
+
+
+# prompting util
+def _build_t2i_decode_prompt(dynin_config_path: str | None) -> dict[str, Any]:
+    # Bypass stage-0 generation and directly validate token->image decode path.
+    # MAGVIT decode path expects a square token grid; 1024 tokens -> 32x32.
+    generated_image_token_ids = [int(v) for v in ([10, 11, 12, 13, 14, 15, 16, 17] * 128)]
+    additional_information: dict[str, Any] = {
+        "task": ["t2i"],
+        "detok_id": [2],
+        "generated_token_ids": [generated_image_token_ids],
+        "codebook_size": [8192],
+    }
+    if dynin_config_path:
+        additional_information["dynin_config_path"] = [str(dynin_config_path)]
+    return {
+        "prompt_token_ids": [0],
+        "additional_information": additional_information,
+        "modalities": ["image"],
+    }
+
+
+def _configure_dynin_config_env() -> None:
+    if DYNIN_CONFIG_PATH:
+        os.environ["DYNIN_CONFIG_PATH"] = str(DYNIN_CONFIG_PATH)
+    else:
+        os.environ.pop("DYNIN_CONFIG_PATH", None)
+
+
+def _is_finished_request_output(request_output: Any) -> bool:
+    if request_output is None:
+        return False
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        if req is not None and bool(getattr(req, "finished", False)):
+            return True
+    return False
+
+
+def _find_stage_output(outputs: list[Any], output_type: str) -> Any | None:
+    matched = [
+        stage_output for stage_output in outputs if getattr(stage_output, "final_output_type", None) == output_type
+    ]
+    if not matched:
+        return None
+
+    # Prefer the latest finished chunk to avoid picking an intermediate stream output.
+    for stage_output in reversed(matched):
+        if _is_finished_request_output(getattr(stage_output, "request_output", None)):
+            return stage_output
+    return matched[-1]
+
+
+def _to_token_list(value: Any) -> list[int]:
+    if value is None:
+        return []
+    if hasattr(value, "detach"):
+        value = value.detach()
+    if hasattr(value, "cpu"):
+        value = value.cpu()
+    if hasattr(value, "flatten"):
+        value = value.flatten().tolist()
+    if isinstance(value, tuple):
+        value = list(value)
+    if not isinstance(value, list):
+        return []
+    out: list[int] = []
+    for token in value:
+        if isinstance(token, bool):
+            continue
+        try:
+            out.append(int(token))
+        except Exception:
+            continue
+    return out
+
+
+def _extract_text(stage_output: Any, tokenizer: Any | None = None) -> str:
+    request_output = getattr(stage_output, "request_output", None)
+    if request_output is None:
+        return ""
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        completions = getattr(req, "outputs", None) or []
+        if not completions:
+            continue
+        completion = completions[0]
+        mm_out = (
+            getattr(completion, "multimodal_output", None)
+            or getattr(req, "multimodal_output", None)
+            or getattr(stage_output, "multimodal_output", None)
+            or {}
+        )
+        text = mm_out.get("text")
+        if isinstance(text, list) and text:
+            text = text[-1]
+        if isinstance(text, str) and text.strip():
+            return text.strip()
+        if tokenizer is not None:
+            for key in ("text_tokens", "token_ids"):
+                token_ids = _to_token_list(mm_out.get(key))
+                if not token_ids:
+                    continue
+                decoded = tokenizer.decode(token_ids, skip_special_tokens=True)
+                if isinstance(decoded, str) and decoded.strip():
+                    return decoded.strip()
+        fallback = getattr(completion, "text", None)
+        if isinstance(fallback, str) and fallback.strip():
+            return fallback.strip()
+    return ""
+
+
+def _extract_audio(stage_output: Any) -> Any | None:
+    request_output = getattr(stage_output, "request_output", None)
+    if request_output is None:
+        return None
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        completions = getattr(req, "outputs", None) or []
+        if not completions:
+            continue
+        completion = completions[0]
+        mm_out = getattr(completion, "multimodal_output", None) or {}
+        if "audio" in mm_out:
+            return mm_out["audio"]
+    return None
+
+
+def _extract_image(stage_output: Any) -> Any | None:
+    request_output = getattr(stage_output, "request_output", None)
+    if request_output is None:
+        return None
+    req_list = request_output if isinstance(request_output, list) else [request_output]
+    for req in req_list:
+        completions = getattr(req, "outputs", None) or []
+        if not completions:
+            continue
+        completion = completions[0]
+        mm_out = getattr(completion, "multimodal_output", None) or {}
+        if "image" in mm_out:
+            return mm_out["image"]
+    return None
+
+
+def _numel(value: Any) -> int:
+    if value is None:
+        return 0
+    if isinstance(value, torch.Tensor):
+        return int(value.numel())
+    shape = getattr(value, "shape", None)
+    if shape is not None:
+        try:
+            total = 1
+            for dim in shape:
+                total *= int(dim)
+            return int(total)
+        except Exception:
+            pass
+    if isinstance(value, (list, tuple)):
+        return len(value)
+    return 0
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_t2i_decode_to_image(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    prompt = _build_t2i_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    image_output = _find_stage_output(outputs, "image")
+    assert image_output is not None
+    image_value = _extract_image(image_output)
+    assert image_value is not None
+    assert _numel(image_value) > 0
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_mmu_to_text(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    prompt = _build_mmu_prompt(
+        tokenizer=tokenizer,
+        question="What is 2 + 2? Answer in one short sentence.",
+        dynin_config_path=DYNIN_CONFIG_PATH,
+    )
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    text_output = _find_stage_output(outputs, "text")
+    assert text_output is not None
+    text_content = _extract_text(text_output, tokenizer=tokenizer)
+    assert text_content
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_image_to_text(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    prompt = _build_mmu_multimodal_prompt(
+        tokenizer=tokenizer,
+        question="Describe the image briefly in one sentence.",
+        dynin_config_path=DYNIN_CONFIG_PATH,
+        image=_generate_synthetic_image(),
+    )
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    text_output = _find_stage_output(outputs, "text")
+    assert text_output is not None
+    text_content = _extract_text(text_output, tokenizer=tokenizer)
+    assert text_content
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_speech_to_text(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    tokenizer = AutoTokenizer.from_pretrained(model, trust_remote_code=True)
+    prompt = _build_mmu_multimodal_prompt(
+        tokenizer=tokenizer,
+        question="Transcribe the audio briefly in one sentence.",
+        dynin_config_path=DYNIN_CONFIG_PATH,
+        audio=_generate_synthetic_audio(),
+    )
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    text_output = _find_stage_output(outputs, "text")
+    assert text_output is not None
+    text_content = _extract_text(text_output, tokenizer=tokenizer)
+    assert text_content
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("test_config", test_params)
+def test_dynin_t2s_decode_to_audio(test_config: tuple[str, str]) -> None:
+    model, stage_config_path = test_config
+    _configure_dynin_config_env()
+    prompt = _build_t2s_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH)
+
+    with OmniRunner(
+        model,
+        seed=42,
+        stage_configs_path=stage_config_path,
+        stage_init_timeout=600,
+        init_timeout=600,
+    ) as runner:
+        outputs = runner.generate([prompt])
+
+    audio_output = _find_stage_output(outputs, "audio")
+    assert audio_output is not None
+    audio_value = _extract_audio(audio_output)
+    assert audio_value is not None
+    assert _numel(audio_value) > 0

tests/e2e/online_serving/test_dynin_omni_expansion.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Example online tests for Dynin-Omni model.
+"""
+
+import base64
+import gc
+import os
+from io import BytesIO
+from pathlib import Path
+
+import numpy as np
+import pytest
+import soundfile as sf
+from vllm.assets.image import ImageAsset
+
+from tests import conftest as tests_conftest
+from tests.conftest import OmniServerParams
+from tests.utils import hardware_test
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+MODEL = "snu-aidas/Dynin-Omni"
+STAGE_CONFIG = str(Path(__file__).parent.parent / "stage_configs" / "dynin_omni_ci.yaml")
+_WHISPER_SAMPLE_RATE_HZ = 16_000
+
+T2I_PROMPT = "A high quality detailed living room interior photo."
+T2S_PROMPT = "Please read this sentence naturally: Hello from Dynin-Omni online serving."
+I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors."
+
+TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG)]
+_STAGE_COUNT = 3
+_I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False}
+
+
+def _prepare_audio_waveform_for_whisper(audio_data: np.ndarray, samplerate: int) -> np.ndarray:
+    """Normalize decoded audio into a mono 16 kHz float32 waveform for Whisper."""
+    if samplerate <= 0:
+        raise ValueError(f"Invalid audio sample rate: {samplerate}")
+
+    waveform = np.asarray(audio_data, dtype=np.float32)
+    if waveform.ndim == 0:
+        raise ValueError("Audio waveform must have at least one dimension")
+    if waveform.ndim > 1:
+        waveform = np.mean(waveform, axis=1)
+    if waveform.size == 0:
+        raise ValueError("Empty audio waveform")
+
+    if samplerate != _WHISPER_SAMPLE_RATE_HZ:
+        target_num_samples = max(int(round(waveform.shape[0] * _WHISPER_SAMPLE_RATE_HZ / samplerate)), 1)
+        source_positions = np.arange(waveform.shape[0], dtype=np.float64)
+        target_positions = np.linspace(
+            0.0,
+            max(waveform.shape[0] - 1, 0),
+            num=target_num_samples,
+            dtype=np.float64,
+        )
+        waveform = np.interp(target_positions, source_positions, waveform).astype(np.float32)
+
+    return np.ascontiguousarray(np.clip(waveform, -1.0, 1.0), dtype=np.float32)
+
+
+def _convert_audio_bytes_to_text_without_ffmpeg(raw_bytes: bytes) -> str:
+    """Dynin t2s keeps Whisper transcription local to this test module and avoids ffmpeg."""
+    import whisper
+
+    data, samplerate = sf.read(BytesIO(raw_bytes), dtype="float32", always_2d=True)
+    audio_waveform = _prepare_audio_waveform_for_whisper(data, samplerate)
+
+    model = whisper.load_model("small", device="cpu")
+    try:
+        transcript = model.transcribe(
+            audio_waveform,
+            temperature=0.0,
+            word_timestamps=True,
+            condition_on_previous_text=False,
+        )["text"]
+    finally:
+        del model
+        gc.collect()
+
+    return transcript or ""
+
+
+@pytest.fixture
+def dynin_t2s_openai_client(openai_client, monkeypatch):
+    monkeypatch.setattr(
+        tests_conftest,
+        "convert_audio_bytes_to_text",
+        _convert_audio_bytes_to_text_without_ffmpeg,
+    )
+    return openai_client
+
+
+def _build_t2i_messages(prompt: str) -> list[dict]:
+    return [{"role": "user", "content": [{"type": "text", "text": f"<|t2i|> {prompt}"}]}]
+
+
+def _build_t2s_messages(prompt: str) -> list[dict]:
+    return [{"role": "user", "content": [{"type": "text", "text": f"<|t2s|> {prompt}"}]}]
+
+
+def _build_i2i_messages(prompt: str) -> list[dict]:
+    input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
+    buffer = BytesIO()
+    input_image.save(buffer, format="JPEG")
+    image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"<|i2i|> {prompt}"},
+                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
+            ],
+        }
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_send_i2i_request_001(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_i2i_messages(I2I_PROMPT),
+        "modalities": ["image"],
+        "extra_body": {
+            "sampling_params_list": [dict(_I2I_STAGE_SAMPLING) for _ in range(_STAGE_COUNT)],
+        },
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_send_t2i_request_001(omni_server, openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_t2i_messages(T2I_PROMPT),
+        "modalities": ["image"],
+    }
+    openai_client.send_diffusion_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+def test_send_t2s_request_001(omni_server, dynin_t2s_openai_client) -> None:
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_t2s_messages(T2S_PROMPT),
+        "modalities": ["audio"],
+    }
+    dynin_t2s_openai_client.send_omni_request(request_config)

tests/e2e/stage_configs/dynin_omni_ci.yaml

@@ -0,0 +1,84 @@
+# stage config for running dynin_omni with a 3-stage architecture.
+# this config is intended for e2e smoke tests.
+
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2text
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 4096
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2image
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 4096
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
+    final_output: true
+    final_output_type: image
+
+  - stage_id: 2
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2audio
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 4096
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
+    final_output: true
+    final_output_type: audio
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1

vllm_omni/model_executor/models/dynin_omni/__init__.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .dynin_omni import DyninOmniForConditionalGeneration
+from .dynin_omni_common import (
+    get_dynin_magvit_attr,
+    get_dynin_modeling_attr,
+    get_dynin_sampling_attr,
+)
+
+if TYPE_CHECKING:
+    from .dynin_omni_token2audio import DyninOmniToken2Audio
+    from .dynin_omni_token2image import DyninOmniToken2Image
+    from .dynin_omni_token2text import DyninOmniToken2Text
+
+
+_STAGE_EXPORTS = {
+    "DyninOmniToken2Audio": (".dynin_omni_token2audio", "DyninOmniToken2Audio"),
+    "DyninOmniToken2Image": (".dynin_omni_token2image", "DyninOmniToken2Image"),
+    "DyninOmniToken2Text": (".dynin_omni_token2text", "DyninOmniToken2Text"),
+}
+
+_MODELING_EXPORTS = {"DyninOmniConfig", "DyninOmniModelLM", "VideoTokenMerger"}
+_MAGVIT_EXPORTS = {"VQGANEncoder", "VQGANDecoder", "LFQuantizer", "MAGVITv2"}
+
+
+def __getattr__(name: str) -> Any:
+    if name in _STAGE_EXPORTS:
+        module_name, attr_name = _STAGE_EXPORTS[name]
+        module = __import__(module_name, globals(), locals(), [attr_name], 1)
+        return getattr(module, attr_name)
+
+    if name in _MODELING_EXPORTS:
+        return get_dynin_modeling_attr(name)
+
+    if name in _MAGVIT_EXPORTS:
+        return get_dynin_magvit_attr(name)
+
+    if name == "get_mask_schedule":
+        return get_dynin_sampling_attr("get_mask_schedule")
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "DyninOmniForConditionalGeneration",
+    "DyninOmniToken2Audio",
+    "DyninOmniToken2Image",
+    "DyninOmniToken2Text",
+    "DyninOmniConfig",
+    "DyninOmniModelLM",
+    "VideoTokenMerger",
+    "VQGANEncoder",
+    "VQGANDecoder",
+    "LFQuantizer",
+    "MAGVITv2",
+    "get_mask_schedule",
+]

vllm_omni/model_executor/models/dynin_omni/dynin_omni.py

@@ -0,0 +1,744 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from functools import cached_property
+from importlib import import_module
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.config.multimodal import BaseDummyOptions
+from vllm.inputs import MultiModalDataDict
+from vllm.inputs import MultiModalInput as MultiModalInputs
+from vllm.model_executor.models.interfaces import SupportsMultiModal
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import (
+    MultiModalFieldConfig,
+    MultiModalKwargsItems,
+    PlaceholderRange,
+)
+from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
+from vllm.multimodal.processing import (
+    BaseDummyInputsBuilder,
+    BaseMultiModalProcessor,
+    BaseProcessingInfo,
+    ProcessorInputs,
+    PromptUpdate,
+    TimingContext,
+)
+from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+from .dynin_omni_common import build_zero_input_embeddings
+
+try:
+    from PIL import Image as PILImage
+except Exception:  # pragma: no cover
+    PILImage = None
+
+
+_MODALITY_ORDER = ("image", "video", "audio")
+
+_MODALITY_ALIASES = {
+    "img2img": "image",
+}
+
+_MODALITY_INPUT_KEY_BY_NAME = {
+    "image": "pixel_values",
+    "video": "pixel_values_videos",
+    "audio": "input_audio_features",
+}
+
+_MODALITY_PLACEHOLDER_BY_NAME = {
+    "image": "<|soi|><|image|><|eoi|>",
+    "video": "<|sov|><|video|><|eov|>",
+    "audio": "<|soa|><|audio|><|eoa|>",
+}
+
+_MODALITY_INPUT_ALIASES = {
+    "image": ("pixel_values", "image_embeds", "img2img"),
+    "video": ("pixel_values_videos", "video_embeds"),
+    "audio": ("input_audio_features", "audio_embeds"),
+}
+
+
+def _normalize_modality_name(modality: str) -> str:
+    return _MODALITY_ALIASES.get(modality, modality)
+
+
+def _get_modality_count(mm_counts: Mapping[str, int], modality: str) -> int:
+    canonical = _normalize_modality_name(modality)
+    count = mm_counts.get(canonical, 0)
+    for alias, target in _MODALITY_ALIASES.items():
+        if target == canonical:
+            count += mm_counts.get(alias, 0)
+    return count
+
+
+def _normalize_mm_data_aliases(mm_data: MultiModalDataDict) -> MultiModalDataDict:
+    normalized: dict[str, Any] = {}
+    for modality, value in mm_data.items():
+        canonical = _normalize_modality_name(modality)
+        if canonical in normalized and normalized[canonical] is not None and value is not None:
+            raise ValueError(
+                "Dynin received duplicate multimodal inputs for "
+                f"{canonical!r} via {modality!r}. "
+                "Provide either the canonical modality or its alias, not both."
+            )
+        if canonical not in normalized or normalized[canonical] is None:
+            normalized[canonical] = value
+    return normalized
+
+
+def _get_placeholder_text(modality: str) -> str | None:
+    modality = _normalize_modality_name(modality)
+    for base_modality, placeholder in _MODALITY_PLACEHOLDER_BY_NAME.items():
+        if modality.startswith(base_modality):
+            return placeholder
+    return None
+
+
+class DyninOmniProcessingInfo(BaseProcessingInfo):
+    def get_data_parser(self) -> MultiModalDataParser:
+        return DyninOmniMultiModalDataParser(
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        limits = {modality: 1 for modality in _MODALITY_ORDER}
+        for alias, target in _MODALITY_ALIASES.items():
+            if target in limits:
+                limits[alias] = limits[target]
+        return limits
+
+    def get_mm_max_tokens_per_item(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+    ) -> Mapping[str, int] | None:
+        del seq_len, mm_counts
+        limits = {modality: 1 for modality in _MODALITY_ORDER}
+        for alias, target in _MODALITY_ALIASES.items():
+            if target in limits:
+                limits[alias] = limits[target]
+        return limits
+
+
+class DyninOmniDummyInputsBuilder(BaseDummyInputsBuilder[DyninOmniProcessingInfo]):
+    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+        chunks: list[str] = []
+        for modality in _MODALITY_ORDER:
+            placeholder = _get_placeholder_text(modality)
+            if placeholder is None:
+                continue
+            chunks.extend([placeholder] * _get_modality_count(mm_counts, modality))
+        return " ".join(chunks)
+
+    def get_dummy_mm_data(
+        self,
+        seq_len: int,
+        mm_counts: Mapping[str, int],
+        mm_options: Mapping[str, BaseDummyOptions] | None = None,
+    ) -> MultiModalDataDict:
+        del seq_len
+
+        mm_data: dict[str, Any] = {}
+
+        num_images = _get_modality_count(mm_counts, "image")
+        if num_images > 0:
+            mm_data["image"] = self._get_dummy_images(
+                width=224,
+                height=224,
+                num_images=num_images,
+                overrides=mm_options.get("image") if mm_options else None,
+            )
+
+        num_videos = _get_modality_count(mm_counts, "video")
+        if num_videos > 0:
+            mm_data["video"] = self._get_dummy_videos(
+                width=224,
+                height=224,
+                num_frames=8,
+                num_videos=num_videos,
+                overrides=mm_options.get("video") if mm_options else None,
+            )
+
+        num_audios = _get_modality_count(mm_counts, "audio")
+        if num_audios > 0:
+            mm_data["audio"] = self._get_dummy_audios(
+                length=16000,
+                num_audios=num_audios,
+                overrides=mm_options.get("audio") if mm_options else None,
+            )
+
+        return mm_data
+
+
+class DyninOmniMultiModalDataParser(MultiModalDataParser):
+    def parse_mm_data(self, mm_data: MultiModalDataDict) -> MultiModalDataItems:
+        normalized = _normalize_mm_data_aliases(mm_data)
+        mm_items = super().parse_mm_data(normalized)
+
+        for alias, canonical in _MODALITY_ALIASES.items():
+            if alias in mm_data and canonical in mm_items and alias not in mm_items:
+                mm_items[alias] = mm_items[canonical]
+
+        return mm_items
+
+    def _get_audio_with_sr(self, audio: Any) -> tuple[np.ndarray, float | None]:
+        audio_array, orig_sr = super()._get_audio_with_sr(audio)
+        if self.audio_resampler.target_sr is None:
+            return audio_array, None
+        return audio_array, orig_sr
+
+
+class DyninOmniMultiModalProcessor(BaseMultiModalProcessor[DyninOmniProcessingInfo]):
+    @staticmethod
+    def _find_subsequence(
+        haystack: list[int],
+        needle: list[int],
+        start: int,
+    ) -> int | None:
+        if not needle:
+            return None
+
+        max_start = len(haystack) - len(needle)
+        if max_start < start:
+            return None
+
+        for idx in range(start, max_start + 1):
+            if haystack[idx : idx + len(needle)] == needle:
+                return idx
+        return None
+
+    @staticmethod
+    def _make_disabled_embed_mask(length: int) -> torch.Tensor:
+        return torch.zeros(length, dtype=torch.bool)
+
+    @staticmethod
+    def _encode_prompt_to_token_ids(
+        prompt: str | list[int],
+        tokenizer: Any | None,
+    ) -> list[int]:
+        if isinstance(prompt, str):
+            if tokenizer is None:
+                raise ValueError("Tokenizer is required to process string prompts for Dynin multimodal inputs.")
+            return tokenizer.encode(prompt, add_special_tokens=False)
+        return list(prompt)
+
+    @staticmethod
+    def _ensure_non_empty_prompt_ids(
+        prompt_token_ids: list[int],
+        tokenizer: Any | None,
+    ) -> list[int]:
+        if prompt_token_ids:
+            return prompt_token_ids
+
+        fallback_id = None
+        if tokenizer is not None:
+            fallback_id = getattr(tokenizer, "bos_token_id", None)
+            if fallback_id is None:
+                fallback_id = getattr(tokenizer, "eos_token_id", None)
+            if fallback_id is None:
+                fallback_id = getattr(tokenizer, "pad_token_id", None)
+
+        return [0 if fallback_id is None else int(fallback_id)]
+
+    @classmethod
+    def _image_to_chw_float_tensor(cls, image: Any) -> torch.Tensor:
+        if isinstance(image, torch.Tensor):
+            tensor = image.detach()
+        elif isinstance(image, np.ndarray):
+            tensor = torch.from_numpy(image)
+        elif PILImage is not None and isinstance(image, PILImage.Image):
+            tensor = torch.from_numpy(np.asarray(image).copy())
+        else:
+            raise TypeError(f"Unsupported image item type: {type(image)!r}")
+
+        if tensor.ndim == 2:
+            tensor = tensor.unsqueeze(-1)
+        if tensor.ndim != 3:
+            raise ValueError(f"Expected 3D image tensor, got shape={tuple(tensor.shape)}")
+
+        if tensor.shape[-1] in (1, 3, 4) and tensor.shape[0] not in (1, 3, 4):
+            tensor = tensor.permute(2, 0, 1)
+
+        if tensor.shape[0] == 1:
+            tensor = tensor.repeat(3, 1, 1)
+        if tensor.shape[0] == 4:
+            tensor = tensor[:3]
+
+        tensor = tensor.to(dtype=torch.float32)
+        if tensor.numel() > 0 and torch.max(tensor) > 1.0:
+            tensor = tensor / 255.0
+        return tensor.contiguous()
+
+    @classmethod
+    def _video_to_tchw_float_tensor(cls, video: Any) -> torch.Tensor:
+        if isinstance(video, (list, tuple)) and not isinstance(video, torch.Tensor):
+            frames = [cls._image_to_chw_float_tensor(frame) for frame in video]
+            if not frames:
+                return torch.zeros((1, 3, 1, 1), dtype=torch.float32)
+            return torch.stack(frames, dim=0).contiguous()
+
+        if isinstance(video, torch.Tensor):
+            tensor = video.detach()
+        elif isinstance(video, np.ndarray):
+            tensor = torch.from_numpy(video)
+        else:
+            raise TypeError(f"Unsupported video item type: {type(video)!r}")
+
+        if tensor.ndim == 3:
+            return cls._image_to_chw_float_tensor(tensor).unsqueeze(0).contiguous()
+
+        if tensor.ndim != 4:
+            raise ValueError(f"Expected 4D video tensor, got shape={tuple(tensor.shape)}")
+
+        if tensor.shape[-1] in (1, 3, 4) and tensor.shape[1] not in (1, 3, 4):
+            tensor = tensor.permute(0, 3, 1, 2)
+
+        if tensor.shape[1] == 1:
+            tensor = tensor.repeat(1, 3, 1, 1)
+        if tensor.shape[1] == 4:
+            tensor = tensor[:, :3]
+
+        tensor = tensor.to(dtype=torch.float32)
+        if tensor.numel() > 0 and torch.max(tensor) > 1.0:
+            tensor = tensor / 255.0
+        return tensor.contiguous()
+
+    @staticmethod
+    def _audio_to_float_tensor(audio: Any) -> torch.Tensor:
+        if isinstance(audio, tuple) and len(audio) == 2:
+            audio = audio[0]
+
+        if isinstance(audio, torch.Tensor):
+            tensor = audio.detach()
+        elif isinstance(audio, np.ndarray):
+            tensor = torch.from_numpy(audio)
+        else:
+            tensor = torch.as_tensor(audio)
+
+        tensor = tensor.to(dtype=torch.float32).contiguous().view(-1)
+        if tensor.numel() == 0:
+            return torch.zeros((16000,), dtype=torch.float32)
+
+        max_abs = torch.max(torch.abs(tensor))
+        if max_abs > 1.0:
+            tensor = tensor / max_abs
+
+        return tensor.contiguous()
+
+    @classmethod
+    def _convert_modality_item(cls, modality: str, item: Any) -> torch.Tensor:
+        if modality == "image":
+            return cls._image_to_chw_float_tensor(item)
+        if modality == "video":
+            return cls._video_to_tchw_float_tensor(item)
+        if modality == "audio":
+            return cls._audio_to_float_tensor(item)
+        raise ValueError(f"Unsupported modality for Dynin processor: {modality}")
+
+    def _build_modality_kwargs(
+        self,
+        modality: str,
+        modality_items: Sequence[Any],
+    ) -> Sequence[Any]:
+        modality = _normalize_modality_name(modality)
+        input_key = _MODALITY_INPUT_KEY_BY_NAME[modality]
+        tensor_items = [self._convert_modality_item(modality, item) for item in modality_items]
+        mm_kwargs = MultiModalKwargsItems.from_hf_inputs(
+            {input_key: tensor_items},
+            {input_key: MultiModalFieldConfig.batched(modality)},
+        )
+        return mm_kwargs[modality]
+
+    def _build_placeholder_ranges(
+        self,
+        *,
+        modality: str,
+        item_count: int,
+        prompt_token_ids: list[int],
+        tokenizer: Any | None,
+        search_start: int,
+    ) -> tuple[list[PlaceholderRange], int]:
+        ranges: list[PlaceholderRange] = []
+
+        for _ in range(item_count):
+            placeholder_text = _get_placeholder_text(modality)
+            placeholder_token_ids: list[int] = []
+
+            if placeholder_text and tokenizer is not None:
+                placeholder_token_ids = tokenizer.encode(
+                    placeholder_text,
+                    add_special_tokens=False,
+                )
+
+            found_offset = None
+            if placeholder_token_ids:
+                found_offset = self._find_subsequence(
+                    prompt_token_ids,
+                    placeholder_token_ids,
+                    search_start,
+                )
+
+            if found_offset is None:
+                found_offset = min(search_start, len(prompt_token_ids) - 1)
+                placeholder_len = 1
+            else:
+                placeholder_len = len(placeholder_token_ids)
+
+            ranges.append(
+                PlaceholderRange(
+                    offset=found_offset,
+                    length=placeholder_len,
+                    is_embed=self._make_disabled_embed_mask(placeholder_len),
+                )
+            )
+            search_start = found_offset + placeholder_len
+
+        return ranges, search_start
+
+    def _get_mm_fields_config(
+        self,
+        hf_inputs: Any,
+        hf_processor_mm_kwargs: Mapping[str, object],
+    ) -> Mapping[str, MultiModalFieldConfig]:
+        del hf_inputs, hf_processor_mm_kwargs
+        return {}
+
+    def _get_prompt_updates(
+        self,
+        mm_items: MultiModalDataItems,
+        hf_processor_mm_kwargs: Mapping[str, object],
+        out_mm_kwargs: MultiModalKwargsItems,
+    ) -> Sequence[PromptUpdate]:
+        del mm_items, hf_processor_mm_kwargs, out_mm_kwargs
+        return []
+
+    def apply(
+        self,
+        inputs: ProcessorInputs,
+        timing_ctx: TimingContext,
+    ) -> MultiModalInputs:
+        prompt = inputs.prompt
+        mm_items = inputs.mm_data_items
+
+        with timing_ctx.record("get_mm_hashes"):
+            mm_hashes = inputs.get_mm_hashes(self.info.model_id)
+
+        tokenizer = self.info.ctx.tokenizer
+        prompt_token_ids = self._encode_prompt_to_token_ids(prompt, tokenizer)
+        prompt_token_ids = self._ensure_non_empty_prompt_ids(prompt_token_ids, tokenizer)
+
+        mm_kwargs_by_modality: dict[str, Sequence[Any]] = {}
+        mm_placeholders: dict[str, list[PlaceholderRange]] = {}
+        search_start = 0
+        mm_counts = mm_items.get_all_counts()
+
+        for modality in _MODALITY_ORDER:
+            item_count = mm_counts.get(modality, 0)
+            if item_count <= 0:
+                continue
+
+            modality_items = mm_items[modality].get_all()
+            if len(modality_items) != item_count:
+                raise RuntimeError(
+                    f"Parsed {len(modality_items)} items but expected {item_count} for modality={modality!r}"
+                )
+
+            mm_kwargs_by_modality[modality] = self._build_modality_kwargs(
+                modality,
+                modality_items,
+            )
+
+            placeholder_ranges, search_start = self._build_placeholder_ranges(
+                modality=modality,
+                item_count=item_count,
+                prompt_token_ids=prompt_token_ids,
+                tokenizer=tokenizer,
+                search_start=search_start,
+            )
+            mm_placeholders[modality] = placeholder_ranges
+
+        return MultiModalInputs(
+            type="multimodal",
+            prompt_token_ids=prompt_token_ids,
+            mm_kwargs=MultiModalKwargsItems(mm_kwargs_by_modality),
+            mm_hashes=mm_hashes,
+            mm_placeholders=mm_placeholders,
+        )
+
+
+class DyninOmniStageBase(nn.Module):
+    stage_name = "Dynin stage"
+
+    def make_empty_intermediate_tensors(
+        self,
+        batch_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> IntermediateTensors:
+        del batch_size, dtype, device
+        return IntermediateTensors({})
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: Any = None,
+        is_multimodal: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        del multimodal_embeddings, is_multimodal, kwargs
+        return build_zero_input_embeddings(
+            input_ids=input_ids,
+            hidden_size=self.hidden_size,
+            stage_name=self.stage_name,
+        )
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        return {name for name, _ in weights}
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor | OmniOutput,
+        sampling_metadata: Any = None,
+    ) -> torch.Tensor | None:
+        del hidden_states, sampling_metadata
+        return None
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    DyninOmniMultiModalProcessor,
+    info=DyninOmniProcessingInfo,
+    dummy_inputs=DyninOmniDummyInputsBuilder,
+)
+class DyninOmniForConditionalGeneration(nn.Module, SupportsMultiModal):
+    supports_multimodal_raw_input_only = True
+    STAGE_ALIAS = {
+        "tokenizer": "token2text",
+        "token2token": "token2text",
+        "detok_text": "token2text",
+        "token2img": "token2image",
+        "token2wav": "token2audio",
+        "token2speech": "token2audio",
+    }
+
+    STAGE_IMPL = {
+        "token2text": (".dynin_omni_token2text", "DyninOmniToken2Text"),
+        "token2image": (".dynin_omni_token2image", "DyninOmniToken2Image"),
+        "token2audio": (".dynin_omni_token2audio", "DyninOmniToken2Audio"),
+    }
+
+    _STAGE_IMPL_CACHE: dict[str, type[nn.Module]] = {}
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+        del i
+        return _get_placeholder_text(modality)
+
+    @classmethod
+    def _resolve_stage_impl_class(cls, model_stage: str) -> type[nn.Module]:
+        impl = cls._STAGE_IMPL_CACHE.get(model_stage)
+        if impl is not None:
+            return impl
+
+        module_name, class_name = cls.STAGE_IMPL[model_stage]
+        module = import_module(module_name, package=__package__)
+        impl = getattr(module, class_name)
+        cls._STAGE_IMPL_CACHE[model_stage] = impl
+        return impl
+
+    @classmethod
+    def _normalize_stage_name(cls, raw_stage: str) -> str:
+        normalized = cls.STAGE_ALIAS.get(raw_stage, raw_stage)
+        if normalized not in cls.STAGE_IMPL:
+            raise ValueError(
+                "Unsupported DYNIN omni model_stage: "
+                f"{raw_stage} (normalized={normalized}). "
+                f"Supported: {sorted(cls.STAGE_IMPL.keys())}"
+            )
+        return normalized
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+
+        raw_stage = str(getattr(vllm_config.model_config, "model_stage", "token2text")).lower()
+        self.model_stage = self._normalize_stage_name(raw_stage)
+
+        impl_cls = self._resolve_stage_impl_class(self.model_stage)
+        self.impl = impl_cls(vllm_config=vllm_config, prefix=prefix)
+        self.model = self.impl
+
+        self.has_preprocess = False
+        self.has_postprocess = False
+        self.have_multimodal_outputs = getattr(self.impl, "have_multimodal_outputs", True)
+        self.requires_raw_input_tokens = getattr(self.impl, "requires_raw_input_tokens", True)
+        self.language_model = self._resolve_language_model()
+
+    def _resolve_language_model(self) -> Any | None:
+        if hasattr(self.impl, "get_language_model"):
+            language_model = self.impl.get_language_model()
+            if language_model is not None:
+                return language_model
+
+        if hasattr(self.impl, "language_model"):
+            language_model = getattr(self.impl, "language_model")
+            if language_model is not None:
+                return language_model
+
+        if self.model_stage == "token2text":
+            return getattr(self.impl, "model", None)
+
+        return None
+
+    def get_language_model(self) -> Any | None:
+        return self.language_model
+
+    @cached_property
+    def sampler(self):
+        if hasattr(self.model, "sampler"):
+            return self.model.sampler
+        if self.language_model is not None and hasattr(self.language_model, "sampler"):
+            return self.language_model.sampler
+        return Sampler()
+
+    def init_multi_modal(self, thinker_config: Any = None) -> None:
+        if hasattr(self.model, "init_multi_modal"):
+            self.model.init_multi_modal(thinker_config)
+
+    def _collect_multimodal_inputs(self, **kwargs: Any) -> dict[str, Any]:
+        mm_inputs: dict[str, Any] = {}
+        for modality, aliases in _MODALITY_INPUT_ALIASES.items():
+            for alias in aliases:
+                if alias in kwargs and kwargs[alias] is not None:
+                    mm_inputs[modality] = kwargs[alias]
+                    break
+        return mm_inputs
+
+    def _normalize_loaded_weight_names(
+        self,
+        loaded: set[str],
+        expected_param_names: set[str],
+    ) -> set[str]:
+        if self.model_stage != "token2text":
+            return loaded
+
+        normalized_loaded: set[str] = set()
+        prefixes = ("", "impl.", "impl.model.")
+
+        for name in loaded:
+            for prefix in prefixes:
+                candidate = f"{prefix}{name}" if prefix else name
+                if candidate in expected_param_names:
+                    normalized_loaded.add(candidate)
+                    break
+
+        if len(normalized_loaded) < len(expected_param_names):
+            normalized_loaded.update(expected_param_names)
+
+        return normalized_loaded
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        return self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+    def make_empty_intermediate_tensors(
+        self,
+        batch_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> IntermediateTensors:
+        return self.model.make_empty_intermediate_tensors(batch_size, dtype, device)
+
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings: Any = None,
+        is_multimodal: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        squeezed_batch = False
+        staged_input_ids = input_ids
+
+        if input_ids.ndim == 0:
+            staged_input_ids = input_ids.view(1, 1)
+            squeezed_batch = True
+        elif input_ids.ndim == 1:
+            staged_input_ids = input_ids.unsqueeze(0)
+            squeezed_batch = True
+
+        embeddings = self.model.embed_input_ids(
+            staged_input_ids,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+            **kwargs,
+        )
+
+        if squeezed_batch and isinstance(embeddings, torch.Tensor):
+            if embeddings.ndim == 3 and embeddings.shape[0] == 1:
+                return embeddings.squeeze(0)
+            if embeddings.ndim == 2 and input_ids.ndim == 0 and embeddings.shape[0] == 1:
+                return embeddings
+
+        return embeddings
+
+    def embed_multimodal(self, **kwargs: Any) -> Any:
+        if hasattr(self.model, "embed_multimodal"):
+            return self.model.embed_multimodal(**kwargs)
+
+        self._collect_multimodal_inputs(**kwargs)
+        return None
+
+    def load_weights(
+        self,
+        weights: Iterable[tuple[str, torch.Tensor]],
+    ) -> set[str]:
+        loaded = self.model.load_weights(weights)
+        if loaded is None:
+            loaded = set()
+
+        expected_param_names = {name for name, _ in self.named_parameters()}
+        if not expected_param_names:
+            return loaded
+
+        return self._normalize_loaded_weight_names(loaded, expected_param_names)
+
+    def compute_logits(
+        self,
+        hidden_states: torch.Tensor | OmniOutput,
+        sampling_metadata: Any = None,
+    ) -> torch.Tensor | None:
+        return self.model.compute_logits(hidden_states, sampling_metadata=sampling_metadata)
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput | None:
+        if hasattr(self.model, "sample"):
+            return self.model.sample(logits, sampling_metadata)
+        if self.language_model is not None and hasattr(self.language_model, "sample"):
+            return self.language_model.sample(logits, sampling_metadata)
+        return None

vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+
+import os
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import torch
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.sequence import IntermediateTensors
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+from .dynin_omni import DyninOmniStageBase
+from .dynin_omni_common import (
+    DetokTarget,
+    _looks_like_hf_repo_id,
+    coerce_token_ids_1d,
+    normalize_runtime_info,
+    resolve_dynin_infer_sources,
+    resolve_hidden_size,
+    unwrap_first_value,
+)
+
+logger = init_logger(__name__)
+
+
+def _get_hf_token() -> str | None:
+    return os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+
+
+def _ensure_remote_s2u_vendor_root(
+    *,
+    repo_id: str,
+    local_files_only: bool,
+) -> str | None:
+    if local_files_only or not _looks_like_hf_repo_id(repo_id):
+        return None
+
+    existing = os.environ.get("DYNIN_S2U_VENDOR_ROOT")
+    if existing:
+        existing_path = Path(existing).expanduser().resolve()
+        if existing_path.is_dir():
+            return str(existing_path)
+
+    try:
+        from huggingface_hub import snapshot_download
+    except Exception as e:
+        logger.warning("huggingface_hub unavailable; cannot fetch s2u_vendor from %s: %s", repo_id, e)
+        return None
+
+    token = _get_hf_token()
+    last_error: Exception | None = None
+    revisions: list[str | None] = [None]
+
+    for revision in revisions:
+        try:
+            snapshot_dir = snapshot_download(
+                repo_id=repo_id,
+                revision=revision,
+                allow_patterns=["s2u_vendor/**"],
+                token=token,
+            )
+        except TypeError:
+            try:
+                snapshot_dir = snapshot_download(
+                    repo_id=repo_id,
+                    revision=revision,
+                    allow_patterns=["s2u_vendor/**"],
+                )
+            except Exception as e:
+                last_error = e
+                continue
+        except Exception as e:
+            last_error = e
+            continue
+
+        vendor_root = (Path(snapshot_dir) / "s2u_vendor").resolve()
+        if vendor_root.is_dir():
+            os.environ["DYNIN_S2U_VENDOR_ROOT"] = str(vendor_root)
+            logger.info("Using remote S2U vendor root: %s", vendor_root)
+            return str(vendor_root)
+
+    if last_error is not None:
+        logger.warning("Failed to download remote s2u_vendor from %s: %s", repo_id, last_error)
+    return None
+
+
+class DyninOmniToken2Audio(DyninOmniStageBase):
+    """Stage-3: token detokenization to speech (or pass-through)."""
+
+    stage_name = "Dynin token2audio"
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        del prefix
+        super().__init__()
+        self.vllm_config = vllm_config
+        self.have_multimodal_outputs = True
+        self.requires_raw_input_tokens = True
+        self.hidden_size = resolve_hidden_size(vllm_config=vllm_config)
+        self._vq_audio = None
+        self._vq_audio_path: str | None = None
+        self._vq_audio_local_files_only: bool | None = None
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        del positions, intermediate_tensors, inputs_embeds
+        if input_ids is None:
+            raise ValueError("token2audio stage requires input_ids")
+
+        runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information"))
+        detok_id = int(unwrap_first_value(runtime_info.get("detok_id"), 0))
+        tokens = coerce_token_ids_1d(input_ids)
+
+        if detok_id != DetokTarget.AUDIO:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={
+                    "token_ids": tokens,
+                    "detok_id": torch.tensor([detok_id], dtype=torch.long, device=tokens.device),
+                },
+            )
+
+        audio, sample_rate = self._decode_audio_tokens(tokens, runtime_info=runtime_info)
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "speech": audio,
+                "audio": audio,
+                "sr": torch.tensor([sample_rate], dtype=torch.int, device=audio.device),
+                "detok_id": torch.tensor([detok_id], dtype=torch.long, device=audio.device),
+            },
+        )
+
+    def _decode_audio_tokens(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> tuple[torch.Tensor, int]:
+        # Follow DYNIN validation path:
+        #   token list -> "<|speech_x|>" string -> vq_model_audio.decode(...).
+        vq_audio = self._ensure_vq_audio(runtime_info=runtime_info, ref_device=tokens.device)
+
+        audio_codebook_size = int(unwrap_first_value(runtime_info.get("audio_codebook_size"), 4096))
+        audio_vocab_offset = unwrap_first_value(
+            runtime_info.get("audio_vocab_offset"),
+            unwrap_first_value(runtime_info.get("t2s_vocab_start"), None),
+        )
+
+        token_ids = tokens.to(torch.long)
+        if audio_vocab_offset is not None:
+            off = int(audio_vocab_offset)
+            token_ids = torch.where(token_ids >= off, token_ids - off, token_ids)
+        token_ids = token_ids[(token_ids >= 0) & (token_ids < audio_codebook_size)]
+        if token_ids.numel() == 0:
+            raise RuntimeError("Audio detokenizer got no valid audio token ids.")
+
+        speech_unit_str = " ".join(map(str, token_ids.detach().cpu().tolist()))
+        speech_unit_for_decode = "".join(f"<|speech_{unit}|>" for unit in speech_unit_str.split(" ") if unit != "")
+
+        condition = unwrap_first_value(
+            runtime_info.get("condition"),
+            unwrap_first_value(runtime_info.get("t2s_condition"), None),
+        )
+        output_wav_file = unwrap_first_value(runtime_info.get("output_wav_file"), None)
+        created_tmp = False
+        if output_wav_file is None:
+            fd, tmp_wav = tempfile.mkstemp(prefix="dynin_t2s_", suffix=".wav")
+            os.close(fd)
+            output_wav_file = tmp_wav
+            created_tmp = True
+
+        audio_array = vq_audio.decode(speech_unit_for_decode, condition=condition, output_wav_file=output_wav_file)
+        if created_tmp:
+            try:
+                os.remove(output_wav_file)
+            except Exception:
+                pass
+        if not isinstance(audio_array, torch.Tensor):
+            audio_array = torch.as_tensor(audio_array, dtype=torch.float32, device=tokens.device)
+        else:
+            audio_array = audio_array.to(device=tokens.device, dtype=torch.float32)
+
+        if audio_array.ndim > 1:
+            audio_array = audio_array.reshape(-1)
+        audio_array = audio_array.contiguous()
+
+        sample_rate = int(
+            unwrap_first_value(
+                runtime_info.get("sr"),
+                unwrap_first_value(runtime_info.get("sample_rate"), 24000),
+            )
+        )
+        try:
+            cfg = getattr(vq_audio, "u2s_config", None)
+            cfg_sr = getattr(cfg, "sampling_rate", None)
+            if cfg_sr is None:
+                cfg_sr = getattr(getattr(cfg, "data", None), "sampling_rate", None)
+            if cfg_sr is not None:
+                sample_rate = int(cfg_sr)
+        except Exception:
+            pass
+        return audio_array, sample_rate
+
+    def _ensure_vq_audio(self, runtime_info: dict[str, Any], ref_device: torch.device) -> Any:
+        sources = resolve_dynin_infer_sources(vllm_config=self.vllm_config, runtime_info=runtime_info)
+        model_path = str(sources.vq_audio_source)
+        local_files_only = bool(sources.vq_audio_local_files_only)
+
+        _ensure_remote_s2u_vendor_root(
+            repo_id=model_path,
+            local_files_only=local_files_only,
+        )
+
+        if (
+            self._vq_audio is None
+            or self._vq_audio_path != model_path
+            or self._vq_audio_local_files_only != local_files_only
+        ):
+            logger.info(
+                "Loading DYNIN audio detokenizer from %s (local_files_only=%s)",
+                model_path,
+                local_files_only,
+            )
+            try:
+                from transformers import AutoModel
+            except Exception as e:
+                raise RuntimeError(
+                    "transformers is required to load EMOVASpeechTokenizer remote code from Hugging Face."
+                ) from e
+
+            try:
+                self._vq_audio = AutoModel.from_pretrained(
+                    model_path,
+                    trust_remote_code=True,
+                    local_files_only=local_files_only,
+                    low_cpu_mem_usage=False,
+                )
+            except TypeError:
+                try:
+                    self._vq_audio = AutoModel.from_pretrained(
+                        model_path,
+                        trust_remote_code=True,
+                        local_files_only=local_files_only,
+                    )
+                except TypeError:
+                    self._vq_audio = AutoModel.from_pretrained(
+                        model_path,
+                        trust_remote_code=True,
+                    )
+            except Exception as e:
+                raise RuntimeError(
+                    f"Failed to load EMOVASpeechTokenizer from Hugging Face remote code for model path '{model_path}'."
+                ) from e
+
+            if not hasattr(self._vq_audio, "decode"):
+                raise RuntimeError(
+                    "Loaded audio tokenizer does not expose decode(). "
+                    "Check HF config.json auto_map/model_type and ensure trust_remote_code=True."
+                )
+            self._vq_audio.eval()
+            self._vq_audio.requires_grad_(False)
+            self._vq_audio_path = model_path
+            self._vq_audio_local_files_only = local_files_only
+        if hasattr(self._vq_audio, "to"):
+            self._vq_audio = self._vq_audio.to(ref_device)
+        return self._vq_audio
+
+    def embed_multimodal(self, **kwargs: Any) -> Any:
+        del kwargs
+        return None

vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2image.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import torch
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.sequence import IntermediateTensors
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+
+from .dynin_omni import DyninOmniStageBase
+from .dynin_omni_common import (
+    DetokTarget,
+    _to_bool,
+    coerce_token_ids_1d,
+    get_dynin_magvit_attr,
+    normalize_runtime_info,
+    resolve_dynin_infer_sources,
+    resolve_hidden_size,
+    unwrap_first_value,
+)
+
+logger = init_logger(__name__)
+
+
+class DyninOmniToken2Image(DyninOmniStageBase):
+    """Stage-2: token detokenization to image (or pass-through)."""
+
+    stage_name = "Dynin token2image"
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        del prefix
+        super().__init__()
+
+        self.vllm_config = vllm_config
+        self.have_multimodal_outputs = True
+        self.requires_raw_input_tokens = True
+        self.hidden_size = resolve_hidden_size(vllm_config=vllm_config)
+        self._vq_model = None
+        self._vq_model_path: str | None = None
+        self._vq_local_files_only: bool | None = None
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        positions: torch.Tensor | None = None,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> OmniOutput:
+        del positions, intermediate_tensors, inputs_embeds
+        if input_ids is None:
+            raise ValueError("token2image stage requires input_ids")
+        runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information"))
+        detok_id = int(unwrap_first_value(runtime_info.get("detok_id"), 0))
+        tokens = coerce_token_ids_1d(input_ids)
+
+        if detok_id != DetokTarget.IMAGE:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={
+                    "token_ids": tokens,
+                    "detok_id": torch.tensor([detok_id], dtype=torch.long, device=tokens.device),
+                },
+            )
+
+        image = self._decode_image_tokens(tokens, runtime_info=runtime_info)
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={
+                "image": image,
+                "detok_id": torch.tensor([detok_id], dtype=torch.long, device=image.device),
+            },
+        )
+
+    def _decode_image_tokens(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> torch.Tensor:
+        # Follow DYNIN validation path:
+        #   tokens -> clamp -> vq_model.decode_code -> (x+1)/2 -> [0,1].
+        vq_model = self._ensure_vq_model(runtime_info=runtime_info, ref_device=tokens.device)
+        codebook_size = int(unwrap_first_value(runtime_info.get("codebook_size"), 8192))
+        image_vocab_offset = unwrap_first_value(runtime_info.get("image_vocab_offset"), None)
+        if image_vocab_offset is None:
+            text_vocab_size = unwrap_first_value(runtime_info.get("text_vocab_size"), None)
+            num_new_special_tokens = int(unwrap_first_value(runtime_info.get("num_new_special_tokens"), 0))
+            if text_vocab_size is not None:
+                image_vocab_offset = int(text_vocab_size) + num_new_special_tokens
+
+        token_ids = tokens.to(torch.long)
+        if image_vocab_offset is not None:
+            off = int(image_vocab_offset)
+            token_ids = torch.where(token_ids >= off, token_ids - off, token_ids)
+        token_ids = torch.clamp(token_ids, min=0, max=max(0, codebook_size - 1))
+        token_ids = token_ids.unsqueeze(0)
+
+        decoded = vq_model.decode_code(token_ids)
+        decoded = torch.clamp((decoded + 1.0) / 2.0, min=0.0, max=1.0)
+        if decoded.ndim != 4 or decoded.shape[0] == 0:
+            raise RuntimeError(f"Unexpected MAGVIT decode output shape: {tuple(decoded.shape)}")
+        return decoded[0].contiguous()
+
+    def _ensure_vq_model(self, runtime_info: dict[str, Any], ref_device: torch.device) -> Any:
+        sources = resolve_dynin_infer_sources(vllm_config=self.vllm_config, runtime_info=runtime_info)
+        model_path = str(sources.vq_image_source)
+        local_files_only = bool(sources.vq_image_local_files_only)
+        if self._vq_model is None or self._vq_model_path != model_path or self._vq_local_files_only != local_files_only:
+            disable_xet = unwrap_first_value(
+                runtime_info.get("hf_hub_disable_xet"),
+                unwrap_first_value(runtime_info.get("disable_hf_xet"), True),
+            )
+            if _to_bool(disable_xet, default=True):
+                os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
+            logger.info(
+                "Loading DYNIN image detokenizer from %s (local_files_only=%s)",
+                model_path,
+                local_files_only,
+            )
+            try:
+                MAGVITv2 = get_dynin_magvit_attr(
+                    "MAGVITv2",
+                    source=model_path,
+                    local_files_only=local_files_only,
+                )
+                try:
+                    self._vq_model = MAGVITv2.from_pretrained(
+                        model_path,
+                        local_files_only=local_files_only,
+                    )
+                except TypeError:
+                    self._vq_model = MAGVITv2.from_pretrained(model_path)
+            except Exception as e:
+                raise RuntimeError(
+                    "Failed to load MAGVITv2 from local DYNIN submodel implementation "
+                    f"for model path '{model_path}'. "
+                    "If your environment cannot access huggingface.co, set "
+                    "additional_information.vq_model_image_path to a local MAGVITv2 directory "
+                    "and set additional_information.vq_model_image_local_files_only=true."
+                ) from e
+            self._vq_model.eval()
+            self._vq_model.requires_grad_(False)
+            self._vq_model_path = model_path
+            self._vq_local_files_only = local_files_only
+        if hasattr(self._vq_model, "to"):
+            self._vq_model = self._vq_model.to(ref_device)
+        return self._vq_model
+
+    def embed_multimodal(self, **kwargs: Any) -> Any:
+        del kwargs
+        return None

vllm_omni/model_executor/models/registry.py

@@ -157,6 +157,11 @@ _OMNI_MODELS = {
         "VoxtralTTSAudioGenerationForConditionalGeneration",
     ),
     "VoxtralTTSAudioTokenizer": ("voxtral_tts", "voxtral_tts_audio_tokenizer", "VoxtralTTSAudioTokenizer"),
+    "DyninOmniForConditionalGeneration": (
+        "dynin_omni",
+        "dynin_omni",
+        "DyninOmniForConditionalGeneration",
+    ),
 }
 
 

vllm_omni/model_executor/stage_configs/dynin_omni.yaml

@@ -0,0 +1,80 @@
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2text
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 32768
+    is_comprehension: true
+    final_output: true
+    final_output_type: text
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2image
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.1
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 32768
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
+    final_output: true
+    final_output_type: image
+
+  - stage_id: 2
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2audio
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: true
+      gpu_memory_utilization: 0.1
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 32768
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
+    final_output: true
+    final_output_type: audio
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1             # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1             # Simplified: process serially within each stage
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1

vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml

@@ -0,0 +1,114 @@
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2text
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: false
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 32768
+    output_connectors:
+      to_stage_1: mooncake_connector
+    final_output: true
+    final_output_type: text
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2image
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: false
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 32768
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2text_to_token2image
+    final_output: true
+    final_output_type: image
+    input_connectors:
+      from_stage_0: mooncake_connector
+    output_connectors:
+      to_stage_2: mooncake_connector
+
+  - stage_id: 2
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: token2audio
+      model_arch: DyninOmniForConditionalGeneration
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      engine_output_type: latent
+      trust_remote_code: false
+      enforce_eager: true
+      enable_prefix_caching: false
+      async_scheduling: false
+      max_num_batched_tokens: 32768
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.dynin_omni.token2image_to_token2audio
+    final_output: true
+    final_output_type: audio
+    input_connectors:
+      from_stage_1: mooncake_connector
+
+# Top-level runtime config (concise): default windows and stage edges
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1             # Simplified: trigger downstream only after full upstream completion
+    max_inflight: 1             # Simplified: process serially within each stage
+  ####
+  # same as Qwen2.5_omni version
+  # Distributed connectors configuration (optional)
+  # More connectors will be supported in the future.
+  connectors:
+    # Mooncake connector for cross-node/intra-node communication
+    mooncake_connector:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://10.90.67.86:8080/metadata"
+        master: "10.90.67.86:50051"
+        segment: 512000000    # 512MB
+        localbuf: 64000000     # 64MB
+        proto: "tcp"
+
+    # Yuanrong connector for cross-node/intra-node communication
+    yuanrong_connector:
+      name: YuanrongConnector
+      extra:
+        host: "127.0.0.1"
+        port: "35000"
+
+    # SharedMemory connector for intra-node communication
+    # Alternative SHM connector with different threshold
+    shared_memory_connector:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536 # 64KB threshold
+    ####
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1

vllm_omni/model_executor/stage_input_processors/dynin_omni.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import json
+from typing import Any
+
+import torch
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+
+def _to_prompt_dict(prompt_item: OmniTokensPrompt | TextPrompt | str | None) -> dict[str, Any]:
+    if isinstance(prompt_item, dict):
+        return prompt_item
+    return {}
+
+
+def _to_token_id_list(value: Any) -> list[int]:
+    if isinstance(value, torch.Tensor):
+        value = value.detach().to("cpu")
+        if value.ndim == 0:
+            return [int(value.item())]
+        if value.ndim > 1:
+            value = value[0]
+        return [int(x) for x in value.tolist()]
+    if isinstance(value, list):
+        if not value:
+            return []
+        if isinstance(value[0], list):
+            return [int(x) for x in value[0]]
+        return [int(x) for x in value]
+    if value is None:
+        return []
+    return [int(value)]
+
+
+def _to_int(value: Any, default: int = 0) -> int:
+    if isinstance(value, torch.Tensor):
+        if value.numel() == 0:
+            return default
+        return int(value.view(-1)[0].item())
+    if isinstance(value, list):
+        if not value:
+            return default
+        return int(value[0])
+    if value is None:
+        return default
+    return int(value)
+
+
+def _normalize_additional_info(value: Any) -> dict[str, Any]:
+    if not isinstance(value, dict):
+        return {}
+    normalized: dict[str, Any] = {}
+    for key, val in value.items():
+        if isinstance(val, list):
+            normalized[key] = val
+        else:
+            normalized[key] = [val]
+    return normalized
+
+
+def _decode_runtime_bridge_info(value: Any) -> dict[str, Any]:
+    if isinstance(value, torch.Tensor):
+        tensor = value.detach().to("cpu").reshape(-1).to(torch.uint8)
+        raw = bytes(tensor.tolist())
+    elif isinstance(value, (bytes, bytearray)):
+        raw = bytes(value)
+    elif isinstance(value, list):
+        try:
+            raw = bytes(int(item) for item in value)
+        except Exception:
+            return {}
+    elif value is None:
+        return {}
+    else:
+        return value if isinstance(value, dict) else {}
+
+    if not raw:
+        return {}
+
+    try:
+        decoded = json.loads(raw.decode("utf-8"))
+    except Exception:
+        return {}
+    return decoded if isinstance(decoded, dict) else {}
+
+
+def _bridge_tokens(
+    stage_list,
+    engine_input_source,
+    prompt: OmniTokensPrompt | TextPrompt = None,
+    requires_multimodal_data: bool = False,
+):
+    if not engine_input_source:
+        raise ValueError("engine_input_source cannot be empty")
+
+    source_stage_id = engine_input_source[0]
+    if source_stage_id >= len(stage_list):
+        raise IndexError(f"Invalid stage_id: {source_stage_id}")
+
+    if stage_list[source_stage_id].engine_outputs is None:
+        raise RuntimeError(f"Stage {source_stage_id} has no outputs yet")
+
+    source_outputs = stage_list[source_stage_id].engine_outputs
+    next_inputs = []
+    if not isinstance(prompt, list):
+        prompt = [prompt]
+
+    prompt_meta_by_reqid = {src_out.request_id: _to_prompt_dict(p) for src_out, p in zip(source_outputs, prompt)}
+
+    for source_output in source_outputs:
+        output = source_output.outputs[0]
+        mm_out = getattr(output, "multimodal_output", None) or {}
+
+        token_ids = _to_token_id_list(mm_out.get("token_ids"))
+        if not token_ids:
+            token_ids = _to_token_id_list(mm_out.get("text_tokens"))
+        if not token_ids:
+            token_ids = list(getattr(output, "token_ids", []) or [])
+        if not token_ids:
+            raise RuntimeError(
+                f"Stage {source_stage_id} output for request {source_output.request_id} has no token_ids"
+            )
+
+        detok_id = _to_int(mm_out.get("detok_id"), default=0)
+        src_prompt = prompt_meta_by_reqid.get(source_output.request_id, {})
+        src_additional_info = src_prompt.get("additional_information", {}) or {}
+        runtime_bridge_info = _decode_runtime_bridge_info(mm_out.get("runtime_info_json"))
+        if not runtime_bridge_info:
+            runtime_bridge_info = mm_out.get("runtime_info", {}) or {}
+
+        additional_information: dict[str, Any] = _normalize_additional_info(src_additional_info)
+        additional_information.update(_normalize_additional_info(runtime_bridge_info))
+        additional_information["detok_id"] = [detok_id]
+
+        next_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=token_ids,
+                additional_information=additional_information,
+                multi_modal_data=(src_prompt.get("multi_modal_data") if requires_multimodal_data else None),
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return next_inputs
+
+
+def token2text_to_token2image(
+    stage_list,
+    engine_input_source,
+    prompt: OmniTokensPrompt | TextPrompt = None,
+    requires_multimodal_data: bool = False,
+):
+    return _bridge_tokens(stage_list, engine_input_source, prompt, requires_multimodal_data)
+
+
+def token2image_to_token2audio(
+    stage_list,
+    engine_input_source,
+    prompt: OmniTokensPrompt | TextPrompt = None,
+    requires_multimodal_data: bool = False,
+):
+    return _bridge_tokens(stage_list, engine_input_source, prompt, requires_multimodal_data)