#!/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/deploy/dynin_omni.yaml"), help="Path to deploy 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()