From e2b0ee4b8723366c927d64bdc7e41e7e67cc9495 Mon Sep 17 00:00:00 2001 From: Dogeun Kim <82812668+DOGEUNNKIM@users.noreply.github.com> Date: Thu, 9 Apr 2026 17:37:34 +0900 Subject: [PATCH] [Model] Add Dynin-omni model in vllm-omni (#1759) Signed-off-by: kdg6245@snu.ac.kr Signed-off-by: Yejoon Lee (IPAI) Signed-off-by: aidas (arpa-kt) Signed-off-by: Dogeun Kim <82812668+DOGEUNNKIM@users.noreply.github.com> Co-authored-by: Yejoon Lee (IPAI) Co-authored-by: aidas (arpa-kt) --- docs/models/supported_models.md | 1 + .../offline_inference/dynin_omni/README.md | 110 ++ .../offline_inference/dynin_omni/end2end.py | 1448 +++++++++++++++ examples/online_serving/dynin_omni/README.md | 97 + ...letion_client_for_multimodal_generation.py | 342 ++++ .../e2e/offline_inference/test_dynin_omni.py | 419 +++++ .../test_dynin_omni_expansion.py | 160 ++ tests/e2e/stage_configs/dynin_omni_ci.yaml | 84 + .../models/dynin_omni/__init__.py | 59 + .../models/dynin_omni/dynin_omni.py | 744 ++++++++ .../models/dynin_omni/dynin_omni_common.py | 1241 +++++++++++++ .../dynin_omni/dynin_omni_token2audio.py | 274 +++ .../dynin_omni/dynin_omni_token2image.py | 150 ++ .../dynin_omni/dynin_omni_token2text.py | 1580 +++++++++++++++++ vllm_omni/model_executor/models/registry.py | 5 + .../stage_configs/dynin_omni.yaml | 80 + .../dynin_omni_multiconnector.yaml | 114 ++ .../stage_input_processors/dynin_omni.py | 164 ++ 18 files changed, 7072 insertions(+) create mode 100644 examples/offline_inference/dynin_omni/README.md create mode 100644 examples/offline_inference/dynin_omni/end2end.py create mode 100644 examples/online_serving/dynin_omni/README.md create mode 100644 examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py create mode 100644 tests/e2e/offline_inference/test_dynin_omni.py create mode 100644 tests/e2e/online_serving/test_dynin_omni_expansion.py create mode 100644 tests/e2e/stage_configs/dynin_omni_ci.yaml create mode 100644 vllm_omni/model_executor/models/dynin_omni/__init__.py create mode 100644 vllm_omni/model_executor/models/dynin_omni/dynin_omni.py create mode 100644 vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py create mode 100644 vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py create mode 100644 vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2image.py create mode 100644 vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2text.py create mode 100644 vllm_omni/model_executor/stage_configs/dynin_omni.yaml create mode 100644 vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml create mode 100644 vllm_omni/model_executor/stage_input_processors/dynin_omni.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 8eab20edc..0f9c8fff6 100644 --- a/docs/models/supported_models.md +++ b/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. diff --git a/examples/offline_inference/dynin_omni/README.md b/examples/offline_inference/dynin_omni/README.md new file mode 100644 index 000000000..d28b36071 --- /dev/null +++ b/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 +``` + +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 /examples/offline_inference/dynin_omni/end2end.py \ + --task t2t --model snu-aidas/Dynin-Omni --text + +# i2t +python /examples/offline_inference/dynin_omni/end2end.py \ + --task i2t --model snu-aidas/Dynin-Omni --image --text "Please describe this image in detail." + +# s2t +python /examples/offline_inference/dynin_omni/end2end.py \ + --task s2t --model snu-aidas/Dynin-Omni --audio --text "Transcribe the given audio." + +# t2i +python /examples/offline_inference/dynin_omni/end2end.py \ + --task t2i --model snu-aidas/Dynin-Omni --text + +# v2t +python /examples/offline_inference/dynin_omni/end2end.py \ + --task v2t --model snu-aidas/Dynin-Omni --video --text "Describe this video in detail." + +# i2i +python /examples/offline_inference/dynin_omni/end2end.py \ + --task i2i --model snu-aidas/Dynin-Omni --image --text + +# t2s +python /examples/offline_inference/dynin_omni/end2end.py \ + --task t2s --model snu-aidas/Dynin-Omni --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 `. +- 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). diff --git a/examples/offline_inference/dynin_omni/end2end.py b/examples/offline_inference/dynin_omni/end2end.py new file mode 100644 index 000000000..66047934d --- /dev/null +++ b/examples/offline_inference/dynin_omni/end2end.py @@ -0,0 +1,1448 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import argparse +import json +import os +import re +import sys +import time +import types +from importlib.machinery import ModuleSpec +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image + +TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t") + +TASK_DEFAULT_RUNTIME = { + "t2t": ("mmu", "mmu", 0, "text"), + "t2i": ("t2i", "t2i_gen", 2, "image"), + "t2s": ("t2s_mmu_like", "t2s_gen", 1, "audio"), + "i2i": ("i2i", "i2i", 2, "image"), + "i2t": ("mmu", "mmu", 0, "text"), + "s2t": ("s2t", "s2t", 0, "text"), + "v2t": ("v2t", "v2t", 0, "text"), +} + +TASK_RUNTIME_FALLBACKS: dict[str, dict[str, Any]] = { + "t2t": { + "output_dir": "/tmp/dynin_end2end_outputs", + "prompt_max_text_len": 1024, + "max_new_tokens": 1024, + "steps": 1024, + "block_length": 16, + "temperature": 0.0, + "cfg_scale": 0.0, + }, + "t2i": { + "output_dir": "/tmp/dynin_t2i_outputs", + "prompt_max_text_len": 128, + "image_token_count": 1024, + "mask_token_id": 126336, + "codebook_size": 8192, + "timesteps": 20, + "guidance_scale": 3.5, + "temperature": 1.0, + }, + "i2i": { + "output_dir": "/tmp/dynin_i2i_outputs", + "prompt_max_text_len": 128, + "mask_token_id": 126336, + "codebook_size": 8192, + "timesteps": 64, + "guidance_scale": 3.5, + "temperature": 1.0, + "image_resolution": 336, + "use_train_i2i_prompt": True, + }, + "i2t": { + "output_dir": "/tmp/dynin_i2t_outputs", + "prompt_max_text_len": 128, + "max_new_tokens": 128, + "steps": 128, + "block_length": 2, + "temperature": 0.0, + "cfg_scale": 0.0, + "mask_token_id": 126336, + "codebook_size": 8192, + "image_resolution": 480, + "remasking": "low_confidence", + }, + "s2t": { + "output_dir": "/tmp/dynin_s2t_outputs", + "prompt_max_text_len": 1024, + "max_new_tokens": 128, + "steps": 128, + "block_length": 2, + "temperature": 0.0, + "cfg_scale": 0.0, + "mask_token_id": 126336, + "codebook_size": 8192, + "remasking": "low_confidence", + }, + "t2s": { + "output_dir": "/tmp/dynin_t2s_outputs", + "runtime_task": "t2s_mmu_like", + "prompting_task": "t2s_gen", + "prompt_max_text_len": 1024, + "t2s_token_length": 512, + "mask_token_id": 126336, + "codebook_size": 8192, + "audio_codebook_size": 4096, + "steps": 512, + "block_length": 128, + "temperature": 1.0, + "cfg_scale": 2.5, + "t2s_condition": "gender-female_emotion-neutral_speed-normal_pitch-normal", + }, + "v2t": { + "output_dir": "/tmp/dynin_v2t_outputs", + "prompt_max_text_len": 1024, + "max_new_tokens": 128, + "steps": 128, + "block_length": 2, + "temperature": 0.0, + "cfg_scale": 0.0, + "mask_token_id": 126336, + "codebook_size": 8192, + "image_resolution": 224, + "num_frames": 5, + "remasking": "low_confidence", + }, +} + +DEFAULT_I2T_QUESTION = "Please describe this image in detail." +DEFAULT_S2T_INSTRUCTION = "Transcribe the given audio." +DEFAULT_V2T_QUESTION = "Please provide a detailed description of the video." +DEFAULT_T2T_PROMPT = "Explain multimodal LLM inference in 3 sentences." +DEFAULT_T2S_INSTRUCTION = "Convert the given text into spoken audio." +DEFAULT_T2S_PROMPT = "Hello. This is a default text-to-speech sample." + +DYNIN_SPECIAL_TOKENS = ( + "<|soi|>", + "<|eoi|>", + "<|sov|>", + "<|eov|>", + "<|t2i|>", + "<|mmu|>", + "<|t2v|>", + "<|v2v|>", + "<|lvg|>", + "<|i2i|>", + "<|ti2ti|>", + "<|v2t|>", + "<|v2s|>", + "<|s2t|>", + "<|t2s|>", + "<|s2s|>", + "<|soa|>", + "<|eoa|>", +) + + +def bootstrap_repo_path() -> Path: + repo_root = Path(__file__).resolve().parents[3] + repo_root_str = str(repo_root) + if repo_root_str not in sys.path: + sys.path.insert(0, repo_root_str) + return repo_root + + +def ensure_safe_import_for_vllm() -> None: + os.environ.setdefault("TRANSFORMERS_NO_TORCHVISION", "1") + try: + import torchvision # noqa: F401 + + return + except Exception: + pass + + import enum + + class _InterpolationMode(enum.Enum): + NEAREST = 0 + BILINEAR = 2 + BICUBIC = 3 + LANCZOS = 1 + HAMMING = 4 + BOX = 5 + + tv_mod = types.ModuleType("torchvision") + tv_mod.__dict__["__version__"] = "0.0-stub" + tv_mod.__spec__ = ModuleSpec(name="torchvision", loader=None) + transforms_mod = types.ModuleType("torchvision.transforms") + transforms_mod.__spec__ = ModuleSpec(name="torchvision.transforms", loader=None) + transforms_mod.InterpolationMode = _InterpolationMode + tv_mod.transforms = transforms_mod + sys.modules["torchvision"] = tv_mod + sys.modules["torchvision.transforms"] = transforms_mod + + +def sanitize_repo_id(repo_id: str) -> str: + return re.sub(r"[^a-zA-Z0-9._-]+", "_", repo_id) + + +def is_hf_repo_id(value: str) -> bool: + return isinstance(value, str) and value.count("/") == 1 and all(value.split("/", 1)) + + +def ensure_local_model_dir(model: str, cache_dir: Path, localize: bool) -> Path: + model_path = Path(model).expanduser() + if model_path.is_dir(): + return model_path.resolve() + if not localize: + return Path(model) + + from huggingface_hub import snapshot_download + + cache_dir.mkdir(parents=True, exist_ok=True) + os.environ.setdefault("HF_HOME", str(cache_dir / ".hf_home")) + local_dir = cache_dir / sanitize_repo_id(model) + if not local_dir.exists(): + print(f"[end2end] Downloading model into local cache: {local_dir}") + snapshot_download( + repo_id=model, + local_dir=str(local_dir), + local_dir_use_symlinks=True, + resume_download=True, + ) + return local_dir.resolve() + + +def resolve_local_only( + override: bool | None, + source: str, + default: bool, +) -> bool: + if override is not None: + return bool(override) + return default or Path(source).expanduser().is_dir() + + +def load_text_tokenizer(tokenizer_source: str, local_files_only: bool): + from transformers import AutoTokenizer + + kwargs = { + "trust_remote_code": True, + "padding_side": "left", + "local_files_only": bool(local_files_only), + } + try: + tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, **kwargs) + except TypeError: + kwargs.pop("local_files_only", None) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, **kwargs) + return tokenizer + + +def preprocess_image(image: Image.Image, resolution: int) -> torch.Tensor: + w, h = image.size + short_side = min(w, h) + scale = resolution / short_side + new_w, new_h = round(w * scale), round(h * scale) + image = image.resize((new_w, new_h), Image.BICUBIC) + left = (new_w - resolution) // 2 + top = (new_h - resolution) // 2 + image = image.crop((left, top, left + resolution, top + resolution)) + arr = np.array(image, dtype=np.float32) / 255.0 + tensor = torch.from_numpy(arr).permute(2, 0, 1) + return (tensor - 0.5) / 0.5 + + +def load_vq_image_encoder(source: str, local_files_only: bool, device: torch.device) -> Any: + from vllm_omni.model_executor.models.dynin_omni.dynin_omni_common import get_dynin_magvit_attr + + MAGVITv2 = get_dynin_magvit_attr("MAGVITv2", source=source, local_files_only=local_files_only) + vq_model = MAGVITv2.from_pretrained(source, local_files_only=local_files_only).to(device) + vq_model.requires_grad_(False) + vq_model.eval() + return vq_model + + +def encode_image_tokens( + image_path: Path, + vq_model: Any, + device: torch.device, + resolution: int, +) -> torch.Tensor: + image = Image.open(image_path).convert("RGB") + image_tensor = preprocess_image(image, resolution=resolution).unsqueeze(0).to(device) + with torch.no_grad(): + token_ids = vq_model.get_code(image_tensor) + token_ids = torch.as_tensor(token_ids, dtype=torch.long).detach().cpu() + if token_ids.ndim == 2 and token_ids.shape[0] == 1: + token_ids = token_ids[0] + return token_ids.contiguous() + + +def encode_video_tokens( + video_path: Path, + vq_model: Any, + device: torch.device, + resolution: int, + num_frames: int, +) -> torch.Tensor: + import cv2 + + cap = cv2.VideoCapture(str(video_path)) + frames: list[np.ndarray] = [] + while True: + ok, frame = cap.read() + if not ok: + break + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frames.append(frame) + cap.release() + if not frames: + raise ValueError(f"Video has no readable frames: {video_path}") + if len(frames) < num_frames: + raise ValueError(f"Video has {len(frames)} frames, requires >= {num_frames}: {video_path}") + + indices = np.linspace(0, len(frames) - 1, num_frames).astype(int) + token_list: list[torch.Tensor] = [] + for idx in indices: + pil = Image.fromarray(frames[int(idx)]) + frame_tensor = preprocess_image(pil, resolution=resolution).unsqueeze(0).to(device) + with torch.no_grad(): + token_list.append(torch.as_tensor(vq_model.get_code(frame_tensor), dtype=torch.long)) + merged = torch.cat(token_list, dim=1).detach().cpu() + if merged.ndim == 2 and merged.shape[0] == 1: + merged = merged[0] + return merged.contiguous() + + +def load_vq_audio_encoder(source: str, local_files_only: bool, device: torch.device) -> Any: + from transformers import AutoModel + + kwargs = { + "trust_remote_code": True, + "local_files_only": bool(local_files_only), + "low_cpu_mem_usage": False, + } + try: + model = AutoModel.from_pretrained(source, **kwargs) + except TypeError: + kwargs.pop("low_cpu_mem_usage", None) + try: + model = AutoModel.from_pretrained(source, **kwargs) + except TypeError: + kwargs.pop("local_files_only", None) + model = AutoModel.from_pretrained(source, **kwargs) + model.requires_grad_(False) + model.eval() + if hasattr(model, "to"): + model = model.to(device) + return model + + +def encode_audio_tokens(audio_path: Path, vq_audio_model: Any) -> torch.Tensor: + encoded = vq_audio_model.encode(str(audio_path)) + if isinstance(encoded, dict): + for key in ("input_ids", "token_ids", "codes", "tokens"): + if key in encoded: + encoded = encoded[key] + break + encoded = torch.as_tensor(encoded, dtype=torch.long).detach().cpu() + if encoded.ndim == 1: + encoded = encoded.unsqueeze(0) + elif encoded.ndim > 2: + encoded = encoded.view(encoded.shape[0], -1) + return encoded.contiguous() + + +def build_chat_prompt(content: str) -> str: + return ( + f"<|start_header_id|>user<|end_header_id|>\n{content}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" + ) + + +def resolve_task_text( + *, + task_name: str, + text: str, + instruction: str = "", + raw_prompt: bool = False, +) -> str: + text = str(text or "").strip() + + if task_name == "t2t" and not text: + return DEFAULT_T2T_PROMPT + if task_name == "i2t" and not text: + return DEFAULT_I2T_QUESTION + if task_name == "s2t" and not text: + return DEFAULT_S2T_INSTRUCTION + if task_name == "v2t" and not text: + return DEFAULT_V2T_QUESTION + if task_name in {"t2i", "i2i"} and not text: + return "A high quality detailed image." + + if task_name != "t2s": + return text + + if not text: + text = DEFAULT_T2S_PROMPT + + if raw_prompt: + return text + + instruction = str(instruction or "").strip() or DEFAULT_T2S_INSTRUCTION + return build_chat_prompt(f"{instruction}\n{text}") + + +def load_universal_prompting( + *, + tokenizer: Any, + tokenizer_source: str, + max_text_len: int, + cond_dropout_prob: float, + local_files_only: bool, + max_audio_len: int = 512, + max_audio_len_short: int = 256, +) -> Any: + from vllm_omni.model_executor.models.dynin_omni.dynin_omni_common import ( + DYNIN_REMOTE_SETTINGS, + resolve_remote_attr, + ) + + UniversalPrompting = resolve_remote_attr( + "UniversalPrompting", + module_name="prompting_utils", + settings=DYNIN_REMOTE_SETTINGS, + source=tokenizer_source, + local_files_only=bool(local_files_only), + fallback_module_names=("modeling_dynin_omni",), + ) + init_kwargs: dict[str, Any] = { + "max_text_len": int(max_text_len), + "special_tokens": DYNIN_SPECIAL_TOKENS, + "ignore_id": -100, + "cond_dropout_prob": float(cond_dropout_prob), + "use_reserved_token": True, + "max_audio_len": int(max_audio_len), + "max_audio_len_short": int(max_audio_len_short), + } + try: + return UniversalPrompting(tokenizer, **init_kwargs) + except TypeError: + init_kwargs.pop("max_audio_len", None) + init_kwargs.pop("max_audio_len_short", None) + return UniversalPrompting(tokenizer, **init_kwargs) + + +def _runtime_fallback(task: str, key: str, value: Any) -> Any: + if isinstance(value, str): + if value.strip() != "": + return value + elif value is not None: + return value + return TASK_RUNTIME_FALLBACKS.get(task, {}).get(key) + + +def _validate_generation_args(*, task: str, max_new_tokens: int, steps: int, block_length: int) -> None: + # Keep i2t/v2t generation constraints aligned with i2t.py/v2t.py. + if task not in {"i2t", "v2t"}: + return + if max_new_tokens <= 0: + raise ValueError(f"{task} requires max_new_tokens > 0.") + if block_length <= 0: + raise ValueError(f"{task} requires block_length > 0.") + if steps <= 0: + raise ValueError(f"{task} requires steps > 0.") + if max_new_tokens % block_length != 0: + raise ValueError(f"{task} requires max_new_tokens % block_length == 0, got {max_new_tokens} % {block_length}") + num_blocks = max_new_tokens // block_length + if num_blocks <= 0: + raise ValueError(f"{task} has invalid num_blocks.") + if steps % num_blocks != 0: + raise ValueError( + f"{task} requires steps % (max_new_tokens // block_length) == 0, " + f"got steps={steps}, max_new_tokens={max_new_tokens}, block_length={block_length}" + ) + + +def make_prompt_payload( + *, + task: str, + text: str, + image_tokens: torch.Tensor | None, + audio_tokens: torch.Tensor | None, + video_tokens: torch.Tensor | None, + image_placeholder_tokens: int, + audio_placeholder_tokens: int, + image_token_offset: int, + speech_token_offset: int, + mask_token_id: int, + use_train_i2i_prompt: bool, +) -> tuple[Any, str]: + runtime_task, prompting_task, _, _ = TASK_DEFAULT_RUNTIME[task] + del runtime_task + + if task == "t2t": + payload = ([[]], [build_chat_prompt(text)]) + return payload, prompting_task + + if task == "i2t": + if image_tokens is None: + raise ValueError("i2t requires image tokens") + img = image_tokens.view(-1).long() + int(image_token_offset) + payload = ([[img]], [build_chat_prompt(text)]) + return payload, prompting_task + + if task == "s2t": + if audio_tokens is None: + raise ValueError("s2t requires audio tokens") + aud = audio_tokens.long() + int(speech_token_offset) + if aud.ndim == 1: + aud = aud.unsqueeze(0) + payload = ([aud], [build_chat_prompt(text)]) + return payload, prompting_task + + if task == "v2t": + if video_tokens is None: + raise ValueError("v2t requires video tokens") + vid = video_tokens.view(-1).long() + int(image_token_offset) + payload = (vid.unsqueeze(0), [build_chat_prompt(text)]) + return payload, prompting_task + + if task == "t2i": + image_placeholder = torch.full( + (1, int(image_placeholder_tokens)), + fill_value=int(mask_token_id), + dtype=torch.long, + ) + payload = ([text], image_placeholder) + return payload, prompting_task + + if task == "i2i": + if image_tokens is None: + raise ValueError("i2i requires image tokens") + src = image_tokens.view(1, -1).long() + int(image_token_offset) + target_len = int(image_placeholder_tokens) if image_placeholder_tokens > 0 else int(src.shape[1]) + image_placeholder = torch.full( + (1, target_len), + fill_value=int(mask_token_id), + dtype=torch.long, + ) + if use_train_i2i_prompt: + labels_placeholder = torch.full( + (1, target_len), + fill_value=-100, + dtype=torch.long, + ) + payload = ([text], src, image_placeholder, labels_placeholder) + return payload, "i2i" + payload = ([text], src, image_placeholder) + return payload, "i2i_gen" + + if task == "t2s": + audio_placeholder = torch.full( + (1, int(audio_placeholder_tokens)), + fill_value=int(mask_token_id), + dtype=torch.long, + ) + payload = ([text], audio_placeholder) + return payload, prompting_task + + raise ValueError(f"Unsupported task: {task}") + + +def _to_1d_int_list(value: Any) -> list[int]: + if value is None: + return [] + if isinstance(value, torch.Tensor): + tensor = value.detach().to(device="cpu", dtype=torch.long) + else: + tensor = torch.as_tensor(value, dtype=torch.long) + if tensor.ndim == 0: + tensor = tensor.view(1) + elif tensor.ndim >= 2: + tensor = tensor.view(tensor.shape[0], -1)[0] + return [int(v) for v in tensor.tolist()] + + +def _run_uni_prompting(uni_prompting: Any, payload: Any, prompting_task: str) -> tuple[list[int], list[int]]: + prepared = uni_prompting(payload, prompting_task) + if isinstance(prepared, tuple): + prepared_input_ids = prepared[0] if len(prepared) > 0 else None + prepared_attention_mask = prepared[1] if len(prepared) > 1 else None + else: + prepared_input_ids = prepared + prepared_attention_mask = None + + input_ids = _to_1d_int_list(prepared_input_ids) + attention_mask = _to_1d_int_list(prepared_attention_mask) + if not input_ids: + raise RuntimeError(f"UniversalPrompting returned empty input_ids for task={prompting_task}") + return input_ids, attention_mask + + +def _get_special_token_id(uni_prompting: Any, token: str) -> int: + sptids = getattr(uni_prompting, "sptids_dict", None) or {} + if token not in sptids: + raise KeyError(f"Special token not found in UniversalPrompting.sptids_dict: {token}") + token_ids = _to_1d_int_list(sptids[token]) + if not token_ids: + raise ValueError(f"Special token id is empty for token: {token}") + return int(token_ids[0]) + + +def _tokenize_chat_query(tokenizer: Any, text: str) -> list[int]: + encoded = tokenizer(build_chat_prompt(text), return_tensors="pt").input_ids[0] + token_ids = _to_1d_int_list(encoded) + if not token_ids: + raise RuntimeError("Failed to tokenize chat query text.") + return token_ids + + +def _flatten_media_token_ids_with_offset(token_ids: Any, token_offset: int) -> list[int]: + media_ids = token_ids + if isinstance(media_ids, torch.Tensor): + media_ids = media_ids.detach().cpu().reshape(-1).tolist() + else: + media_ids = np.asarray(media_ids).reshape(-1).tolist() + return [int(x) + int(token_offset) for x in media_ids] + + +def _scalar_token_id(value: Any) -> int: + if isinstance(value, torch.Tensor): + if value.numel() == 0: + raise ValueError("Empty special-token tensor.") + return int(value.view(-1)[0].item()) + if isinstance(value, (list, tuple)): + if not value: + raise ValueError("Empty special-token list.") + return int(value[0]) + return int(value) + + +def build_v2t_input_ids( + *, + video_token_ids: Any, + tokenizer: Any, + uni_prompting: Any, + question: str, + image_token_offset: int, +) -> tuple[list[int], str]: + media_ids = video_token_ids + if isinstance(media_ids, torch.Tensor): + media_ids = media_ids.detach().cpu().reshape(-1).tolist() + else: + media_ids = np.asarray(media_ids).reshape(-1).tolist() + media_ids = [int(x) + int(image_token_offset) for x in media_ids] + + sptids = uni_prompting.sptids_dict + task_id = _scalar_token_id(sptids["<|v2t|>"]) + soi_id = _scalar_token_id(sptids["<|soi|>"]) + eoi_id = _scalar_token_id(sptids["<|eoi|>"]) + sot_id = _scalar_token_id(sptids["<|sot|>"]) + + prompt_text = build_v2t_chat_prompt(question) + query_ids = tokenizer(prompt_text, return_tensors="pt").input_ids[0].detach().cpu().tolist() + input_ids = [task_id, soi_id] + media_ids + [eoi_id, sot_id] + [int(v) for v in query_ids] + return input_ids, prompt_text + + +def build_i2t_input_ids( + *, + image_token_ids: Any, + tokenizer: Any, + uni_prompting: Any, + question: str, + image_token_offset: int, +) -> tuple[list[int], str]: + image_ids = image_token_ids + if isinstance(image_ids, torch.Tensor): + image_ids = image_ids.detach().cpu().reshape(-1).tolist() + else: + image_ids = np.asarray(image_ids).reshape(-1).tolist() + image_ids = [int(x) + int(image_token_offset) for x in image_ids] + + sptids = uni_prompting.sptids_dict + task_id = _scalar_token_id(sptids["<|mmu|>"]) + soi_id = _scalar_token_id(sptids["<|soi|>"]) + eoi_id = _scalar_token_id(sptids["<|eoi|>"]) + sot_id = _scalar_token_id(sptids["<|sot|>"]) + + prompt_text = build_i2t_chat_prompt(question) + query_ids = tokenizer(prompt_text, return_tensors="pt").input_ids[0].detach().cpu().tolist() + input_ids = [task_id, soi_id] + image_ids + [eoi_id, sot_id] + [int(v) for v in query_ids] + return input_ids, prompt_text + + +def build_v2t_chat_prompt(question: str) -> str: + return ( + f"<|start_header_id|>user<|end_header_id|>\n{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" + ) + + +def build_i2t_chat_prompt(question: str) -> str: + return ( + f"<|start_header_id|>user<|end_header_id|>\n{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" + ) + + +def make_mmu_prompt( + *, + task: str, + text: str, + tokenizer: Any, + uni_prompting: Any, + image_tokens: torch.Tensor | None, + audio_tokens: torch.Tensor | None, + video_tokens: torch.Tensor | None, + image_token_offset: int, + speech_token_offset: int, +) -> tuple[list[int], list[int]]: + query_ids = _tokenize_chat_query(tokenizer, text) + + if task == "i2t": + token_ids, _ = build_i2t_input_ids( + image_token_ids=image_tokens, + tokenizer=tokenizer, + uni_prompting=uni_prompting, + question=text, + image_token_offset=int(image_token_offset), + ) + token_ids = [int(v) for v in token_ids] + return token_ids, [1] * len(token_ids) + + if task == "v2t": + token_ids, _ = build_v2t_input_ids( + video_token_ids=video_tokens, + tokenizer=tokenizer, + uni_prompting=uni_prompting, + question=text, + image_token_offset=int(image_token_offset), + ) + token_ids = [int(v) for v in token_ids] + return token_ids, [1] * len(token_ids) + + if task == "s2t": + if audio_tokens is None: + raise ValueError("s2t requires audio tokens") + audio_ids = _to_1d_int_list(audio_tokens.long() + int(speech_token_offset)) + token_ids = [ + _get_special_token_id(uni_prompting, "<|s2t|>"), + _get_special_token_id(uni_prompting, "<|soa|>"), + *audio_ids, + _get_special_token_id(uni_prompting, "<|eoa|>"), + *query_ids, + ] + return token_ids, [1] * len(token_ids) + + raise ValueError(f"Unsupported task for validation-style MMU prompt: {task}") + + +def iter_mm_outputs(outputs: list[Any]): + for omni_out in outputs: + req_out = getattr(omni_out, "request_output", None) + req_list = req_out if isinstance(req_out, list) else [req_out] + for item in req_list: + if item is None: + continue + mm_out = getattr(item, "multimodal_output", None) or {} + if mm_out: + yield mm_out + completions = getattr(item, "outputs", None) or [] + for completion in completions: + c_mm_out = getattr(completion, "multimodal_output", None) or {} + if c_mm_out: + yield c_mm_out + omni_mm = getattr(omni_out, "multimodal_output", None) or {} + if omni_mm: + yield omni_mm + + +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_output(outputs: list[Any], tokenizer: Any) -> str: + for mm_out in iter_mm_outputs(outputs): + text = mm_out.get("text") + if isinstance(text, list) and text: + text = text[-1] + if isinstance(text, str) and text.strip(): + return text.strip() + 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() + return "" + + +def extract_image_output(outputs: list[Any]) -> torch.Tensor | None: + for mm_out in iter_mm_outputs(outputs): + image = mm_out.get("image") + if isinstance(image, list) and image: + image = image[-1] + if isinstance(image, torch.Tensor): + return image + return None + + +def tensor_to_pil_image(image: torch.Tensor) -> Image.Image: + arr = image.detach().cpu().numpy() + if arr.ndim == 4: + arr = arr[0] + if arr.ndim == 3 and arr.shape[0] in (1, 3, 4): + arr = np.transpose(arr, (1, 2, 0)) + if arr.dtype != np.uint8: + arr = arr.astype(np.float32) + if arr.max() <= 1.0: + arr = arr * 255.0 + arr = np.clip(arr, 0.0, 255.0).astype(np.uint8) + if arr.ndim == 3 and arr.shape[-1] == 1: + arr = arr[..., 0] + return Image.fromarray(arr) + + +def extract_audio_output(outputs: list[Any]) -> tuple[np.ndarray, int] | None: + for mm_out in iter_mm_outputs(outputs): + audio = mm_out.get("audio") + if audio is None: + audio = mm_out.get("speech") + if audio is None: + continue + + def _to_wav_array(value: Any) -> np.ndarray: + if isinstance(value, torch.Tensor): + return value.detach().cpu().numpy().reshape(-1).astype(np.float32) + return np.asarray(value).reshape(-1).astype(np.float32) + + if isinstance(audio, list): + chunks = [_to_wav_array(chunk) for chunk in audio] + wav = np.concatenate(chunks, axis=0) if chunks else np.zeros((0,), dtype=np.float32) + else: + wav = _to_wav_array(audio) + sr = mm_out.get("sr", 24000) + if hasattr(sr, "item"): + try: + sr = int(sr.item()) + except Exception: + sr = 24000 + elif isinstance(sr, list): + sr = int(sr[0]) if sr else 24000 + else: + sr = int(sr) + return wav, sr + return None + + +def save_audio_wav(path: Path, wav: np.ndarray, sr: int) -> None: + try: + import soundfile as sf + + sf.write(str(path), wav, int(sr), format="WAV") + except Exception: + from scipy.io import wavfile + + wav_i16 = np.clip(wav, -1.0, 1.0) + wav_i16 = (wav_i16 * 32767.0).astype(np.int16) + wavfile.write(str(path), int(sr), wav_i16) + + +def parse_args(repo_root: Path) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Dynin-Omni unified offline end2end example.") + parser.add_argument("--task", type=str, required=True, choices=TASK_CHOICES) + parser.add_argument("--model", type=str, required=True, help="HF repo id or local model directory.") + parser.add_argument( + "--stage-config-path", + type=str, + default=str(repo_root / "vllm_omni/model_executor/stage_configs/dynin_omni.yaml"), + help="Path to stage config yaml.", + ) + parser.add_argument( + "--dynin-config-path", + type=str, + default="", + help="Path to DYNIN config yaml (passed through additional_information).", + ) + parser.add_argument( + "--model-cache-dir", + type=str, + default="/tmp/dynin_localized_models", + help="Cache directory used when --model is HF repo id.", + ) + parser.add_argument( + "--localize-model", + action=argparse.BooleanOptionalAction, + default=True, + help="If true and --model is HF repo id, snapshot it under --model-cache-dir.", + ) + parser.add_argument("--text", type=str, default="", help="Prompt/edit/question text.") + parser.add_argument("--instruction", type=str, default="", help="Optional extra instruction.") + parser.add_argument("--raw-prompt", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--image", type=str, default="", help="Input image path for i2i/i2t.") + parser.add_argument("--audio", type=str, default="", help="Input audio path for s2t.") + parser.add_argument("--video", type=str, default="", help="Input video path for v2t.") + parser.add_argument("--image-resolution", type=int, default=None) + parser.add_argument("--num-frames", type=int, default=None) + parser.add_argument( + "--output-dir", + type=str, + default="", + help="Directory for generated outputs.", + ) + parser.add_argument("--output-prefix", type=str, default="") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--dtype", type=str, default="auto") + parser.add_argument("--max-tokens-per-stage", type=int, default=1) + + parser.add_argument("--runtime-task", type=str, default="", help="Override runtime task key.") + parser.add_argument("--prompting-task", type=str, default="", help="Override prompting task key.") + parser.add_argument("--detok-id", type=int, default=None, help="Override detok id.") + + parser.add_argument("--prompt-max-text-len", type=int, default=None) + parser.add_argument("--cond-dropout-prob", type=float, default=0.0) + parser.add_argument("--max-new-tokens", type=int, default=None) + parser.add_argument("--steps", type=int, default=None) + parser.add_argument("--block-length", type=int, default=None) + parser.add_argument("--temperature", type=float, default=None) + parser.add_argument("--cfg-scale", type=float, default=None) + parser.add_argument("--remasking", type=str, default="low_confidence") + + parser.add_argument("--timesteps", type=int, default=None) + parser.add_argument("--guidance-scale", type=float, default=None) + parser.add_argument("--noise-type", type=str, default="mask") + parser.add_argument("--noise-schedule-name", type=str, default="cosine") + parser.add_argument("--noise-schedule-params", type=str, default="{}") + + parser.add_argument("--mask-token-id", type=int, default=None) + parser.add_argument("--codebook-size", type=int, default=None) + parser.add_argument("--audio-codebook-size", type=int, default=None) + parser.add_argument("--image-token-count", type=int, default=None) + parser.add_argument("--t2s-token-length", type=int, default=None) + parser.add_argument( + "--t2s-condition", + type=str, + default="", + ) + parser.add_argument( + "--use-train-i2i-prompt", + action="store_true", + help="Use i2i training prompt template (default behavior of i2i.py).", + ) + parser.add_argument( + "--no-use-train-i2i-prompt", + dest="use_train_i2i_prompt", + action="store_false", + help="Use i2i_gen prompt template.", + ) + parser.set_defaults(use_train_i2i_prompt=None) + + parser.add_argument("--tokenizer-path", type=str, default="") + parser.add_argument("--model-local-files-only", action=argparse.BooleanOptionalAction, default=None) + parser.add_argument("--tokenizer-local-files-only", action=argparse.BooleanOptionalAction, default=None) + + parser.add_argument("--vq-model-image-path", type=str, default="") + parser.add_argument("--vq-model-image-local-files-only", action=argparse.BooleanOptionalAction, default=None) + parser.add_argument("--vq-model-audio-path", type=str, default="") + parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None) + + parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True) + return parser.parse_args() + + +def main() -> None: + repo_root = bootstrap_repo_path() + ensure_safe_import_for_vllm() + from vllm_omni.model_executor.models.dynin_omni.dynin_omni_common import ( + DYNIN_PROMPT_SOURCE_KEY, + DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT, + ) + + args = parse_args(repo_root) + + if args.disable_hf_xet: + os.environ.setdefault("HF_HUB_DISABLE_XET", "1") + + np.random.seed(args.seed) + torch.manual_seed(args.seed) + + model_dir = ensure_local_model_dir( + model=args.model, + cache_dir=Path(args.model_cache_dir).expanduser(), + localize=bool(args.localize_model), + ) + model_source = str(model_dir) + + task_name = str(args.task) + dynin_config_path = str(Path(args.dynin_config_path).expanduser()) + os.environ["DYNIN_CONFIG_PATH"] = dynin_config_path + default_runtime_task, default_prompting_task, default_detok_id, final_modality = TASK_DEFAULT_RUNTIME[task_name] + runtime_task = args.runtime_task.strip() or str( + _runtime_fallback(task_name, "runtime_task", None) or default_runtime_task + ) + prompting_task = args.prompting_task.strip() or str( + _runtime_fallback(task_name, "prompting_task", None) or default_prompting_task + ) + detok_id_default = _runtime_fallback(task_name, "detok_id", None) + if detok_id_default is None: + detok_id_default = default_detok_id + detok_id = int(detok_id_default if args.detok_id is None else args.detok_id) + + output_dir_default = _runtime_fallback(task_name, "output_dir", args.output_dir) + resolved_output_dir = str(output_dir_default or "/tmp/dynin_end2end_outputs") + + image_resolution_value = _runtime_fallback( + task_name, + "image_resolution", + args.image_resolution, + ) + if image_resolution_value is None: + image_resolution_value = 336 + image_resolution = int(image_resolution_value) + + num_frames_value = _runtime_fallback( + task_name, + "num_frames", + args.num_frames, + ) + if num_frames_value is None: + num_frames_value = 8 + num_frames = int(num_frames_value) + + prompt_max_text_len_value = _runtime_fallback( + task_name, + "prompt_max_text_len", + args.prompt_max_text_len, + ) + if prompt_max_text_len_value is None: + prompt_max_text_len_value = 1024 + prompt_max_text_len = int(prompt_max_text_len_value) + + max_new_tokens_value = _runtime_fallback( + task_name, + "max_new_tokens", + args.max_new_tokens, + ) + if max_new_tokens_value is None: + max_new_tokens_value = 256 + max_new_tokens = int(max_new_tokens_value) + + steps_value = _runtime_fallback( + task_name, + "steps", + args.steps, + ) + if steps_value is None: + steps_value = 256 + steps = int(steps_value) + + block_length_value = _runtime_fallback( + task_name, + "block_length", + args.block_length, + ) + if block_length_value is None: + block_length_value = 2 + block_length = int(block_length_value) + + temperature_value = _runtime_fallback( + task_name, + "temperature", + args.temperature, + ) + if temperature_value is None: + temperature_value = 0.0 + temperature = float(temperature_value) + + cfg_scale_value = _runtime_fallback( + task_name, + "cfg_scale", + args.cfg_scale, + ) + if cfg_scale_value is None: + cfg_scale_value = 0.0 + cfg_scale = float(cfg_scale_value) + + remasking = str(_runtime_fallback(task_name, "remasking", args.remasking) or "low_confidence") + + timesteps_value = _runtime_fallback( + task_name, + "timesteps", + args.timesteps, + ) + if timesteps_value is None: + timesteps_value = 20 + timesteps = int(timesteps_value) + + guidance_scale_value = _runtime_fallback( + task_name, + "guidance_scale", + args.guidance_scale, + ) + if guidance_scale_value is None: + guidance_scale_value = 0.0 + guidance_scale = float(guidance_scale_value) + + mask_token_id_value = _runtime_fallback( + task_name, + "mask_token_id", + args.mask_token_id, + ) + if mask_token_id_value is None: + mask_token_id_value = 126336 + mask_token_id = int(mask_token_id_value) + + codebook_size_value = _runtime_fallback( + task_name, + "codebook_size", + args.codebook_size, + ) + if codebook_size_value is None: + codebook_size_value = 8192 + codebook_size = int(codebook_size_value) + + audio_codebook_size_value = _runtime_fallback( + task_name, + "audio_codebook_size", + args.audio_codebook_size, + ) + if audio_codebook_size_value is None: + audio_codebook_size_value = 4096 + audio_codebook_size = int(audio_codebook_size_value) + + image_token_count_value = _runtime_fallback( + task_name, + "image_token_count", + args.image_token_count, + ) + image_token_count = int(image_token_count_value) if image_token_count_value is not None else 0 + + t2s_token_length_value = _runtime_fallback( + task_name, + "t2s_token_length", + args.t2s_token_length, + ) + if t2s_token_length_value is None: + t2s_token_length_value = 383 + t2s_token_length = int(t2s_token_length_value) + + t2s_condition = str( + _runtime_fallback(task_name, "t2s_condition", args.t2s_condition) + or "gender-female_emotion-neutral_speed-normal_pitch-normal" + ) + + _validate_generation_args( + task=task_name, + max_new_tokens=max_new_tokens, + steps=steps, + block_length=block_length, + ) + + use_train_i2i_prompt = _runtime_fallback(task_name, "use_train_i2i_prompt", args.use_train_i2i_prompt) + if use_train_i2i_prompt is None: + use_train_i2i_prompt = bool(task_name == "i2i") + use_train_i2i_prompt = bool(use_train_i2i_prompt) + + if task_name in {"i2i", "i2t"} and not args.image: + raise ValueError(f"--task {task_name} requires --image") + if task_name == "s2t" and not args.audio: + raise ValueError("--task s2t requires --audio") + if task_name == "v2t" and not args.video: + raise ValueError("--task v2t requires --video") + + text = resolve_task_text( + task_name=task_name, + text=args.text, + instruction=args.instruction, + raw_prompt=bool(args.raw_prompt), + ) + + tokenizer_source = args.tokenizer_path.strip() or model_source + model_local_only = resolve_local_only( + args.model_local_files_only, model_source, default=Path(model_source).is_dir() + ) + tokenizer_local_only = resolve_local_only( + args.tokenizer_local_files_only, + tokenizer_source, + default=model_local_only, + ) + tokenizer = load_text_tokenizer(tokenizer_source, local_files_only=tokenizer_local_only) + text_vocab_size = int(len(tokenizer)) + + image_tokens: torch.Tensor | None = None + audio_tokens: torch.Tensor | None = None + video_tokens: torch.Tensor | None = None + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + vq_image_source = args.vq_model_image_path.strip() or "snu-aidas/magvitv2" + vq_audio_source = args.vq_model_audio_path.strip() or "snu-aidas/emova_speech_tokenizer_vllm" + vq_image_local_only = resolve_local_only(args.vq_model_image_local_files_only, vq_image_source, default=False) + vq_audio_local_only = resolve_local_only(args.vq_model_audio_local_files_only, vq_audio_source, default=False) + + if task_name in {"i2i", "i2t", "v2t"}: + vq_image = load_vq_image_encoder(vq_image_source, vq_image_local_only, device) + if task_name in {"i2i", "i2t"}: + image_tokens = encode_image_tokens( + Path(args.image).expanduser().resolve(), + vq_model=vq_image, + device=device, + resolution=int(image_resolution), + ) + if task_name == "v2t": + video_tokens = encode_video_tokens( + Path(args.video).expanduser().resolve(), + vq_model=vq_image, + device=device, + resolution=int(image_resolution), + num_frames=int(num_frames), + ) + if hasattr(vq_image, "cpu"): + vq_image = vq_image.cpu() + + if task_name == "s2t": + vq_audio = load_vq_audio_encoder(vq_audio_source, vq_audio_local_only, device) + audio_tokens = encode_audio_tokens(Path(args.audio).expanduser().resolve(), vq_audio) + if hasattr(vq_audio, "cpu"): + vq_audio = vq_audio.cpu() + + noise_schedule_params: dict[str, Any] = {} + try: + parsed = json.loads(args.noise_schedule_params) + if isinstance(parsed, dict): + noise_schedule_params = {str(k): v for k, v in parsed.items()} + except Exception: + noise_schedule_params = {} + + image_token_count = int(image_token_count) + if image_token_count <= 0: + if image_tokens is not None: + image_token_count = int(image_tokens.numel()) + else: + base_res = int(image_resolution) + image_token_count = max(1, (base_res // 16) ** 2) + + uncond_input_ids: list[int] | None = None + uncond_attention_mask: list[int] | None = None + if task_name == "t2t": + messages = [{"role": "user", "content": text}] + if getattr(tokenizer, "chat_template", None): + prompt_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + encoded = tokenizer(prompt_text, return_tensors="pt", add_special_tokens=False) + else: + encoded = tokenizer(text, return_tensors="pt", add_special_tokens=True) + prompt_token_ids = _to_1d_int_list(encoded["input_ids"]) + prompt_attention_mask = _to_1d_int_list(encoded.get("attention_mask")) + if not prompt_attention_mask: + prompt_attention_mask = [1] * len(prompt_token_ids) + else: + max_audio_len_for_prompt = int(max(t2s_token_length, 512)) + if audio_tokens is not None: + max_audio_len_for_prompt = max(max_audio_len_for_prompt, int(audio_tokens.numel())) + max_audio_len_short_for_prompt = max(256, max_audio_len_for_prompt // 2) + + uni_prompting = load_universal_prompting( + tokenizer=tokenizer, + tokenizer_source=tokenizer_source, + max_text_len=int(prompt_max_text_len), + cond_dropout_prob=float(args.cond_dropout_prob), + local_files_only=bool(tokenizer_local_only), + max_audio_len=int(max_audio_len_for_prompt), + max_audio_len_short=int(max_audio_len_short_for_prompt), + ) + prompting_text_vocab_size = int(len(uni_prompting.text_tokenizer)) + + is_mmu_task = task_name in {"i2t", "s2t", "v2t"} and not args.prompting_task.strip() + if is_mmu_task: + prompt_token_ids, prompt_attention_mask = make_mmu_prompt( + task=task_name, + text=text, + tokenizer=uni_prompting.text_tokenizer, + uni_prompting=uni_prompting, + image_tokens=image_tokens, + audio_tokens=audio_tokens, + video_tokens=video_tokens, + image_token_offset=prompting_text_vocab_size, + speech_token_offset=prompting_text_vocab_size + int(codebook_size), + ) + else: + prompt_payload, prompting_task = make_prompt_payload( + task=task_name, + text=text, + image_tokens=image_tokens, + audio_tokens=audio_tokens, + video_tokens=video_tokens, + image_placeholder_tokens=image_token_count, + audio_placeholder_tokens=int(t2s_token_length), + image_token_offset=text_vocab_size, + speech_token_offset=text_vocab_size + int(codebook_size), + mask_token_id=int(mask_token_id), + use_train_i2i_prompt=use_train_i2i_prompt, + ) + if args.prompting_task.strip(): + prompting_task = args.prompting_task.strip() + + prompt_token_ids, prompt_attention_mask = _run_uni_prompting( + uni_prompting, + prompt_payload, + prompting_task, + ) + + if task_name in {"i2t", "s2t", "v2t"}: + prompt_attention_mask = [1] * len(prompt_token_ids) + if not prompt_attention_mask: + prompt_attention_mask = [1] * len(prompt_token_ids) + + if task_name in {"t2i", "i2i"} and guidance_scale > 0: + uncond_payload, uncond_prompting_task = make_prompt_payload( + task=task_name, + text="", + image_tokens=image_tokens, + audio_tokens=audio_tokens, + video_tokens=video_tokens, + image_placeholder_tokens=image_token_count, + audio_placeholder_tokens=int(t2s_token_length), + image_token_offset=text_vocab_size, + speech_token_offset=text_vocab_size + int(codebook_size), + mask_token_id=int(mask_token_id), + use_train_i2i_prompt=use_train_i2i_prompt, + ) + uncond_input_ids, uncond_attention_mask = _run_uni_prompting( + uni_prompting, + uncond_payload, + args.prompting_task.strip() or uncond_prompting_task, + ) + if not uncond_attention_mask: + uncond_attention_mask = [1] * len(uncond_input_ids) + + runtime_info: dict[str, Any] = { + "task": [runtime_task], + "detok_id": [int(detok_id)], + DYNIN_PROMPT_SOURCE_KEY: [DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT], + "dynin_config_path": [str(dynin_config_path)], + "attention_mask": [prompt_attention_mask], + "prompt_max_text_len": [int(prompt_max_text_len)], + "prompting_max_text_len": [int(prompt_max_text_len)], + "cond_dropout_prob": [float(args.cond_dropout_prob)], + "prompting_cond_dropout_prob": [float(args.cond_dropout_prob)], + "tokenizer_path": [str(tokenizer_source)], + "text_vocab_size": [int(text_vocab_size)], + "model_local_files_only": [bool(model_local_only)], + "max_new_tokens": [int(max_new_tokens)], + "steps": [int(steps)], + "block_length": [int(block_length)], + "temperature": [float(temperature)], + "cfg_scale": [float(cfg_scale)], + "remasking": [str(remasking)], + "mask_id": [int(mask_token_id)], + "mask_token_id": [int(mask_token_id)], + "codebook_size": [int(codebook_size)], + "audio_codebook_size": [int(audio_codebook_size)], + "timesteps": [int(timesteps)], + "guidance_scale": [float(guidance_scale)], + "noise_type": [str(args.noise_type)], + "noise_schedule_name": [str(args.noise_schedule_name)], + "noise_schedule_params": [noise_schedule_params], + "seq_len": [int(image_token_count)], + "condition": [str(t2s_condition)], + "vq_model_image_path": [str(vq_image_source)], + "vq_model_image_local_files_only": [bool(vq_image_local_only)], + "vq_model_audio_path": [str(vq_audio_source)], + "vq_model_audio_local_files_only": [bool(vq_audio_local_only)], + } + + if task_name in {"t2t", "i2t", "s2t", "v2t"}: + runtime_info["prompt_length"] = [int(len(prompt_token_ids))] + if uncond_input_ids is not None: + runtime_info["uncond_input_ids"] = [uncond_input_ids] + if uncond_attention_mask is not None: + runtime_info["uncond_attention_mask"] = [uncond_attention_mask] + + if task_name == "t2s": + runtime_info["max_new_tokens"] = [int(t2s_token_length)] + + prompt = { + "prompt_token_ids": [int(v) for v in prompt_token_ids], + "additional_information": runtime_info, + "modalities": [final_modality], + } + + from vllm import SamplingParams + + from vllm_omni.entrypoints.omni import Omni + + stage_config_path = str(Path(args.stage_config_path).expanduser()) + omni = Omni(model=model_source, stage_configs_path=stage_config_path, dtype=args.dtype) + sampling_params_list = [ + SamplingParams(max_tokens=int(args.max_tokens_per_stage), temperature=0.0, top_p=1.0, detokenize=False) + for _ in range(omni.num_stages) + ] + + try: + outputs = list(omni.generate(prompt, sampling_params_list)) + finally: + omni.close() + + out_dir = Path(resolved_output_dir).expanduser() + out_dir.mkdir(parents=True, exist_ok=True) + stamp = time.strftime("%Y%m%d_%H%M%S") + prefix = args.output_prefix.strip() or f"{task_name}_{stamp}" + + if final_modality == "text": + text_out = extract_text_output(outputs, tokenizer=tokenizer) + if not text_out: + raise RuntimeError("No text output found.") + out_path = out_dir / f"{prefix}.txt" + out_path.write_text(text_out + "\n", encoding="utf-8") + print(f"[end2end] text saved: {out_path}") + print(text_out) + return + + if final_modality == "image": + image_out = extract_image_output(outputs) + if image_out is None: + raise RuntimeError("No image output found.") + pil = tensor_to_pil_image(image_out) + out_path = out_dir / f"{prefix}.png" + pil.save(out_path) + print(f"[end2end] image saved: {out_path}") + return + + if final_modality == "audio": + audio_out = extract_audio_output(outputs) + if audio_out is None: + raise RuntimeError("No audio output found.") + wav, sr = audio_out + out_path = out_dir / f"{prefix}.wav" + save_audio_wav(out_path, wav, sr) + print(f"[end2end] audio saved: {out_path} (sr={sr}, samples={wav.shape[0]})") + return + + raise RuntimeError(f"Unsupported final modality: {final_modality}") + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/dynin_omni/README.md b/examples/online_serving/dynin_omni/README.md new file mode 100644 index 000000000..d8526d423 --- /dev/null +++ b/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=:$PYTHONPATH +export HF_MODULES_CACHE= +``` + +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. diff --git a/examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py new file mode 100644 index 000000000..972855543 --- /dev/null +++ b/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() diff --git a/tests/e2e/offline_inference/test_dynin_omni.py b/tests/e2e/offline_inference/test_dynin_omni.py new file mode 100644 index 000000000..d17e7b817 --- /dev/null +++ b/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 diff --git a/tests/e2e/online_serving/test_dynin_omni_expansion.py b/tests/e2e/online_serving/test_dynin_omni_expansion.py new file mode 100644 index 000000000..4648c424f --- /dev/null +++ b/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) diff --git a/tests/e2e/stage_configs/dynin_omni_ci.yaml b/tests/e2e/stage_configs/dynin_omni_ci.yaml new file mode 100644 index 000000000..024000751 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/dynin_omni/__init__.py b/vllm_omni/model_executor/models/dynin_omni/__init__.py new file mode 100644 index 000000000..2a3bae8a9 --- /dev/null +++ b/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", +] diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni.py new file mode 100644 index 000000000..0caae158e --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py new file mode 100644 index 000000000..6166d8615 --- /dev/null +++ b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_common.py @@ -0,0 +1,1241 @@ +from __future__ import annotations + +import hashlib +import importlib.util +import os +import sys +import threading +import types +from collections.abc import Iterable +from dataclasses import dataclass +from enum import IntEnum +from functools import lru_cache +from pathlib import Path +from typing import Any + +import torch +from vllm.config import VllmConfig +from vllm.logger import init_logger + +logger = init_logger(__name__) + +try: + from huggingface_hub import snapshot_download +except Exception: # pragma: no cover + snapshot_download = None + + +class DetokTarget(IntEnum): + TEXT = 0 + AUDIO = 1 + IMAGE = 2 + + +TASK_TO_DETOK = { + "mmu": DetokTarget.TEXT, + "s2t": DetokTarget.TEXT, + "mmu_fast": DetokTarget.TEXT, + "mmu_fastdllm_v1": DetokTarget.TEXT, + "v2t": DetokTarget.TEXT, + "t2s": DetokTarget.AUDIO, + "t2s_mmu_like": DetokTarget.AUDIO, + "t2s_fixed": DetokTarget.AUDIO, + "s2s": DetokTarget.AUDIO, + "v2s": DetokTarget.AUDIO, + "t2i": DetokTarget.IMAGE, + "i2i": DetokTarget.IMAGE, + "ti2ti": DetokTarget.IMAGE, +} + +DEFAULT_VQ_IMAGE_SOURCE = "snu-aidas/magvitv2" +DEFAULT_VQ_AUDIO_SOURCE = "snu-aidas/emova_speech_tokenizer_vllm" +DEFAULT_MAGVIT_REMOTE_CODE_REPO = "snu-aidas/magvitv2" +DEFAULT_DYNIN_REMOTE_CODE_REPO = "snu-aidas/Dynin-Omni" +DYNIN_PROMPT_SOURCE_KEY = "dynin_prompt_source" +DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT = "offline_prebuilt" + +DYNIN_TASK_DEFAULT_RUNTIME = { + "t2t": ("mmu", "mmu", 0, "text"), + "t2i": ("t2i", "t2i_gen", 2, "image"), + "t2s": ("t2s_mmu_like", "t2s_gen", 1, "audio"), + "i2i": ("i2i", "i2i", 2, "image"), +} + +DYNIN_TASK_RUNTIME_FALLBACKS: dict[str, dict[str, Any]] = { + "t2t": { + "prompt_max_text_len": 1024, + "max_new_tokens": 1024, + "steps": 1024, + "block_length": 16, + "temperature": 0.0, + "cfg_scale": 0.0, + }, + "t2i": { + "prompt_max_text_len": 128, + "image_token_count": 1024, + "mask_token_id": 126336, + "codebook_size": 8192, + "timesteps": 20, + "guidance_scale": 3.5, + "temperature": 1.0, + }, + "i2i": { + "prompt_max_text_len": 128, + "mask_token_id": 126336, + "codebook_size": 8192, + "timesteps": 64, + "guidance_scale": 3.5, + "temperature": 1.0, + "image_resolution": 336, + "use_train_i2i_prompt": True, + }, + "t2s": { + "runtime_task": "t2s_mmu_like", + "prompting_task": "t2s_gen", + "prompt_max_text_len": 1024, + "t2s_token_length": 512, + "mask_token_id": 126336, + "codebook_size": 8192, + "audio_codebook_size": 4096, + "steps": 512, + "block_length": 128, + "temperature": 1.0, + "cfg_scale": 2.5, + "t2s_condition": "gender-female_emotion-neutral_speed-normal_pitch-normal", + }, +} + +DEFAULT_DYNIN_T2S_INSTRUCTION = "Please read the following text naturally." + +DYNIN_SPECIAL_TOKENS = ( + "<|soi|>", + "<|eoi|>", + "<|sov|>", + "<|eov|>", + "<|t2i|>", + "<|mmu|>", + "<|t2v|>", + "<|v2v|>", + "<|lvg|>", + "<|i2i|>", + "<|ti2ti|>", + "<|v2t|>", + "<|v2s|>", + "<|s2t|>", + "<|t2s|>", + "<|s2s|>", + "<|soa|>", + "<|eoa|>", +) + +_DYNIN_ONLINE_PROMPT_TOKEN_BY_TASK = { + "t2i": "<|t2i|>", + "i2i": "<|i2i|>", + "t2s": "<|t2s|>", +} + +_DYNIN_MODALITY_PLACEHOLDERS = ( + "<|soi|><|image|><|eoi|>", + "<|sov|><|video|><|eov|>", + "<|soa|><|audio|><|eoa|>", +) + +_DYNIN_CONFIG_CANDIDATE_RELPATHS = ( + "configs/dynin_omni.yaml", + "models/configs/dynin_omni.yaml", + "vllm_omni/model_executor/models/dynin_omni/configs/dynin_omni.yaml", + "vllm_omni/model_executor/stage_configs/dynin_omni.yaml", + "dynin_omni.yaml", +) + +_DYNIN_REMOTE_ALLOW_PATTERNS = ("*.py", "*.json", "*.yaml", "*.yml") + +_DYNIN_REMOTE_CACHE_LOCK = threading.Lock() +_DYNIN_REMOTE_PACKAGE_BY_SNAPSHOT: dict[str, str] = {} +_DYNIN_REMOTE_ATTR_CACHE: dict[tuple[str, str, str, str | None, bool], Any] = {} + + +@dataclass(frozen=True) +class DyninInferSources: + model_source: str + tokenizer_source: str + vq_image_source: str + vq_audio_source: str + model_local_files_only: bool + vq_image_local_files_only: bool + vq_audio_local_files_only: bool + config_path: str | None = None + + @property + def local_files_only(self) -> bool: + return self.model_local_files_only + + +@dataclass(frozen=True) +class RemoteCodeSettings: + default_repo: str + repo_env: str + revision_env: str + local_only_env: str + + +DYNIN_REMOTE_SETTINGS = RemoteCodeSettings( + default_repo=DEFAULT_DYNIN_REMOTE_CODE_REPO, + repo_env="DYNIN_REMOTE_CODE_REPO_ID", + revision_env="DYNIN_REMOTE_CODE_REVISION", + local_only_env="DYNIN_REMOTE_CODE_LOCAL_FILES_ONLY", +) + +MAGVIT_REMOTE_SETTINGS = RemoteCodeSettings( + default_repo=DEFAULT_MAGVIT_REMOTE_CODE_REPO, + repo_env="DYNIN_MAGVIT_REMOTE_CODE_REPO_ID", + revision_env="DYNIN_MAGVIT_REMOTE_CODE_REVISION", + local_only_env="DYNIN_MAGVIT_REMOTE_CODE_LOCAL_FILES_ONLY", +) + + +def unwrap_first_value(value: Any, default: Any = None) -> Any: + if value is None: + return default + if isinstance(value, list): + return default if not value else value[0] + if isinstance(value, torch.Tensor): + if value.numel() == 0: + return default + if value.numel() == 1: + return value.item() + return value + return value + + +def normalize_runtime_info(runtime_additional_information: Any) -> dict[str, Any]: + if isinstance(runtime_additional_information, list): + if not runtime_additional_information: + return {} + first = runtime_additional_information[0] + return first if isinstance(first, dict) else {} + if isinstance(runtime_additional_information, dict): + return runtime_additional_information + return {} + + +def logical_dynin_task(task: Any) -> str: + task_text = str(unwrap_first_value(task, "") or "").strip().lower() + if task_text in ("t2s", "t2s_mmu_like", "t2s_fixed"): + return "t2s" + if task_text in ("t2i", "i2i"): + return task_text + return "t2t" + + +def dynin_runtime_fallback(task: str, key: str, value: Any = None) -> Any: + if isinstance(value, str): + if value.strip() != "": + return value + elif value is not None: + return value + return DYNIN_TASK_RUNTIME_FALLBACKS.get(task, {}).get(key) + + +def coerce_token_ids_1d( + value: Any, + ref_device: torch.device | None = None, +) -> torch.Tensor: + if isinstance(value, tuple): + value = value[0] + + if isinstance(value, list): + if not value: + device = ref_device or torch.device("cpu") + return torch.empty(0, dtype=torch.long, device=device) + if isinstance(value[0], torch.Tensor): + value = value[0] + else: + value = torch.tensor( + value[0] if isinstance(value[0], list) else value, + dtype=torch.long, + ) + + if not isinstance(value, torch.Tensor): + value = torch.tensor(value, dtype=torch.long) + + if value.ndim == 0: + value = value.unsqueeze(0) + if value.ndim > 1: + value = value[0] + + if ref_device is not None and value.device != ref_device: + value = value.to(ref_device) + + return value.to(dtype=torch.long).contiguous() + + +def _first_positive_int(value: Any) -> int | None: + if value is None: + return None + if isinstance(value, torch.Tensor): + if value.numel() != 1: + return None + value = value.item() + try: + value = int(value) + except (TypeError, ValueError): + return None + return value if value > 0 else None + + +def resolve_hidden_size( + *, + vllm_config: VllmConfig, + model: Any | None = None, + default: int = 1024, +) -> int: + if model is not None: + try: + embeddings = model.get_input_embeddings() + weight = getattr(embeddings, "weight", None) + if isinstance(weight, torch.Tensor) and weight.ndim >= 2: + hidden_size = _first_positive_int(weight.shape[-1]) + if hidden_size is not None: + return hidden_size + except Exception: + pass + + model_cfg = getattr(model, "config", None) + for key in ("hidden_size", "d_model", "n_embd", "dim", "model_dim", "embed_dim"): + hidden_size = _first_positive_int(getattr(model_cfg, key, None)) + if hidden_size is not None: + return hidden_size + + for config_obj in ( + getattr(vllm_config.model_config, "hf_config", None), + getattr(vllm_config.model_config, "hf_text_config", None), + ): + if config_obj is None: + continue + for key in ("hidden_size", "d_model", "n_embd", "dim", "model_dim", "embed_dim"): + value = config_obj.get(key) if isinstance(config_obj, dict) else getattr(config_obj, key, None) + hidden_size = _first_positive_int(value) + if hidden_size is not None: + return hidden_size + + return default + + +def build_zero_input_embeddings( + *, + input_ids: torch.Tensor, + hidden_size: int, + stage_name: str, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + if input_ids.ndim == 0: + shape = (1, hidden_size) + elif input_ids.ndim == 1: + shape = (input_ids.shape[0], hidden_size) + elif input_ids.ndim == 2: + shape = (input_ids.shape[0], input_ids.shape[1], hidden_size) + else: + raise ValueError(f"Unsupported input_ids rank for {stage_name}: {input_ids.ndim}") + return torch.zeros(shape, dtype=dtype, device=input_ids.device) + + +def _to_bool(value: Any, default: bool = False) -> bool: + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + + text = str(value).strip().lower() + if text in ("1", "true", "yes", "y", "on"): + return True + if text in ("0", "false", "no", "n", "off", "", "none", "null"): + return False + return default + + +def _runtime_value(runtime_info: dict[str, Any], key: str) -> Any: + return unwrap_first_value(runtime_info.get(key), None) + + +def _runtime_first_value(runtime_info: dict[str, Any], keys: tuple[str, ...]) -> Any: + for key in keys: + value = _runtime_value(runtime_info, key) + if value is not None: + return value + return None + + +def _node_value(node: Any, key: str, default: Any = None) -> Any: + if node is None: + return default + if isinstance(node, dict): + return node.get(key, default) + try: + return node.get(key, default) + except Exception: + return getattr(node, key, default) + + +def _looks_like_hf_repo_id(value: str | None) -> bool: + if not isinstance(value, str): + return False + if value.count("/") != 1: + return False + org, name = value.split("/", 1) + return bool(org and name) + + +def _find_dynin_config_under_root(root: Path) -> Path | None: + for rel_path in _DYNIN_CONFIG_CANDIDATE_RELPATHS: + candidate = root.expanduser() / rel_path + if candidate.exists(): + return candidate.resolve() + return None + + +@lru_cache(maxsize=16) +def _resolve_dynin_config_from_hf_repo(repo_id: str) -> str | None: + if not _looks_like_hf_repo_id(repo_id) or snapshot_download is None: + return None + + try: + snapshot_dir = ( + Path( + snapshot_download( + repo_id=repo_id, + repo_type="model", + allow_patterns=list(_DYNIN_CONFIG_CANDIDATE_RELPATHS), + local_files_only=True, + ) + ) + .expanduser() + .resolve() + ) + except Exception: + return None + + found = _find_dynin_config_under_root(snapshot_dir) + return str(found) if found is not None else None + + +def _resolve_existing_path(path_like: Any, source_name: str) -> str | None: + if path_like is None: + return None + text = str(path_like).strip() + if not text: + return None + + path = Path(text).expanduser() + if path.is_file(): + return str(path.resolve()) + + logger.warning( + "DYNIN config path from %s does not exist: %s. Falling back to auto-discovery.", + source_name, + path, + ) + return None + + +def _resolve_config_path(vllm_config: VllmConfig, runtime_info: dict[str, Any]) -> str | None: + for value, name in ( + (_runtime_value(runtime_info, "dynin_config_path"), "runtime_info.dynin_config_path"), + (os.getenv("DYNIN_CONFIG_PATH"), "DYNIN_CONFIG_PATH"), + (getattr(vllm_config.model_config, "dynin_config_path", None), "vllm_config.model_config.dynin_config_path"), + ): + resolved = _resolve_existing_path(value, name) + if resolved: + return resolved + + model_source = str(getattr(vllm_config.model_config, "model", "") or "") + tokenizer_source = str(getattr(vllm_config.model_config, "tokenizer", "") or "") + hf_config = getattr(vllm_config.model_config, "hf_config", None) + hf_name_or_path = ( + hf_config.get("_name_or_path") if isinstance(hf_config, dict) else getattr(hf_config, "_name_or_path", None) + ) + + hf_repo_candidates: list[str] = [] + for source in (model_source, tokenizer_source, hf_name_or_path): + if not _looks_like_hf_repo_id(source): + continue + source = str(source) + if source not in hf_repo_candidates: + hf_repo_candidates.append(source) + + for source in hf_repo_candidates: + resolved = _resolve_dynin_config_from_hf_repo(source) + if resolved is not None: + logger.info("Resolved dynin config from Hugging Face cache for %s: %s", source, resolved) + return resolved + + for source in (model_source, tokenizer_source): + source_path = Path(source).expanduser() + if source_path.is_dir(): + found = _find_dynin_config_under_root(source_path) + if found is not None: + return str(found) + + module_root = Path(__file__).resolve().parent + for bundled in ( + module_root / "configs" / "dynin_omni.yaml", + module_root / "models" / "configs" / "dynin_omni.yaml", + module_root.parent / "stage_configs" / "dynin_omni.yaml", + ): + if bundled.exists(): + return str(bundled) + + return None + + +@lru_cache(maxsize=16) +def _load_omega_config(config_path: str) -> Any: + try: + from omegaconf import OmegaConf + except ImportError as e: + raise ImportError( + f"omegaconf is required to load Dynin config files. Install it to read config: {config_path}" + ) from e + return OmegaConf.load(config_path) + + +def resolve_dynin_infer_sources( + *, + vllm_config: VllmConfig, + runtime_info: dict[str, Any] | None = None, +) -> DyninInferSources: + runtime_info = runtime_info or {} + + base_model_source = str(getattr(vllm_config.model_config, "model", "")) + base_model_path = Path(base_model_source).expanduser() + local_vllm_model_source = str(base_model_path) if base_model_path.is_dir() else None + + model_source = base_model_source + tokenizer_source = model_source + vq_image_source = DEFAULT_VQ_IMAGE_SOURCE + vq_audio_source = DEFAULT_VQ_AUDIO_SOURCE + model_local_files_only = False + vq_image_local_files_only = False + vq_audio_local_files_only = False + + resolver_source: str | None = base_model_source if base_model_source else None + resolver_local_files_only: bool | None = True if base_model_path.is_dir() else None + resolve_model_pretrained_source_fn = get_dynin_config_resolver_attr( + "resolve_model_pretrained_source", + source=resolver_source, + local_files_only=resolver_local_files_only, + ) + resolve_tokenizer_source_fn = get_dynin_config_resolver_attr( + "resolve_tokenizer_source", + source=resolver_source, + local_files_only=resolver_local_files_only, + ) + resolve_model_local_files_only_fn = get_dynin_config_resolver_attr( + "resolve_model_local_files_only", + source=resolver_source, + local_files_only=resolver_local_files_only, + ) + resolve_vq_cfg_block_fn = get_dynin_config_resolver_attr( + "resolve_vq_cfg_block", + source=resolver_source, + local_files_only=resolver_local_files_only, + ) + resolve_vq_repo_source_fn = get_dynin_config_resolver_attr( + "resolve_vq_repo_source", + source=resolver_source, + local_files_only=resolver_local_files_only, + ) + + config_path = _resolve_config_path(vllm_config, runtime_info) + if config_path: + config_file = Path(config_path).expanduser() + if config_file.exists(): + try: + dynin_cfg = _load_omega_config(str(config_file)) + model_source = resolve_model_pretrained_source_fn( + dynin_cfg, + default=model_source, + ) + tokenizer_source = resolve_tokenizer_source_fn( + dynin_cfg, + default=tokenizer_source, + ) + model_local_files_only = resolve_model_local_files_only_fn( + dynin_cfg, + default=model_local_files_only, + ) + vq_image_cfg = resolve_vq_cfg_block_fn(dynin_cfg, modality="image") + vq_audio_cfg = resolve_vq_cfg_block_fn(dynin_cfg, modality="audio") + vq_image_source = resolve_vq_repo_source_fn( + vq_image_cfg, + default=vq_image_source, + ) + vq_audio_source = resolve_vq_repo_source_fn( + vq_audio_cfg, + default=vq_audio_source, + ) + vq_image_local_files_only = _to_bool( + _node_value(vq_image_cfg, "local_files_only", None), + default=model_local_files_only, + ) + vq_audio_local_files_only = _to_bool( + _node_value(vq_audio_cfg, "local_files_only", None), + default=model_local_files_only, + ) + except Exception as e: + logger.warning( + "Failed to resolve DYNIN inference config from %s: %s", + config_file, + e, + ) + else: + logger.warning("DYNIN config path does not exist: %s", config_file) + + runtime_model_source = _runtime_value(runtime_info, "dynin_model_path") + if runtime_model_source: + model_source = str(runtime_model_source) + + runtime_tokenizer_source = _runtime_value(runtime_info, "tokenizer_path") + if runtime_tokenizer_source: + tokenizer_source = str(runtime_tokenizer_source) + + runtime_vq_image_source = _runtime_value(runtime_info, "vq_model_image_path") + if runtime_vq_image_source is None: + runtime_vq_image_source = _runtime_value(runtime_info, "vq_model_path_image") + if runtime_vq_image_source: + vq_image_source = str(runtime_vq_image_source) + + runtime_vq_audio_source = _runtime_value(runtime_info, "vq_model_audio_path") + if runtime_vq_audio_source is None: + runtime_vq_audio_source = _runtime_value(runtime_info, "vq_model_path_audio") + if runtime_vq_audio_source: + vq_audio_source = str(runtime_vq_audio_source) + + runtime_local_global = _runtime_value(runtime_info, "local_files_only") + runtime_local_model = _runtime_first_value( + runtime_info, + ("model_local_files_only", "local_files_only_model"), + ) + runtime_local_vq_image = _runtime_first_value( + runtime_info, + ("vq_model_image_local_files_only", "local_files_only_vq_image"), + ) + runtime_local_vq_audio = _runtime_first_value( + runtime_info, + ("vq_model_audio_local_files_only", "local_files_only_vq_audio"), + ) + + if runtime_local_global is not None: + global_local = _to_bool(runtime_local_global, default=False) + if runtime_local_model is None: + model_local_files_only = global_local + if runtime_local_vq_image is None: + vq_image_local_files_only = global_local + if runtime_local_vq_audio is None: + vq_audio_local_files_only = global_local + + if runtime_local_model is not None: + model_local_files_only = _to_bool( + runtime_local_model, + default=model_local_files_only, + ) + if runtime_local_vq_image is not None: + vq_image_local_files_only = _to_bool( + runtime_local_vq_image, + default=vq_image_local_files_only, + ) + if runtime_local_vq_audio is not None: + vq_audio_local_files_only = _to_bool( + runtime_local_vq_audio, + default=vq_audio_local_files_only, + ) + + if runtime_local_global is None and runtime_local_model is None and local_vllm_model_source is not None: + model_local_files_only = True + + if local_vllm_model_source is not None: + if not runtime_model_source: + if model_source != local_vllm_model_source: + logger.info( + "DYNIN infer model source overridden to local vLLM model path: %s (from %s)", + local_vllm_model_source, + model_source, + ) + model_source = local_vllm_model_source + if not runtime_tokenizer_source: + tokenizer_source = local_vllm_model_source + + return DyninInferSources( + model_source=model_source, + tokenizer_source=tokenizer_source, + vq_image_source=vq_image_source, + vq_audio_source=vq_audio_source, + model_local_files_only=model_local_files_only, + vq_image_local_files_only=vq_image_local_files_only, + vq_audio_local_files_only=vq_audio_local_files_only, + config_path=config_path, + ) + + +def _resolve_remote_source(source: str | None, settings: RemoteCodeSettings) -> str: + if isinstance(source, str): + stripped = source.strip() + if stripped: + source_path = Path(stripped).expanduser() + if source_path.is_dir(): + return str(source_path.resolve()) + if _looks_like_hf_repo_id(stripped): + return stripped + + env_repo = os.getenv(settings.repo_env) + if _looks_like_hf_repo_id(env_repo): + return str(env_repo).strip() + + return settings.default_repo + + +def _resolve_remote_revision(revision: str | None, settings: RemoteCodeSettings) -> str | None: + if isinstance(revision, str) and revision.strip(): + return revision.strip() + env_revision = os.getenv(settings.revision_env) + if isinstance(env_revision, str) and env_revision.strip(): + return env_revision.strip() + return None + + +def _resolve_remote_local_only(local_files_only: bool | None, settings: RemoteCodeSettings) -> bool: + if local_files_only is not None: + return bool(local_files_only) + return _to_bool(os.getenv(settings.local_only_env), default=False) + + +def _resolve_remote_snapshot_dir( + *, + source: str, + revision: str | None, + local_files_only: bool, +) -> str: + source_path = Path(source).expanduser() + if source_path.is_dir(): + return str(source_path.resolve()) + + if snapshot_download is None: + raise RuntimeError("huggingface_hub is required to load remote code.") + + kwargs: dict[str, Any] = { + "repo_id": source, + "repo_type": "model", + "allow_patterns": list(_DYNIN_REMOTE_ALLOW_PATTERNS), + "local_files_only": bool(local_files_only), + } + if revision is not None: + kwargs["revision"] = revision + + try: + return str(snapshot_download(**kwargs)) + except TypeError: + kwargs.pop("local_files_only", None) + return str(snapshot_download(**kwargs)) + + +def _ensure_remote_package(snapshot_dir: str) -> str: + with _DYNIN_REMOTE_CACHE_LOCK: + existing = _DYNIN_REMOTE_PACKAGE_BY_SNAPSHOT.get(snapshot_dir) + if existing is not None: + return existing + + digest = hashlib.sha1(snapshot_dir.encode("utf-8")).hexdigest()[:12] + package_name = f"_dynin_hf_remote_{digest}" + + package = types.ModuleType(package_name) + package.__path__ = [snapshot_dir] # type: ignore[attr-defined] + package.__file__ = str(Path(snapshot_dir) / "__init__.py") + + sys.modules.setdefault(package_name, package) + _DYNIN_REMOTE_PACKAGE_BY_SNAPSHOT[snapshot_dir] = package_name + return package_name + + +def _load_remote_module( + *, + module_name: str, + source: str, + revision: str | None, + local_files_only: bool, +): + snapshot_dir = _resolve_remote_snapshot_dir( + source=source, + revision=revision, + local_files_only=local_files_only, + ) + + module_path = Path(snapshot_dir) / f"{module_name}.py" + if not module_path.is_file(): + raise ImportError(f"Remote code module '{module_name}.py' not found under '{snapshot_dir}'. source={source!r}") + + package_name = _ensure_remote_package(snapshot_dir) + full_name = f"{package_name}.{module_name}" + + existing = sys.modules.get(full_name) + if existing is not None: + return existing + + spec = importlib.util.spec_from_file_location(full_name, module_path) + if spec is None or spec.loader is None: + raise ImportError(f"Failed to create import spec for '{module_path}'.") + + module = importlib.util.module_from_spec(spec) + module.__package__ = package_name + sys.modules[full_name] = module + try: + spec.loader.exec_module(module) + except Exception: + sys.modules.pop(full_name, None) + raise + return module + + +def resolve_remote_attr( + attr_name: str, + *, + module_name: str, + settings: RemoteCodeSettings, + source: str | None = None, + revision: str | None = None, + local_files_only: bool | None = None, + fallback_module_names: Iterable[str] = (), + optional: bool = False, +) -> Any | None: + resolved_source = _resolve_remote_source(source, settings) + resolved_revision = _resolve_remote_revision(revision, settings) + resolved_local_only = _resolve_remote_local_only(local_files_only, settings) + + module_candidates = [module_name, *[m for m in fallback_module_names if m and m != module_name]] + last_error: Exception | None = None + + for candidate in module_candidates: + cache_key = (attr_name, candidate, resolved_source, resolved_revision, resolved_local_only) + cached = _DYNIN_REMOTE_ATTR_CACHE.get(cache_key) + if cached is not None: + return cached + + try: + module = _load_remote_module( + module_name=candidate, + source=resolved_source, + revision=resolved_revision, + local_files_only=resolved_local_only, + ) + if hasattr(module, attr_name): + value = getattr(module, attr_name) + _DYNIN_REMOTE_ATTR_CACHE[cache_key] = value + return value + except Exception as e: + last_error = e + + if optional: + if last_error is not None: + logger.debug( + "Optional remote attr not found: attr=%s source=%s revision=%s err=%s", + attr_name, + resolved_source, + resolved_revision, + last_error, + ) + return None + + raise ImportError( + f"Failed to resolve '{attr_name}' from remote code " + f"(source={resolved_source!r}, revision={resolved_revision!r}, modules={module_candidates})." + ) from last_error + + +_DYNIN_MODELING_REMOTE_EXPORTS = { + "DyninOmniConfig": "DyninOmniConfig", + "DyninOmniModelLM": "DyninOmniModelLM", + "VideoTokenMerger": "VideoTokenMerger", +} + +_DYNIN_SAMPLING_REMOTE_EXPORTS = { + "log": "log", + "gumbel_noise": "gumbel_noise", + "gumbel_sample": "gumbel_sample", + "top_k": "top_k", + "mask_by_random_topk": "mask_by_random_topk", + "cosine_schedule": "cosine_schedule", + "linear_schedule": "linear_schedule", + "pow": "pow", + "sigmoid_schedule": "sigmoid_schedule", + "get_mask_schedule": "get_mask_schedule", + "top_k_top_p_filtering": "top_k_top_p_filtering", +} + +_DYNIN_CONFIG_RESOLVER_REMOTE_EXPORTS = { + "resolve_model_pretrained_source": "resolve_model_pretrained_source", + "resolve_tokenizer_source": "resolve_tokenizer_source", + "resolve_model_local_files_only": "resolve_model_local_files_only", + "resolve_vq_cfg_block": "resolve_vq_cfg_block", + "resolve_vq_repo_source": "resolve_vq_repo_source", +} + +_DYNIN_MAGVIT_REMOTE_EXPORTS = { + "VQGANEncoder": "VQGANEncoder", + "VQGANDecoder": "VQGANDecoder", + "LFQuantizer": "LFQuantizer", + "MAGVITv2": "MAGVITv2", +} + + +def _get_export_attr( + name: str, + export_map: dict[str, str], + *, + module_name: str, + settings: RemoteCodeSettings, + source: str | None = None, + revision: str | None = None, + local_files_only: bool | None = None, + optional: bool = False, +) -> Any | None: + attr_name = export_map.get(name) + if attr_name is None: + raise AttributeError(f"Unsupported export: {name!r}") + + return resolve_remote_attr( + attr_name, + module_name=module_name, + settings=settings, + source=source, + revision=revision, + local_files_only=local_files_only, + optional=optional, + ) + + +def get_dynin_modeling_attr(name: str) -> Any: + return _get_export_attr( + name, + _DYNIN_MODELING_REMOTE_EXPORTS, + module_name="modeling_dynin_omni", + settings=DYNIN_REMOTE_SETTINGS, + ) + + +def get_dynin_sampling_attr(name: str) -> Any: + return _get_export_attr( + name, + _DYNIN_SAMPLING_REMOTE_EXPORTS, + module_name="sampling", + settings=DYNIN_REMOTE_SETTINGS, + ) + + +def get_dynin_config_resolver_attr( + name: str, + *, + source: str | None = None, + revision: str | None = None, + local_files_only: bool | None = None, +) -> Any: + attr_name = _DYNIN_CONFIG_RESOLVER_REMOTE_EXPORTS.get(name) + if attr_name is None: + raise AttributeError(f"Unsupported Dynin config_resolver export: {name!r}") + + if source is not None: + value = resolve_remote_attr( + attr_name, + module_name="config_resolver", + settings=DYNIN_REMOTE_SETTINGS, + source=source, + revision=revision, + local_files_only=local_files_only, + optional=True, + ) + if value is not None: + return value + + return resolve_remote_attr( + attr_name, + module_name="config_resolver", + settings=DYNIN_REMOTE_SETTINGS, + source=DEFAULT_DYNIN_REMOTE_CODE_REPO, + revision=revision, + local_files_only=local_files_only, + optional=False, + ) + + +def get_dynin_magvit_attr( + name: str, + *, + source: str | None = None, + revision: str | None = None, + local_files_only: bool | None = None, +) -> Any: + attr_name = _DYNIN_MAGVIT_REMOTE_EXPORTS.get(name) + if attr_name is None: + raise AttributeError(f"Unsupported Dynin MAGVIT export: {name!r}") + + value = resolve_remote_attr( + attr_name, + module_name="modeling_magvitv2", + settings=MAGVIT_REMOTE_SETTINGS, + source=source, + revision=revision, + local_files_only=local_files_only, + optional=True, + ) + if value is not None: + return value + + resolved_source = _resolve_remote_source(source, MAGVIT_REMOTE_SETTINGS) + resolved_revision = _resolve_remote_revision(revision, MAGVIT_REMOTE_SETTINGS) + resolved_local_only = _resolve_remote_local_only(local_files_only, MAGVIT_REMOTE_SETTINGS) + + if resolved_source != DEFAULT_MAGVIT_REMOTE_CODE_REPO: + return resolve_remote_attr( + attr_name, + module_name="modeling_magvitv2", + settings=MAGVIT_REMOTE_SETTINGS, + source=DEFAULT_MAGVIT_REMOTE_CODE_REPO, + revision=resolved_revision, + local_files_only=resolved_local_only, + optional=False, + ) + + raise ImportError( + f"Failed to resolve MAGVIT attr '{attr_name}' from source={resolved_source!r} (revision={resolved_revision!r})." + ) + + +def build_dynin_chat_prompt(content: str) -> str: + return ( + f"<|start_header_id|>user<|end_header_id|>\n{content}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" + ) + + +def extract_dynin_user_prompt_text(decoded_prompt: str) -> str: + text = str(decoded_prompt or "") + assistant_marker = "<|start_header_id|>assistant<|end_header_id|>" + user_marker = "<|start_header_id|>user<|end_header_id|>" + end_header_marker = "<|end_header_id|>" + eot_marker = "<|eot_id|>" + + if assistant_marker in text: + text = text.rsplit(assistant_marker, 1)[0] + if eot_marker in text: + text = text.rsplit(eot_marker, 1)[0] + if user_marker in text: + text = text.rsplit(user_marker, 1)[-1] + if end_header_marker in text: + text = text.split(end_header_marker, 1)[-1] + return text.strip() + + +def normalize_dynin_online_prompt_text(task: str, decoded_prompt: str) -> str: + text = extract_dynin_user_prompt_text(decoded_prompt) + if not text: + text = str(decoded_prompt or "") + + for placeholder in _DYNIN_MODALITY_PLACEHOLDERS: + text = text.replace(placeholder, " ") + + task_token = _DYNIN_ONLINE_PROMPT_TOKEN_BY_TASK.get(task) + if task_token: + text = text.replace(task_token, " ", 1) + + text = " ".join(text.split()).strip() + + if task == "t2s": + if not text: + text = "Hello. This is a default text-to-speech sample." + text = build_dynin_chat_prompt(f"{DEFAULT_DYNIN_T2S_INSTRUCTION}\n{text}") + elif task in {"t2i", "i2i"} and not text: + text = "A high quality detailed image." + + return text + + +def infer_dynin_online_task( + *, + decoded_prompt: str, + has_image: bool = False, + has_audio: bool = False, + has_video: bool = False, +) -> str: + prompt = str(decoded_prompt or "") + if "<|i2i|>" in prompt: + return "i2i" + if "<|t2i|>" in prompt and not has_audio and not has_video: + return "t2i" + if "<|t2s|>" in prompt and not has_audio and not has_video: + return "t2s" + return "t2t" + + +def build_dynin_prompt_payload( + *, + task: str, + text: str, + image_tokens: torch.Tensor | None, + image_placeholder_tokens: int, + audio_placeholder_tokens: int, + image_token_offset: int, + mask_token_id: int, + use_train_i2i_prompt: bool, +) -> tuple[Any, str]: + _, prompting_task, _, _ = DYNIN_TASK_DEFAULT_RUNTIME[task] + + if task == "t2t": + payload = ([[]], [build_dynin_chat_prompt(text)]) + return payload, prompting_task + + if task == "t2i": + image_placeholder = torch.full( + (1, int(image_placeholder_tokens)), + fill_value=int(mask_token_id), + dtype=torch.long, + ) + payload = ([text], image_placeholder) + return payload, prompting_task + + if task == "i2i": + if image_tokens is None: + raise ValueError("i2i requires image tokens") + src = image_tokens.view(1, -1).long() + int(image_token_offset) + target_len = int(image_placeholder_tokens) if image_placeholder_tokens > 0 else int(src.shape[1]) + image_placeholder = torch.full( + (1, target_len), + fill_value=int(mask_token_id), + dtype=torch.long, + ) + if use_train_i2i_prompt: + labels_placeholder = torch.full( + (1, target_len), + fill_value=-100, + dtype=torch.long, + ) + payload = ([text], src, image_placeholder, labels_placeholder) + return payload, "i2i" + payload = ([text], src, image_placeholder) + return payload, "i2i_gen" + + if task == "t2s": + audio_placeholder = torch.full( + (1, int(audio_placeholder_tokens)), + fill_value=int(mask_token_id), + dtype=torch.long, + ) + payload = ([text], audio_placeholder) + return payload, prompting_task + + raise ValueError(f"Unsupported Dynin online bootstrap task: {task}") + + +def _wrap_runtime_field(value: Any) -> list[Any]: + return [value] + + +def build_dynin_online_runtime_info( + *, + task: str, + text_vocab_size: int, + infer_sources: DyninInferSources, + dynin_config_path: str | None = None, + prompting_input: Any | None = None, + attention_mask: list[int] | None = None, + prompt_length: int | None = None, + uncond_prompting_input: Any | None = None, + image_token_count: int = 0, + t2s_token_length: int | None = None, + use_train_i2i_prompt: bool | None = None, +) -> dict[str, Any]: + runtime_task, prompting_task, detok_id, _ = DYNIN_TASK_DEFAULT_RUNTIME[task] + + prompt_max_text_len = int(dynin_runtime_fallback(task, "prompt_max_text_len", None) or 1024) + max_new_tokens = int(dynin_runtime_fallback(task, "max_new_tokens", None) or 256) + steps = int(dynin_runtime_fallback(task, "steps", None) or 256) + block_length = int(dynin_runtime_fallback(task, "block_length", None) or 2) + temperature = float(dynin_runtime_fallback(task, "temperature", None) or 0.0) + cfg_scale = float(dynin_runtime_fallback(task, "cfg_scale", None) or 0.0) + remasking = str(dynin_runtime_fallback(task, "remasking", None) or "low_confidence") + timesteps = int(dynin_runtime_fallback(task, "timesteps", None) or 20) + guidance_scale = float(dynin_runtime_fallback(task, "guidance_scale", None) or 0.0) + mask_token_id = int(dynin_runtime_fallback(task, "mask_token_id", None) or 126336) + codebook_size = int(dynin_runtime_fallback(task, "codebook_size", None) or 8192) + audio_codebook_size = int(dynin_runtime_fallback(task, "audio_codebook_size", None) or 4096) + image_resolution = int(dynin_runtime_fallback(task, "image_resolution", None) or 336) + if image_token_count <= 0 and task in {"t2i", "i2i"}: + fallback_count = dynin_runtime_fallback(task, "image_token_count", None) + if fallback_count is not None: + image_token_count = int(fallback_count) + else: + image_token_count = max(1, (image_resolution // 16) ** 2) + + if t2s_token_length is None: + t2s_token_length = int(dynin_runtime_fallback(task, "t2s_token_length", None) or 383) + t2s_condition = str( + dynin_runtime_fallback( + task, + "t2s_condition", + None, + ) + or "gender-female_emotion-neutral_speed-normal_pitch-normal" + ) + if use_train_i2i_prompt is None: + use_train_i2i_prompt = bool(dynin_runtime_fallback(task, "use_train_i2i_prompt", task == "i2i")) + + runtime_info: dict[str, Any] = { + "task": _wrap_runtime_field(runtime_task), + "prompting_task": _wrap_runtime_field(prompting_task), + "detok_id": _wrap_runtime_field(int(detok_id)), + "prompt_max_text_len": _wrap_runtime_field(prompt_max_text_len), + "prompting_max_text_len": _wrap_runtime_field(prompt_max_text_len), + "cond_dropout_prob": _wrap_runtime_field(0.0), + "prompting_cond_dropout_prob": _wrap_runtime_field(0.0), + "tokenizer_path": _wrap_runtime_field(str(infer_sources.tokenizer_source)), + "text_vocab_size": _wrap_runtime_field(int(text_vocab_size)), + "model_local_files_only": _wrap_runtime_field(bool(infer_sources.model_local_files_only)), + "max_new_tokens": _wrap_runtime_field(int(t2s_token_length if task == "t2s" else max_new_tokens)), + "steps": _wrap_runtime_field(steps), + "block_length": _wrap_runtime_field(block_length), + "temperature": _wrap_runtime_field(temperature), + "cfg_scale": _wrap_runtime_field(cfg_scale), + "remasking": _wrap_runtime_field(remasking), + "mask_id": _wrap_runtime_field(mask_token_id), + "mask_token_id": _wrap_runtime_field(mask_token_id), + "codebook_size": _wrap_runtime_field(codebook_size), + "audio_codebook_size": _wrap_runtime_field(audio_codebook_size), + "timesteps": _wrap_runtime_field(timesteps), + "guidance_scale": _wrap_runtime_field(guidance_scale), + "noise_type": _wrap_runtime_field("mask"), + "noise_schedule_name": _wrap_runtime_field("cosine"), + "noise_schedule_params": _wrap_runtime_field({}), + "seq_len": _wrap_runtime_field(int(image_token_count)), + "condition": _wrap_runtime_field(t2s_condition), + "t2s_condition": _wrap_runtime_field(t2s_condition), + "vq_model_image_path": _wrap_runtime_field(str(infer_sources.vq_image_source)), + "vq_model_image_local_files_only": _wrap_runtime_field(bool(infer_sources.vq_image_local_files_only)), + "vq_model_audio_path": _wrap_runtime_field(str(infer_sources.vq_audio_source)), + "vq_model_audio_local_files_only": _wrap_runtime_field(bool(infer_sources.vq_audio_local_files_only)), + "image_resolution": _wrap_runtime_field(image_resolution), + "t2s_token_length": _wrap_runtime_field(int(t2s_token_length)), + "use_train_i2i_prompt": _wrap_runtime_field(bool(use_train_i2i_prompt)), + } + + if dynin_config_path: + runtime_info["dynin_config_path"] = _wrap_runtime_field(str(dynin_config_path)) + if prompting_input is not None: + runtime_info["prompting_input"] = _wrap_runtime_field(prompting_input) + if uncond_prompting_input is not None: + runtime_info["uncond_prompting_input"] = _wrap_runtime_field(uncond_prompting_input) + if attention_mask: + runtime_info["attention_mask"] = _wrap_runtime_field(list(attention_mask)) + if prompt_length is None and attention_mask: + prompt_length = len(attention_mask) + if prompt_length is not None: + runtime_info["prompt_length"] = _wrap_runtime_field(int(prompt_length)) + + return runtime_info diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2audio.py new file mode 100644 index 000000000..8b4063d07 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2image.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2image.py new file mode 100644 index 000000000..6b5110a77 --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2text.py b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2text.py new file mode 100644 index 000000000..fb5ac1702 --- /dev/null +++ b/vllm_omni/model_executor/models/dynin_omni/dynin_omni_token2text.py @@ -0,0 +1,1580 @@ +from __future__ import annotations + +import inspect +import json +from contextlib import contextmanager +from typing import Any + +import torch +import torch.nn.functional as F +from transformers import AutoTokenizer +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 ( + DYNIN_PROMPT_SOURCE_KEY, + DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT, + DYNIN_REMOTE_SETTINGS, + DYNIN_SPECIAL_TOKENS, + TASK_TO_DETOK, + DetokTarget, + _to_bool, + build_dynin_online_runtime_info, + build_dynin_prompt_payload, + coerce_token_ids_1d, + dynin_runtime_fallback, + get_dynin_magvit_attr, + get_dynin_modeling_attr, + get_dynin_sampling_attr, + infer_dynin_online_task, + logical_dynin_task, + normalize_dynin_online_prompt_text, + normalize_runtime_info, + resolve_dynin_infer_sources, + resolve_hidden_size, + resolve_remote_attr, + unwrap_first_value, +) + +logger = init_logger(__name__) + +TASK_TO_PROMPTING_TASK = { + "t2i": "t2i_gen", + "i2i": "i2i_gen", + "ti2ti": "ti2ti_gen", + "t2s": "t2s_gen", + "t2s_mmu_like": "t2s_gen", + "t2s_fixed": "t2s_fixed_gen", + "s2s": "s2s_gen", + "v2s": "v2s_gen", + "mmu": "mmu", + "mmu_fast": "mmu", + "mmu_fastdllm_v1": "mmu", + "s2t": "s2t", + "v2t": "v2t", +} + +TASK_TO_GENERATE_FN = { + "t2i": "t2i_generate", + "i2i": "i2i_generate", + "ti2ti": "ti2ti_generate", + "t2s": "t2s_generate", + "t2s_mmu_like": "t2s_generate_mmu_like", + "t2s_fixed": "t2s_fixed_generate", + "s2s": "t2s_generate_mmu_like", + "v2s": "t2s_generate_mmu_like", + "s2t": "s2t_generate", + "mmu": "mmu_generate", + "t2t": "generate", + "mmu_fast": "mmu_generate_fast", + "mmu_fastdllm_v1": "mmu_generate_fastdllm_v1", + "v2t": "mmu_generate", +} + +TASKS_USING_UNI_PROMPTING = set(TASK_TO_PROMPTING_TASK.keys()) +PROMPT_PAYLOAD_REQUIRED_TASKS = { + "t2i", + "i2i", + "ti2ti", + "t2s", + "t2s_mmu_like", + "t2s_fixed", + "s2s", + "v2s", +} + +GENERATE_RUNTIME_KWARG_KEYS = ( + "uncond_input_ids", + "uncond_attention_mask", + "noise_schedule", + "generator", + "config", + "uni_prompting", + "resolution", + "max_new_tokens", + "steps", + "block_length", + "temperature", + "top_k", + "eot_token", + "cfg_scale", + "remasking", + "mask_id", + "attention_mask", + "timesteps", + "guidance_scale", + "noise_type", + "seq_len", + "mask_token_id", + "codebook_size", + "audio_codebook_size", + "use_cache", + "threshold", + "factor", +) + +PASSTHROUGH_GENERATE_KWARG_KEYS = ( + "attention_mask", + "uncond_input_ids", + "uncond_attention_mask", + "noise_schedule", + "uni_prompting", + "generator", + "noise_type", +) + +PROMPTING_PAYLOAD_KEYS = ( + "prompting_input", + "prompting_inputs", + "dynin_inputs", + "model_inputs", + "raw_inputs", +) + +UNCOND_PROMPTING_PAYLOAD_KEYS = ( + "uncond_prompting_input", + "uncond_prompting_inputs", +) + +PROMPTING_META_KEYS = ( + "uncond_prompting_input", + "uncond_prompting_inputs", + "uni_prompting", + "prompting_task", + "prompting_config", +) + +MM_INPUT_ALIASES = { + "image": ("pixel_values", "image_embeds", "img2img"), + "video": ("pixel_values_videos", "video_embeds"), + "audio": ("input_audio_features", "audio_embeds"), +} + + +class DyninOmniToken2Text(DyninOmniStageBase): + """Stage-1: DYNIN generation + text detokenization or pass-through.""" + + stage_name = "Dynin token2text" + + 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._infer_sources = resolve_dynin_infer_sources(vllm_config=vllm_config) + if self._infer_sources.config_path: + logger.info( + "DYNIN token2text using inference config: %s", + self._infer_sources.config_path, + ) + + self.model = self._load_text_model( + self._infer_sources.model_source, + local_files_only=self._infer_sources.model_local_files_only, + ) + self.model.eval() + self.model.requires_grad_(False) + + self.hidden_size = resolve_hidden_size( + vllm_config=vllm_config, + model=self.model, + ) + + self.tokenizer: Any | None = None + self._tokenizer_path: str | None = None + self._uni_prompting: Any | None = None + self._uni_prompting_init_spec: tuple[Any, ...] | None = None + self._prompt_vq_model: Any | None = None + self._prompt_vq_model_path: str | None = None + self._prompt_vq_local_files_only: bool | None = None + self._cached_mm_inputs: dict[str, Any] = {} + + try: + self._set_tokenizer( + self._infer_sources.tokenizer_source, + local_files_only=self._infer_sources.model_local_files_only, + ) + except Exception: + self.tokenizer = None + self._tokenizer_path = None + + @staticmethod + def _load_text_model(model_path: str, *, local_files_only: bool = False) -> Any: + try: + dynin_model_cls = get_dynin_modeling_attr("DyninOmniModelLM") + try: + return dynin_model_cls.from_pretrained( + model_path, + torch_dtype=torch.bfloat16, + local_files_only=local_files_only, + ) + except TypeError: + return dynin_model_cls.from_pretrained( + model_path, + torch_dtype=torch.bfloat16, + ) + except Exception as e: + raise RuntimeError( + f"Failed to load DyninOmniModelLM via remote Dynin code for model path '{model_path}'." + ) from e + + @staticmethod + def _load_tokenizer_from_source( + source: str, + *, + local_files_only: bool = False, + trust_remote_code: bool = False, + ) -> Any: + load_kwargs = { + "trust_remote_code": trust_remote_code, + "local_files_only": _to_bool(local_files_only, default=False), + } + try: + return AutoTokenizer.from_pretrained(source, **load_kwargs) + except TypeError: + load_kwargs.pop("local_files_only", None) + return AutoTokenizer.from_pretrained(source, **load_kwargs) + + def _set_tokenizer(self, source: str, *, local_files_only: bool) -> None: + try: + tokenizer = self._load_tokenizer_from_source( + source, + local_files_only=local_files_only, + trust_remote_code=False, + ) + except Exception as e: + logger.info( + "Falling back to trust_remote_code=True tokenizer loading for %s: %s", + source, + e, + ) + tokenizer = self._load_tokenizer_from_source( + source, + local_files_only=local_files_only, + trust_remote_code=True, + ) + + self.tokenizer = tokenizer + self._tokenizer_path = source + self._reset_uni_prompting_cache() + + def _reset_uni_prompting_cache(self) -> None: + self._uni_prompting = None + self._uni_prompting_init_spec = None + + def get_language_model(self) -> Any: + return self.model + + @staticmethod + def _merge_runtime_info_missing_values( + runtime_info: dict[str, Any], + fallback_info: dict[str, Any], + ) -> dict[str, Any]: + merged = dict(runtime_info) + for key, value in fallback_info.items(): + if unwrap_first_value(merged.get(key), None) is None: + merged[key] = value + return merged + + def _runtime_info_needs_bootstrap( + self, + runtime_info: dict[str, Any], + logical_task_name: str, + ) -> bool: + task = str(unwrap_first_value(runtime_info.get("task"), "") or "").lower() + detok_id = unwrap_first_value(runtime_info.get("detok_id"), None) + prompt_length = unwrap_first_value(runtime_info.get("prompt_length"), None) + + if not task or detok_id is None: + return True + if prompt_length is None: + return True + if ( + task in PROMPT_PAYLOAD_REQUIRED_TASKS + and self._find_first_payload( + runtime_info=runtime_info, + kwargs={}, + keys=PROMPTING_PAYLOAD_KEYS, + ) + is None + ): + return True + if logical_task_name in {"t2i", "i2i"}: + for key in ("codebook_size", "text_vocab_size", "vq_model_image_path"): + if unwrap_first_value(runtime_info.get(key), None) is None: + return True + if logical_task_name == "t2s": + for key in ("audio_codebook_size", "condition", "vq_model_audio_path"): + if unwrap_first_value(runtime_info.get(key), None) is None: + return True + return False + + def _decode_prompt_for_bootstrap( + self, + input_ids: torch.Tensor, + runtime_info: dict[str, Any], + ) -> str: + self._maybe_load_runtime_tokenizer(runtime_info) + if self.tokenizer is None: + return "" + token_ids = coerce_token_ids_1d(input_ids).detach().cpu().tolist() + try: + return str(self.tokenizer.decode(token_ids, skip_special_tokens=False)) + except Exception: + return "" + + def _bootstrap_runtime_info_if_needed( + self, + *, + input_ids: torch.Tensor, + runtime_info: dict[str, Any], + kwargs: dict[str, Any], + ) -> dict[str, Any]: + if unwrap_first_value(runtime_info.get(DYNIN_PROMPT_SOURCE_KEY), None) == DYNIN_PROMPT_SOURCE_OFFLINE_PREBUILT: + return runtime_info + + mm_inputs = self._collect_mm_inputs(**kwargs) + decoded_prompt = "" + + task_value = unwrap_first_value(runtime_info.get("task"), None) + if task_value is None: + decoded_prompt = self._decode_prompt_for_bootstrap(input_ids, runtime_info) + logical_task_name = infer_dynin_online_task( + decoded_prompt=decoded_prompt, + has_image="image" in mm_inputs, + has_audio="audio" in mm_inputs, + has_video="video" in mm_inputs, + ) + else: + logical_task_name = logical_dynin_task(task_value) + + if not self._runtime_info_needs_bootstrap(runtime_info, logical_task_name): + return runtime_info + + self._maybe_load_runtime_tokenizer(runtime_info) + if self.tokenizer is None: + logger.warning("Unable to bootstrap Dynin runtime info because tokenizer is unavailable.") + return runtime_info + + if not decoded_prompt: + decoded_prompt = self._decode_prompt_for_bootstrap(input_ids, runtime_info) + + text_vocab_size = int(len(self.tokenizer)) + prompt_len = int(coerce_token_ids_1d(input_ids).numel()) + dynin_config_path = self._infer_sources.config_path + + base_runtime_info = build_dynin_online_runtime_info( + task=logical_task_name, + text_vocab_size=text_vocab_size, + infer_sources=self._infer_sources, + dynin_config_path=dynin_config_path, + attention_mask=([1] * prompt_len) if logical_task_name == "t2t" else None, + prompt_length=prompt_len if logical_task_name == "t2t" else None, + ) + merged_runtime_info = self._merge_runtime_info_missing_values(runtime_info, base_runtime_info) + + payload_required = logical_task_name in {"t2i", "i2i", "t2s"} + existing_prompt_payload = self._find_first_payload( + runtime_info=merged_runtime_info, + kwargs=kwargs, + keys=PROMPTING_PAYLOAD_KEYS, + ) + has_prompt_payload = existing_prompt_payload is not None + needs_prompt_length = unwrap_first_value(merged_runtime_info.get("prompt_length"), None) is None + if not payload_required: + return merged_runtime_info + + use_train_i2i_prompt = _to_bool( + unwrap_first_value( + merged_runtime_info.get("use_train_i2i_prompt"), + dynin_runtime_fallback(logical_task_name, "use_train_i2i_prompt", logical_task_name == "i2i"), + ), + default=logical_task_name == "i2i", + ) + t2s_token_length = int( + dynin_runtime_fallback( + logical_task_name, + "t2s_token_length", + unwrap_first_value(merged_runtime_info.get("t2s_token_length"), None), + ) + or 383 + ) + image_resolution = int( + dynin_runtime_fallback( + logical_task_name, + "image_resolution", + unwrap_first_value(merged_runtime_info.get("image_resolution"), None), + ) + or 336 + ) + + image_token_count = int( + dynin_runtime_fallback( + logical_task_name, + "image_token_count", + unwrap_first_value(merged_runtime_info.get("seq_len"), None), + ) + or 0 + ) + image_tokens: torch.Tensor | None = None + if logical_task_name == "i2i" and (not has_prompt_payload or image_token_count <= 0): + image_tokens = self._encode_prompt_image_tokens( + runtime_info=merged_runtime_info, + mm_inputs=mm_inputs, + resolution=image_resolution, + ) + image_token_count = int(image_tokens.numel()) + + mask_token_id = int(unwrap_first_value(merged_runtime_info.get("mask_token_id"), 126336)) + prompting_input = self._unwrap_singleton(existing_prompt_payload) + prompting_task = str( + unwrap_first_value( + merged_runtime_info.get("prompting_task"), + TASK_TO_PROMPTING_TASK.get( + str(unwrap_first_value(merged_runtime_info.get("task"), "mmu")).lower(), + "mmu", + ), + ) + ) + if not has_prompt_payload: + prompt_text = normalize_dynin_online_prompt_text(logical_task_name, decoded_prompt) + prompting_input, prompting_task = build_dynin_prompt_payload( + task=logical_task_name, + text=prompt_text, + image_tokens=image_tokens, + image_placeholder_tokens=image_token_count, + audio_placeholder_tokens=t2s_token_length, + image_token_offset=text_vocab_size, + mask_token_id=mask_token_id, + use_train_i2i_prompt=use_train_i2i_prompt, + ) + + prompt_runtime_info = build_dynin_online_runtime_info( + task=logical_task_name, + text_vocab_size=text_vocab_size, + infer_sources=self._infer_sources, + dynin_config_path=dynin_config_path, + image_token_count=image_token_count, + t2s_token_length=t2s_token_length, + use_train_i2i_prompt=use_train_i2i_prompt, + ) + prompt_runtime_info["prompting_task"] = [str(prompting_task)] + prompt_runtime_info["prompting_input"] = [prompting_input] + merged_runtime_info = self._merge_runtime_info_missing_values(merged_runtime_info, prompt_runtime_info) + + if not needs_prompt_length and has_prompt_payload: + return merged_runtime_info + + uni_prompting = self._get_or_create_uni_prompting( + runtime_info=merged_runtime_info, + kwargs=kwargs, + ) + if uni_prompting is not None: + prepared_input_ids, prepared_attention_mask = self._prepare_prompting_input( + payload=prompting_input, + task=str(unwrap_first_value(merged_runtime_info.get("task"), "mmu")), + runtime_info=merged_runtime_info, + kwargs=kwargs, + uni_prompting=uni_prompting, + ref_device=input_ids.device, + ) + if prepared_input_ids is not None: + prepared_prompt_len = int(prepared_input_ids.shape[-1]) + prepared_attention_list: list[int] | None = None + if prepared_attention_mask is not None: + prepared_attention_list = prepared_attention_mask.view(-1).detach().cpu().tolist() + final_runtime_info = build_dynin_online_runtime_info( + task=logical_task_name, + text_vocab_size=text_vocab_size, + infer_sources=self._infer_sources, + dynin_config_path=dynin_config_path, + prompting_input=prompting_input, + attention_mask=prepared_attention_list, + prompt_length=prepared_prompt_len, + image_token_count=image_token_count, + t2s_token_length=t2s_token_length, + use_train_i2i_prompt=use_train_i2i_prompt, + ) + final_runtime_info["prompting_task"] = [str(prompting_task)] + + guidance_scale = float(unwrap_first_value(merged_runtime_info.get("guidance_scale"), 0.0)) + if logical_task_name in {"t2i", "i2i"} and guidance_scale > 0: + uncond_prompting_input, _ = build_dynin_prompt_payload( + task=logical_task_name, + text="", + image_tokens=image_tokens, + image_placeholder_tokens=image_token_count, + audio_placeholder_tokens=t2s_token_length, + image_token_offset=text_vocab_size, + mask_token_id=mask_token_id, + use_train_i2i_prompt=use_train_i2i_prompt, + ) + final_runtime_info["uncond_prompting_input"] = [uncond_prompting_input] + + merged_runtime_info = self._merge_runtime_info_missing_values( + merged_runtime_info, + final_runtime_info, + ) + + return merged_runtime_info + + @staticmethod + def _build_downstream_runtime_info(runtime_info: dict[str, Any]) -> dict[str, Any]: + bridge_keys = ( + "task", + "detok_id", + "dynin_config_path", + "codebook_size", + "audio_codebook_size", + "text_vocab_size", + "num_new_special_tokens", + "image_vocab_offset", + "audio_vocab_offset", + "t2s_vocab_start", + "condition", + "t2s_condition", + "vq_model_image_path", + "vq_model_image_local_files_only", + "vq_model_audio_path", + "vq_model_audio_local_files_only", + "model_local_files_only", + "local_files_only", + "hf_hub_disable_xet", + "disable_hf_xet", + ) + return {key: runtime_info[key] for key in bridge_keys if key in runtime_info} + + @staticmethod + def _jsonify_runtime_value(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, (list, tuple)): + return [DyninOmniToken2Text._jsonify_runtime_value(item) for item in value] + if isinstance(value, dict): + return {str(key): DyninOmniToken2Text._jsonify_runtime_value(val) for key, val in value.items()} + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return str(value) + + def _encode_runtime_info_tensor( + self, + runtime_info: dict[str, Any], + *, + device: torch.device, + ) -> torch.Tensor | None: + if not runtime_info: + return None + payload = {key: self._jsonify_runtime_value(value) for key, value in runtime_info.items()} + encoded = json.dumps( + payload, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + if not encoded: + return None + return torch.tensor(list(encoded), dtype=torch.uint8, device=device) + + 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("token2text stage requires input_ids") + try: + runtime_info = normalize_runtime_info(kwargs.get("runtime_additional_information")) + runtime_info = self._bootstrap_runtime_info_if_needed( + input_ids=input_ids, + runtime_info=runtime_info, + kwargs=kwargs, + ) + task = str(unwrap_first_value(runtime_info.get("task"), "mmu")).lower() + + detok_id = int( + unwrap_first_value( + runtime_info.get("detok_id"), + TASK_TO_DETOK.get(task, DetokTarget.TEXT), + ) + ) + + token_ids = self._generate_token_ids( + task=task, + input_ids=input_ids, + runtime_info=runtime_info, + kwargs=kwargs, + ) + bridge_runtime_info = self._build_downstream_runtime_info(runtime_info) + runtime_info_tensor = self._encode_runtime_info_tensor( + bridge_runtime_info, + device=token_ids.device, + ) + + if detok_id != int(DetokTarget.TEXT): + multimodal_outputs = { + "token_ids": token_ids, + "detok_id": torch.tensor( + [detok_id], + dtype=torch.long, + device=token_ids.device, + ), + } + if runtime_info_tensor is not None: + multimodal_outputs["runtime_info_json"] = runtime_info_tensor + return OmniOutput( + text_hidden_states=None, + multimodal_outputs=multimodal_outputs, + ) + + decode_tokens = self._extract_decode_tokens(token_ids, runtime_info=runtime_info) + multimodal_outputs = { + "token_ids": token_ids, + "text_tokens": decode_tokens, + "detok_id": torch.tensor( + [detok_id], + dtype=torch.long, + device=token_ids.device, + ), + } + if runtime_info_tensor is not None: + multimodal_outputs["runtime_info_json"] = runtime_info_tensor + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs=multimodal_outputs, + ) + finally: + self._cached_mm_inputs = {} + + def _generate_token_ids( + self, + task: str, + input_ids: torch.Tensor, + runtime_info: dict[str, Any], + kwargs: dict[str, Any], + ) -> torch.Tensor: + precomputed = self._get_precomputed_token_ids(runtime_info) + if precomputed is not None: + return coerce_token_ids_1d(precomputed, ref_device=input_ids.device) + + gen_fn_name = TASK_TO_GENERATE_FN.get(task, "mmu_generate") + gen_fn = self._resolve_generate_fn(gen_fn_name) + + gen_kwargs = self._collect_generate_kwargs(runtime_info=runtime_info, kwargs=kwargs) + + if "noise_schedule" not in gen_kwargs: + noise_schedule = self._resolve_noise_schedule( + runtime_info=runtime_info, + kwargs=kwargs, + ) + if noise_schedule is not None: + gen_kwargs["noise_schedule"] = noise_schedule + + if task in TASKS_USING_UNI_PROMPTING and "uni_prompting" not in gen_kwargs: + uni_prompting = self._get_or_create_uni_prompting( + runtime_info=runtime_info, + kwargs=kwargs, + ) + if uni_prompting is not None: + gen_kwargs["uni_prompting"] = uni_prompting + + should_prepare_prompting_inputs = task in TASKS_USING_UNI_PROMPTING or self._contains_prompting_payload( + runtime_info=runtime_info, kwargs=kwargs + ) + if should_prepare_prompting_inputs: + input_ids, gen_kwargs = self._prepare_prompting_inputs_if_needed( + task=task, + input_ids=input_ids, + runtime_info=runtime_info, + kwargs=kwargs, + gen_kwargs=gen_kwargs, + ) + + input_ids, gen_kwargs = self._normalize_generate_inputs( + input_ids=input_ids, + gen_kwargs=gen_kwargs, + ref_device=input_ids.device, + ) + gen_kwargs = self._filter_supported_generate_kwargs( + gen_fn=gen_fn, + gen_kwargs=gen_kwargs, + fn_name=gen_fn_name, + ) + + generated = self._call_generate_fn( + gen_fn=gen_fn, + input_ids=input_ids, + gen_kwargs=gen_kwargs, + ) + return coerce_token_ids_1d(generated, ref_device=input_ids.device) + + @staticmethod + def _get_precomputed_token_ids(runtime_info: dict[str, Any]) -> Any | None: + precomputed = runtime_info.get("generated_token_ids") + if precomputed is None: + precomputed = runtime_info.get("token_ids") + return precomputed + + def _resolve_generate_fn(self, fn_name: str) -> Any: + if not hasattr(self.model, fn_name): + raise RuntimeError( + f"DYNIN model does not expose '{fn_name}'. " + "Pass additional_information.generated_token_ids or adjust task mapping." + ) + return getattr(self.model, fn_name) + + @staticmethod + def _collect_generate_kwargs( + *, + runtime_info: dict[str, Any], + kwargs: dict[str, Any], + ) -> dict[str, Any]: + gen_kwargs: dict[str, Any] = {} + + for key in GENERATE_RUNTIME_KWARG_KEYS: + if key in runtime_info: + gen_kwargs[key] = unwrap_first_value(runtime_info[key]) + + for key in PASSTHROUGH_GENERATE_KWARG_KEYS: + if key not in gen_kwargs and key in kwargs: + gen_kwargs[key] = kwargs[key] + + return gen_kwargs + + @staticmethod + def _contains_prompting_payload( + runtime_info: dict[str, Any], + kwargs: dict[str, Any], + ) -> bool: + keys = PROMPTING_PAYLOAD_KEYS + PROMPTING_META_KEYS + return any(key in runtime_info for key in keys) or any(key in kwargs for key in keys) + + @staticmethod + def _filter_supported_generate_kwargs( + *, + gen_fn: Any, + gen_kwargs: dict[str, Any], + fn_name: str, + ) -> dict[str, Any]: + if not gen_kwargs: + return gen_kwargs + + try: + signature = inspect.signature(gen_fn) + except (TypeError, ValueError): + return gen_kwargs + + params = signature.parameters + accepts_var_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) + if accepts_var_kwargs: + return gen_kwargs + + allowed_keys = { + name + for name, param in params.items() + if param.kind + in ( + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + } + filtered = {k: v for k, v in gen_kwargs.items() if k in allowed_keys} + + removed_keys = sorted(set(gen_kwargs.keys()) - set(filtered.keys())) + if removed_keys: + logger.debug("Filtered unsupported kwargs for %s: %s", fn_name, removed_keys) + + return filtered + + @staticmethod + def _call_generate_fn( + *, + gen_fn: Any, + input_ids: torch.Tensor, + gen_kwargs: dict[str, Any], + ) -> Any: + try: + signature = inspect.signature(gen_fn) + params = signature.parameters + except (TypeError, ValueError): + params = {} + + if "idx" in params: + return gen_fn(idx=input_ids, **gen_kwargs) + if "input_ids" in params: + return gen_fn(input_ids=input_ids, **gen_kwargs) + + try: + return gen_fn(input_ids, **gen_kwargs) + except TypeError: + try: + return gen_fn(idx=input_ids, **gen_kwargs) + except TypeError: + return gen_fn(input_ids=input_ids, **gen_kwargs) + + def _normalize_generate_inputs( + self, + *, + input_ids: torch.Tensor, + gen_kwargs: dict[str, Any], + ref_device: torch.device, + ) -> tuple[torch.Tensor, dict[str, Any]]: + normalized_input_ids = self._coerce_long_tensor_2d(input_ids, ref_device) + if normalized_input_ids is None: + normalized_input_ids = input_ids + + normalized_kwargs = dict(gen_kwargs) + for key in ("attention_mask", "uncond_input_ids", "uncond_attention_mask"): + if key not in normalized_kwargs: + continue + normalized_value = self._coerce_long_tensor_2d( + normalized_kwargs[key], + ref_device, + ) + if normalized_value is not None: + normalized_kwargs[key] = normalized_value + + return normalized_input_ids, normalized_kwargs + + def _get_or_create_uni_prompting( + self, + runtime_info: dict[str, Any], + kwargs: dict[str, Any], + ) -> Any | None: + runtime_uni_prompting = runtime_info.get("uni_prompting") + if runtime_uni_prompting is not None: + runtime_uni_prompting = self._unwrap_singleton(runtime_uni_prompting) + if runtime_uni_prompting is not None: + return runtime_uni_prompting + + kwargs_uni_prompting = self._unwrap_singleton(kwargs.get("uni_prompting")) + if kwargs_uni_prompting is not None: + return kwargs_uni_prompting + + self._maybe_load_runtime_tokenizer(runtime_info) + if self.tokenizer is None: + return None + + use_reserved_token = _to_bool( + unwrap_first_value( + runtime_info.get("use_reserved_token"), + unwrap_first_value(runtime_info.get("prompting_use_reserved_token"), True), + ), + default=True, + ) + + max_text_len_value = unwrap_first_value( + runtime_info.get("prompt_max_text_len"), + unwrap_first_value( + runtime_info.get("prompting_max_text_len"), + unwrap_first_value(runtime_info.get("max_text_len"), None), + ), + ) + cond_dropout_value = unwrap_first_value( + runtime_info.get("cond_dropout_prob"), + unwrap_first_value(runtime_info.get("prompting_cond_dropout_prob"), None), + ) + max_audio_len_value = unwrap_first_value( + runtime_info.get("max_audio_len"), + unwrap_first_value(runtime_info.get("t2s_token_length"), None), + ) + max_audio_len_short_value = unwrap_first_value( + runtime_info.get("max_audio_len_short"), + None, + ) + + max_text_len: int | None = None + if max_text_len_value is not None: + try: + parsed = int(max_text_len_value) + if parsed > 0: + max_text_len = parsed + except Exception: + pass + + cond_dropout_prob: float | None = None + if cond_dropout_value is not None: + try: + cond_dropout_prob = float(cond_dropout_value) + except Exception: + pass + + max_audio_len: int | None = None + if max_audio_len_value is not None: + try: + parsed = int(max_audio_len_value) + if parsed > 0: + max_audio_len = max(parsed, 512) + except Exception: + pass + + max_audio_len_short: int | None = None + if max_audio_len_short_value is not None: + try: + parsed = int(max_audio_len_short_value) + if parsed > 0: + max_audio_len_short = parsed + except Exception: + pass + elif max_audio_len is not None: + max_audio_len_short = max(256, max_audio_len // 2) + + if self._uni_prompting is not None: + if max_text_len is None and hasattr(self._uni_prompting, "max_text_len"): + try: + existing_max_text_len = int(getattr(self._uni_prompting, "max_text_len")) + if existing_max_text_len > 0: + max_text_len = existing_max_text_len - 1 + except Exception: + pass + if cond_dropout_prob is None and hasattr(self._uni_prompting, "cond_dropout_prob"): + try: + cond_dropout_prob = float(getattr(self._uni_prompting, "cond_dropout_prob")) + except Exception: + pass + + desired_spec = ( + id(self.tokenizer), + use_reserved_token, + max_text_len, + cond_dropout_prob, + max_audio_len, + max_audio_len_short, + ) + + if self._uni_prompting is not None and self._uni_prompting_init_spec != desired_spec: + self._reset_uni_prompting_cache() + + if self._uni_prompting is None: + try: + universal_prompting_cls = resolve_remote_attr( + "UniversalPrompting", + module_name="prompting_utils", + settings=DYNIN_REMOTE_SETTINGS, + source=self._infer_sources.model_source, + local_files_only=self._infer_sources.model_local_files_only, + fallback_module_names=("modeling_dynin_omni",), + optional=True, + ) + except Exception: + universal_prompting_cls = None + + try: + if universal_prompting_cls is None: + raise ImportError("UniversalPrompting is not available in the configured remote Dynin code.") + + init_kwargs: dict[str, Any] = { + "use_reserved_token": use_reserved_token, + "special_tokens": DYNIN_SPECIAL_TOKENS, + "ignore_id": -100, + } + if max_text_len is not None: + init_kwargs["max_text_len"] = max_text_len + if cond_dropout_prob is not None: + init_kwargs["cond_dropout_prob"] = cond_dropout_prob + if max_audio_len is not None: + init_kwargs["max_audio_len"] = max_audio_len + if max_audio_len_short is not None: + init_kwargs["max_audio_len_short"] = max_audio_len_short + + try: + self._uni_prompting = universal_prompting_cls(self.tokenizer, **init_kwargs) + except TypeError: + trimmed_audio_kwargs = dict(init_kwargs) + trimmed_audio_kwargs.pop("max_audio_len", None) + trimmed_audio_kwargs.pop("max_audio_len_short", None) + try: + self._uni_prompting = universal_prompting_cls(self.tokenizer, **trimmed_audio_kwargs) + except TypeError: + minimal_kwargs = dict(trimmed_audio_kwargs) + minimal_kwargs.pop("special_tokens", None) + minimal_kwargs.pop("ignore_id", None) + self._uni_prompting = universal_prompting_cls(self.tokenizer, **minimal_kwargs) + self._uni_prompting_init_spec = desired_spec + except Exception as e: + logger.warning("Failed to initialize UniversalPrompting: %s", e) + self._reset_uni_prompting_cache() + + return self._uni_prompting + + @staticmethod + def _unwrap_singleton(value: Any) -> Any: + if isinstance(value, list) and len(value) == 1: + return value[0] + return value + + @classmethod + def _coerce_schedule_params(cls, value: Any) -> dict[str, Any]: + value = cls._unwrap_singleton(value) + if value is None: + return {} + if isinstance(value, dict): + return {str(k): v for k, v in value.items()} + if hasattr(value, "items"): + try: + return {str(k): v for k, v in dict(value).items()} + except Exception: + return {} + if isinstance(value, str): + text = value.strip() + if not text: + return {} + try: + parsed = json.loads(text) + except Exception: + return {} + if isinstance(parsed, dict): + return {str(k): v for k, v in parsed.items()} + return {} + + def _resolve_noise_schedule( + self, + runtime_info: dict[str, Any], + kwargs: dict[str, Any], + ) -> Any | None: + runtime_noise_schedule = unwrap_first_value( + runtime_info.get("noise_schedule"), + kwargs.get("noise_schedule"), + ) + runtime_noise_schedule = self._unwrap_singleton(runtime_noise_schedule) + if callable(runtime_noise_schedule): + return runtime_noise_schedule + + schedule_name: str | None = None + if isinstance(runtime_noise_schedule, str) and runtime_noise_schedule.strip(): + schedule_name = runtime_noise_schedule.strip() + + if schedule_name is None: + for key in ("noise_schedule_name", "mask_schedule", "schedule"): + value = unwrap_first_value(runtime_info.get(key), None) + if value is None and key in kwargs: + value = self._unwrap_singleton(kwargs.get(key)) + if isinstance(value, str) and value.strip(): + schedule_name = value.strip() + break + + if schedule_name is None: + return None + + schedule_params = self._coerce_schedule_params( + unwrap_first_value( + runtime_info.get("noise_schedule_params"), + kwargs.get("noise_schedule_params"), + ) + ) + + try: + get_mask_schedule = get_dynin_sampling_attr("get_mask_schedule") + return get_mask_schedule(schedule_name, **schedule_params) + except Exception as e: + logger.warning( + "Failed to resolve mask schedule '%s' with params=%s: %s", + schedule_name, + schedule_params, + e, + ) + return None + + @staticmethod + def _coerce_long_tensor_2d( + value: Any, + device: torch.device, + ) -> torch.Tensor | None: + if value is None: + return None + out = value if isinstance(value, torch.Tensor) else torch.as_tensor(value) + if out.ndim == 1: + out = out.unsqueeze(0) + if out.ndim > 2: + out = out.view(out.shape[0], -1) + return out.to(device=device, dtype=torch.long).contiguous() + + @staticmethod + def _config_get(config_obj: Any, key: str) -> Any: + if config_obj is None: + return None + if isinstance(config_obj, dict): + return config_obj.get(key) + if hasattr(config_obj, "get"): + try: + return config_obj.get(key) + except Exception: + return None + return None + + @classmethod + def _is_numeric_token_structure(cls, value: Any) -> bool: + if isinstance(value, torch.Tensor): + return True + if isinstance(value, bool): + return True + if isinstance(value, int): + return True + if isinstance(value, float): + return float(value).is_integer() + if isinstance(value, (list, tuple)): + if not value: + return False + return all(cls._is_numeric_token_structure(v) for v in value) + return False + + @classmethod + def _materialize_prompting_payload(cls, value: Any, ref_device: torch.device) -> Any: + if isinstance(value, torch.Tensor): + return value.to(device=ref_device, dtype=torch.long).contiguous() + if isinstance(value, dict): + return {k: cls._materialize_prompting_payload(v, ref_device) for k, v in value.items()} + if isinstance(value, (list, tuple)): + if cls._is_numeric_token_structure(value): + try: + return torch.as_tensor(value, dtype=torch.long, device=ref_device) + except Exception: + pass + converted = [cls._materialize_prompting_payload(v, ref_device) for v in value] + return tuple(converted) if isinstance(value, tuple) else converted + return value + + @contextmanager + def _temporary_prompting_overrides(self, uni_prompting: Any, prompting_cfg: Any): + restore_values: dict[str, Any] = {} + try: + max_text_len_override = self._config_get(prompting_cfg, "max_text_len_override") + if max_text_len_override is not None and hasattr(uni_prompting, "max_text_len"): + try: + override_int = int(max_text_len_override) + if override_int > 0: + restore_values["max_text_len"] = getattr(uni_prompting, "max_text_len") + setattr(uni_prompting, "max_text_len", override_int + 1) + except Exception: + pass + yield + finally: + for attr_name, original_value in restore_values.items(): + try: + setattr(uni_prompting, attr_name, original_value) + except Exception: + pass + + def _prepare_prompting_input( + self, + *, + payload: Any, + task: str, + runtime_info: dict[str, Any], + kwargs: dict[str, Any], + uni_prompting: Any, + ref_device: torch.device, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if payload is None: + return None, None + + payload = self._unwrap_singleton(payload) + prompting_task = str( + self._unwrap_singleton( + unwrap_first_value( + runtime_info.get("prompting_task"), + TASK_TO_PROMPTING_TASK.get(task, task), + ) + ) + ) + prompting_cfg = self._unwrap_singleton( + unwrap_first_value( + runtime_info.get("prompting_config"), + kwargs.get("prompting_config"), + ) + ) + + if isinstance(payload, dict): + if payload.get("task") is not None: + prompting_task = str(payload["task"]) + if payload.get("config") is not None: + prompting_cfg = payload["config"] + payload = payload.get("input", payload.get("inputs", payload.get("data", payload))) + + payload = self._materialize_prompting_payload(payload, ref_device) + + try: + with self._temporary_prompting_overrides(uni_prompting, prompting_cfg): + prepared = uni_prompting(payload, prompting_task, config=prompting_cfg) + except Exception as e: + logger.warning( + "UniversalPrompting failed for task=%s prompting_task=%s: %s", + task, + prompting_task, + e, + ) + return None, None + + if isinstance(prepared, tuple): + prepared_input_ids = prepared[0] if len(prepared) > 0 else None + prepared_attention_mask = prepared[1] if len(prepared) > 1 else None + else: + prepared_input_ids = prepared + prepared_attention_mask = None + + return ( + self._coerce_long_tensor_2d(prepared_input_ids, ref_device), + self._coerce_long_tensor_2d(prepared_attention_mask, ref_device), + ) + + def _prepare_prompting_inputs_if_needed( + self, + *, + task: str, + input_ids: torch.Tensor, + runtime_info: dict[str, Any], + kwargs: dict[str, Any], + gen_kwargs: dict[str, Any], + ) -> tuple[torch.Tensor, dict[str, Any]]: + uni_prompting = gen_kwargs.get("uni_prompting") + if uni_prompting is None: + uni_prompting = self._get_or_create_uni_prompting( + runtime_info=runtime_info, + kwargs=kwargs, + ) + if uni_prompting is not None: + gen_kwargs["uni_prompting"] = uni_prompting + + if uni_prompting is None: + return input_ids, gen_kwargs + + payload = self._find_first_payload( + runtime_info=runtime_info, + kwargs=kwargs, + keys=PROMPTING_PAYLOAD_KEYS, + ) + + if payload is not None: + prepared_input_ids, prepared_attention_mask = self._prepare_prompting_input( + payload=payload, + task=task, + runtime_info=runtime_info, + kwargs=kwargs, + uni_prompting=uni_prompting, + ref_device=input_ids.device, + ) + if prepared_input_ids is not None: + input_ids = prepared_input_ids + if prepared_attention_mask is not None and "attention_mask" not in gen_kwargs: + gen_kwargs["attention_mask"] = prepared_attention_mask + + uncond_payload = self._find_first_payload( + runtime_info=runtime_info, + kwargs=kwargs, + keys=UNCOND_PROMPTING_PAYLOAD_KEYS, + ) + if uncond_payload is not None and "uncond_input_ids" not in gen_kwargs: + uncond_input_ids, uncond_attention_mask = self._prepare_prompting_input( + payload=uncond_payload, + task=task, + runtime_info=runtime_info, + kwargs=kwargs, + uni_prompting=uni_prompting, + ref_device=input_ids.device, + ) + if uncond_input_ids is not None: + gen_kwargs["uncond_input_ids"] = uncond_input_ids + if uncond_attention_mask is not None and "uncond_attention_mask" not in gen_kwargs: + gen_kwargs["uncond_attention_mask"] = uncond_attention_mask + + return input_ids, gen_kwargs + + @staticmethod + def _find_first_payload( + *, + runtime_info: dict[str, Any], + kwargs: dict[str, Any], + keys: tuple[str, ...], + ) -> Any | None: + for key in keys: + if key in runtime_info: + return runtime_info[key] + if key in kwargs: + return kwargs[key] + return None + + def _extract_decode_tokens( + self, + tokens: torch.Tensor, + runtime_info: dict[str, Any], + ) -> torch.Tensor: + prompt_len = int( + unwrap_first_value( + runtime_info.get("prompt_length"), + unwrap_first_value( + runtime_info.get("prompt_len"), + unwrap_first_value(runtime_info.get("prompt_token_len"), 0), + ), + ) + ) + + decode_tokens = tokens + if 0 < prompt_len < tokens.numel(): + decode_tokens = tokens[prompt_len:] + + text_vocab_size = unwrap_first_value(runtime_info.get("text_vocab_size"), None) + if text_vocab_size is None and self.tokenizer is not None: + text_vocab_size = len(self.tokenizer) + + if text_vocab_size is not None: + vocab_size = int(text_vocab_size) + valid = decode_tokens[(decode_tokens >= 0) & (decode_tokens < vocab_size)] + if valid.numel() > 0: + decode_tokens = valid + + return decode_tokens.contiguous() + + def _decode_text(self, tokens: torch.Tensor, runtime_info: dict[str, Any]) -> str: + self._maybe_load_runtime_tokenizer(runtime_info) + if self.tokenizer is None: + return "" + try: + return self.tokenizer.decode( + tokens.detach().cpu().tolist(), + skip_special_tokens=True, + ) + except Exception: + return "" + + def _maybe_load_runtime_tokenizer(self, runtime_info: dict[str, Any]) -> None: + tokenizer_path = unwrap_first_value(runtime_info.get("tokenizer_path"), None) + if tokenizer_path is not None: + tokenizer_path = str(tokenizer_path) + + runtime_local_files_only = unwrap_first_value( + runtime_info.get("local_files_only_model"), + unwrap_first_value( + runtime_info.get("model_local_files_only"), + unwrap_first_value( + runtime_info.get("local_files_only"), + self._infer_sources.model_local_files_only, + ), + ), + ) + local_only = _to_bool( + runtime_local_files_only, + default=self._infer_sources.model_local_files_only, + ) + + if tokenizer_path and tokenizer_path != self._tokenizer_path: + try: + logger.info("Loading DYNIN text tokenizer from %s", tokenizer_path) + self._set_tokenizer(tokenizer_path, local_files_only=local_only) + except Exception as e: + logger.warning("Failed to load tokenizer from %s: %s", tokenizer_path, e) + + def _ensure_prompt_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._prompt_vq_model is None + or self._prompt_vq_model_path != model_path + or self._prompt_vq_local_files_only != local_files_only + ): + logger.info( + "Loading DYNIN prompt VQ encoder from %s (local_files_only=%s)", + model_path, + local_files_only, + ) + magvit_cls = get_dynin_magvit_attr( + "MAGVITv2", + source=model_path, + local_files_only=local_files_only, + ) + try: + self._prompt_vq_model = magvit_cls.from_pretrained( + model_path, + local_files_only=local_files_only, + ) + except TypeError: + self._prompt_vq_model = magvit_cls.from_pretrained(model_path) + self._prompt_vq_model.eval() + self._prompt_vq_model.requires_grad_(False) + self._prompt_vq_model_path = model_path + self._prompt_vq_local_files_only = local_files_only + if hasattr(self._prompt_vq_model, "to"): + self._prompt_vq_model = self._prompt_vq_model.to(ref_device) + return self._prompt_vq_model + + @staticmethod + def _prepare_prompt_image_tensor( + image: Any, + *, + resolution: int, + device: torch.device, + ) -> torch.Tensor: + tensor = image if isinstance(image, torch.Tensor) else torch.as_tensor(image) + if tensor.ndim == 4: + tensor = tensor[0] + if tensor.ndim != 3: + raise ValueError(f"Unsupported image tensor shape for Dynin bootstrap: {tuple(tensor.shape)}") + + if tensor.shape[0] not in (1, 3, 4) and tensor.shape[-1] 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(device=device, dtype=torch.float32) + if tensor.numel() > 0 and tensor.max() > 1.0: + tensor = tensor / 255.0 + + tensor = tensor.unsqueeze(0) + _, _, height, width = tensor.shape + short_side = max(1, min(int(height), int(width))) + scale = float(resolution) / float(short_side) + new_height = max(1, int(round(height * scale))) + new_width = max(1, int(round(width * scale))) + tensor = F.interpolate( + tensor, + size=(new_height, new_width), + mode="bicubic", + align_corners=False, + ) + top = max(0, (new_height - resolution) // 2) + left = max(0, (new_width - resolution) // 2) + tensor = tensor[:, :, top : top + resolution, left : left + resolution] + if tensor.shape[-2:] != (resolution, resolution): + tensor = F.interpolate( + tensor, + size=(resolution, resolution), + mode="bicubic", + align_corners=False, + ) + tensor = torch.clamp(tensor, min=0.0, max=1.0) + return ((tensor - 0.5) / 0.5).contiguous() + + def _encode_prompt_image_tokens( + self, + *, + runtime_info: dict[str, Any], + mm_inputs: dict[str, Any], + resolution: int, + ) -> torch.Tensor: + image_value = mm_inputs.get("image") + image_items = self._split_mm_items(image_value) + if not image_items: + raise ValueError("Dynin online i2i bootstrap requires an image input.") + + device = self._default_mm_device() + image_tensor = self._prepare_prompt_image_tensor( + image_items[0], + resolution=resolution, + device=device, + ) + vq_model = self._ensure_prompt_vq_model(runtime_info=runtime_info, ref_device=device) + with torch.no_grad(): + token_ids = vq_model.get_code(image_tensor) + token_ids = torch.as_tensor(token_ids, dtype=torch.long).detach().cpu() + if token_ids.ndim == 2 and token_ids.shape[0] == 1: + token_ids = token_ids[0] + return token_ids.contiguous() + + @staticmethod + def _split_mm_items(value: Any) -> list[Any]: + if value is None: + return [] + if isinstance(value, torch.Tensor): + if value.ndim == 0: + return [value] + return [value[i] for i in range(value.shape[0])] + if isinstance(value, list): + return value + if isinstance(value, tuple): + if len(value) == 2 and isinstance(value[1], (int, float)): + return [value] + return list(value) + return [value] + + def _default_mm_device(self) -> torch.device: + try: + return next(self.model.parameters()).device + except StopIteration: + return torch.device("cpu") + + @staticmethod + def _coerce_mm_item_to_float_tensor( + item: Any, + *, + device: torch.device, + ) -> torch.Tensor: + if isinstance(item, tuple) and len(item) == 2 and isinstance(item[1], (int, float)): + item = item[0] + + if isinstance(item, torch.Tensor): + tensor = item.detach().to(device=device, dtype=torch.float32) + else: + tensor = torch.as_tensor(item, dtype=torch.float32, device=device) + + return tensor.contiguous() + + def _build_deterministic_mm_embedding( + self, + item: Any, + *, + device: torch.device, + ) -> torch.Tensor: + tensor = self._coerce_mm_item_to_float_tensor(item, device=device) + if tensor.numel() == 0: + return torch.zeros((1, self.hidden_size), dtype=torch.bfloat16, device=device) + + flattened = tensor.view(-1) + first = flattened[0] + last = flattened[-1] + mean = flattened.mean() + std = flattened.std(unbiased=False) + abs_mean = flattened.abs().mean() + max_abs = flattened.abs().max() + l2 = torch.linalg.vector_norm(flattened) / max(float(flattened.numel()), 1.0) + + base = torch.stack([first, last, mean, std, abs_mean, max_abs, l2], dim=0) + denom = torch.clamp(base.abs().max(), min=1.0) + base = base / denom + + repeats = (self.hidden_size + base.numel() - 1) // base.numel() + embedding = base.repeat(repeats)[: self.hidden_size].to(dtype=torch.bfloat16) + return embedding.unsqueeze(0).contiguous() + + def _collect_mm_inputs(self, **kwargs: Any) -> dict[str, Any]: + mm_inputs: dict[str, Any] = {} + for modality, aliases in MM_INPUT_ALIASES.items(): + for alias in aliases: + if alias in kwargs and kwargs[alias] is not None: + mm_inputs[modality] = kwargs[alias] + break + for modality, value in self._cached_mm_inputs.items(): + if modality not in mm_inputs and value is not None: + mm_inputs[modality] = value + return mm_inputs + + def embed_multimodal(self, **kwargs: Any) -> Any: + mm_inputs = self._collect_mm_inputs(**kwargs) + self._cached_mm_inputs = dict(mm_inputs) + if not mm_inputs: + return None + + device = self._default_mm_device() + mm_embeddings: list[torch.Tensor] = [] + + for modality in ("image", "video", "audio"): + value = mm_inputs.get(modality) + if value is None: + continue + for item in self._split_mm_items(value): + mm_embeddings.append(self._build_deterministic_mm_embedding(item, device=device)) + + return tuple(mm_embeddings) if mm_embeddings else None diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 139892345..3b51f2002 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/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", + ), } diff --git a/vllm_omni/model_executor/stage_configs/dynin_omni.yaml b/vllm_omni/model_executor/stage_configs/dynin_omni.yaml new file mode 100644 index 000000000..0724146aa --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml b/vllm_omni/model_executor/stage_configs/dynin_omni_multiconnector.yaml new file mode 100644 index 000000000..7259daa9e --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/dynin_omni.py b/vllm_omni/model_executor/stage_input_processors/dynin_omni.py new file mode 100644 index 000000000..9ec849799 --- /dev/null +++ b/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)