diff --git a/README.md b/README.md index 66c8725ee..20abbb608 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ Traditional vLLM systems are limited to text-based, autoregressive generation. v - **Multi-modal Models**: Text, image, video, audio, and sensor data processing - **Non-autoregressive Architectures**: Diffusion Transformers (DiT) and other parallel generation models -- **Heterogeneous Outputs**: Beyond traditional text generation to structured, binary, and streaming outputs +- **Heterogeneous Outputs**: Beyond traditional text generation to multimodal outputs ## 🏗️ Architecture @@ -28,119 +28,48 @@ vLLM-omni is built on a modular architecture that extends vLLM's core functional - **Text**: Advanced tokenization and embedding generation - **Image**: Vision encoder integration (CLIP, etc.) - **Audio**: Speech processing and audio embedding -- **Video**: Frame-by-frame and temporal processing -- **Sensor**: IoT and sensor data interpretation - -### Output Formats - -- **Structured Data**: JSON, XML, and custom formats -- **Binary Outputs**: Images, audio, and video generation -- **Streaming**: Real-time progressive generation -- **Multipart**: Combined multi-modal responses ## 📋 Supported Models ### AR + Diffusion Transformer (DiT) Models -- Qwen-Image (Image generation and editing) - Qwen-omni (Thinker-Talker-Codec structure) -- Custom DiT and hiybrid architectures +- HunyunaImage 3.0 (Ongoing) +- Qwen-Image (Ongoing) ## 🛠️ Installation -### Quick Start - -#### Option 1: Docker (Recommended for macOS) +Set up basic environments +```bash +uv venv --python 3.12 --seed +source .venv/bin/activate +``` +Install certain version of vllm with commitid: 808a7b69df479b6b3a16181711cac7ca28a9b941 ```bash -# Clone the repository -git clone https://github.com/hsliuustc0106/vllm-omni.git -cd vllm-omni - -# Run the automated Docker setup -./scripts/docker-setup-macos.sh +git clone https://github.com/vllm-project/vllm.git +cd vllm +git checkout 808a7b69df479b6b3a16181711cac7ca28a9b941 +VLLM_USE_PRECOMPILED=1 uv pip install --editable . ``` -#### Option 2: Local Installation +## Run examples (Qwen2.5-omni) +Get into the example folder ```bash -# Clone the repository -git clone https://github.com/hsliuustc0106/vllm-omni.git -cd vllm-omni - -# Run the installation script -./install.sh +cd vllm_omni +cd examples/offline_inference/qwen2_5_omni ``` - -### Prerequisites - -- Python 3.11+ (recommended) -- Conda or Miniconda -- Git -- CUDA 11.8+ (for GPU acceleration) or CPU-only installation - -### Installation Methods - -#### Method 1: Automated Installation (Recommended) +Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run. ```bash -# Using shell script -./install.sh - -# Or using Python script -python install.py +bash run.sh ``` +The output audio is saved in ./output_audio -#### Method 2: Manual Installation -```bash -# Create conda environment -conda create -n vllm_omni python=3.11 -y -conda activate vllm_omni +## To-do list +- [x] Offline inference example for Qwen2.5-omni with single request +- [ ] Adaptation from current vllm branch to stable vllm v0.11.0 +- [ ] Offline inference example for Qwen2.5-omni with streaming multiple requests +- [ ] Online inference support +- [ ] Support for other models -# Install PyTorch (CPU or GPU) -pip install torch>=2.7 --index-url https://download.pytorch.org/whl/cpu # CPU -# pip install torch>=2.7 --index-url https://download.pytorch.org/whl/cu121 # GPU - -# Install dependencies -pip install -r requirements.txt -pip install "vllm>=0.10.2" - -# Install vLLM-omni -pip install -e . -``` - -### Verify Installation - -```bash -# Test the installation -python test_installation.py - -# Test basic functionality -python -c "import vllm_omni; print('Ready!')" - -# Test CLI -vllm --help -``` - -For detailed installation instructions, see [INSTALL.md](INSTALL.md). - -## 📥 Model Download - -Models are automatically downloaded when first used, or you can pre-download them: - -```bash -# Check downloaded models -python scripts/download_models.py --check-cache - -# Download all default models -python scripts/download_models.py --all - -# Download specific models -python scripts/download_models.py --ar-models Qwen/Qwen3-0.6B -python scripts/download_models.py --dit-models stabilityai/stable-diffusion-2-1 -``` - -**Model Storage Location:** -- Default: `~/.cache/huggingface/hub/` -- AR models: 100MB - 1GB each -- DiT models: 2GB - 7GB each - -For detailed model management, see [MODEL_DOWNLOAD_GUIDE.md](docs/MODEL_DOWNLOAD_GUIDE.md). +For detailed model management, see [vllm_omni_design.md](docs/architecture/vllm_omni_design.md) and [high_level_arch_design.md](docs/architecture/high_level_arch_design.md). \ No newline at end of file diff --git a/examples/offline_inference/qwen_2_5_omni/README.md b/examples/offline_inference/qwen_2_5_omni/README.md new file mode 100644 index 000000000..d1dfde205 --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/README.md @@ -0,0 +1,37 @@ +# Offline Example of vLLM-omni for Qwen2.5-omni + +## Installation + +Set up basic environments +```bash +uv venv --python 3.12 --seed +source .venv/bin/activate +``` +Install certain version of vllm with commitid: 808a7b69df479b6b3a16181711cac7ca28a9b941 + +```bash +git clone https://github.com/vllm-project/vllm.git +cd vllm +git checkout 808a7b69df479b6b3a16181711cac7ca28a9b941 +VLLM_USE_PRECOMPILED=1 uv pip install --editable . +``` + +## Run examples + +Get into the example folder +```bash +cd vllm_omni +cd examples/offline_inference/qwen2_5_omni +``` +Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run. +```bash +bash run.sh +``` +The output audio is saved in ./output_audio + +## To-do list +- [x] Offline inference example for Qwen2.5-omni with single request +- [ ] Adaptation from current vllm branch to stable vllm v0.11.0 +- [ ] Offline inference example for Qwen2.5-omni with streaming multiple requests +- [ ] Online inference support +- [ ] Support for other models \ No newline at end of file diff --git a/examples/offline_inference/qwen_2_5_omni/end2end.py b/examples/offline_inference/qwen_2_5_omni/end2end.py new file mode 100644 index 000000000..030db0b71 --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/end2end.py @@ -0,0 +1,130 @@ +import argparse +import os +import soundfile as sf +import random +import numpy as np +import torch + +from vllm.sampling_params import SamplingParams + +import os as _os_env_toggle +_os_env_toggle.environ["VLLM_USE_V1"] = "1" + +from vllm_omni.entrypoints.omni_llm import OmniLLM +from utils import make_omni_prompt + + +SEED = 42 +# Set all random seeds +random.seed(SEED) +np.random.seed(SEED) +torch.manual_seed(SEED) +torch.cuda.manual_seed(SEED) +torch.cuda.manual_seed_all(SEED) + +# Make PyTorch deterministic +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = False + +# Set environment variables for deterministic behavior +os.environ["PYTHONHASHSEED"] = str(SEED) +os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--model', required=True, help='Path to merged model directory (will be created if downloading).') + parser.add_argument('--thinker-model', type=str, default=None) + parser.add_argument('--talker-model', type=str, default=None) + parser.add_argument('--code2wav-model', type=str, default=None) + parser.add_argument('--hf-hub-id', default='Qwen/Qwen2.5-Omni-7B', help='Hugging Face repo id to download if needed.') + parser.add_argument('--hf-revision', default=None, help='Optional HF revision (branch/tag/commit).') + parser.add_argument('--prompts', required=True, nargs='+', help='Input text prompts.') + parser.add_argument('--voice-type', default='default', help='Voice type, e.g., m02, f030, default.') + parser.add_argument('--code2wav-dir', default=None, help='Path to code2wav folder (contains spk_dict.pt).') + parser.add_argument('--dit-ckpt', default=None, help='Path to DiT checkpoint file (e.g., dit.pt).') + parser.add_argument('--bigvgan-ckpt', default=None, help='Path to BigVGAN checkpoint file.') + parser.add_argument('--dtype', default='bfloat16', choices=['float16', 'bfloat16', 'float32']) + parser.add_argument('--max-model-len', type=int, default=32768) + + parser.add_argument("--thinker-only", action="store_true") + parser.add_argument("--text-only", action="store_true") + parser.add_argument("--do-wave", action="store_true") + parser.add_argument('--prompt_type', + choices=[ + 'text', 'audio', 'audio-long', 'audio-long-chunks', + 'audio-long-expand-chunks', 'image', 'video', + 'video-frames', 'audio-in-video', 'audio-in-video-v2', + "audio-multi-round", "badcase-vl", "badcase-text", + "badcase-image-early-stop", "badcase-two-audios", + "badcase-two-videos", "badcase-multi-round", + "badcase-voice-type", "badcase-voice-type-v2", + "badcase-audio-tower-1", "badcase-audio-only" + ], + default='text') + parser.add_argument('--use-torchvision', action='store_true') + parser.add_argument('--tokenize', action='store_true') + parser.add_argument('--output-wav', default="output.wav", help='Output wav file path.') + parser.add_argument('--thinker-hidden-states-dir', default="thinker_hidden_states", help='Path to thinker hidden states directory.') + args = parser.parse_args() + return args + + +def main(): + args = parse_args() + model_name = args.model + omni_llm = OmniLLM(model=model_name) + thinker_sampling_params = SamplingParams( + temperature=0.0, # Deterministic - no randomness + top_p=1.0, # Disable nucleus sampling + top_k=-1, # Disable top-k sampling + max_tokens=2048, + seed=SEED, # Fixed seed for sampling + detokenize=True, + repetition_penalty=1.1, + ) + talker_sampling_params = SamplingParams( + temperature=0.0, # Deterministic - no randomness + top_p=1.0, # Disable nucleus sampling + top_k=-1, # Disable top-k sampling + max_tokens=2048, + seed=SEED, # Fixed seed for sampling + detokenize=True, + repetition_penalty=1.1, + stop_token_ids=[8294] + ) + code2wav_sampling_params = SamplingParams( + temperature=0.0, # Deterministic - no randomness + top_p=1.0, # Disable nucleus sampling + top_k=-1, # Disable top-k sampling + max_tokens=2048, + seed=SEED, # Fixed seed for sampling + detokenize=True, + repetition_penalty=1.1, + ) + + sampling_params_list = [thinker_sampling_params, + talker_sampling_params, + code2wav_sampling_params] + + prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts] + omni_outputs = omni_llm.generate(prompt, sampling_params_list) + + os.makedirs(args.output_wav, exist_ok=True) + for stage_outputs in omni_outputs: + if stage_outputs.final_output_type == "text": + for output in stage_outputs.request_output: + request_id = output.request_id + text_output = output.outputs[0].text + print(f"Request ID: {request_id}, Text Output: {text_output}") + elif stage_outputs.final_output_type == "audio": + for output in stage_outputs.request_output: + request_id = output.request_id + audio_tensor = output.multimodal_output["audio"] + output_wav = os.path.join(args.output_wav, f"output_{output.request_id}.wav") + sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000) + print(f"Request ID: {request_id}, Saved audio to {output_wav}") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/offline_inference/qwen_2_5_omni/output_audio/output_0.wav b/examples/offline_inference/qwen_2_5_omni/output_audio/output_0.wav new file mode 100644 index 000000000..45caf1d0e Binary files /dev/null and b/examples/offline_inference/qwen_2_5_omni/output_audio/output_0.wav differ diff --git a/examples/offline_inference/qwen_2_5_omni/processing_omni.py b/examples/offline_inference/qwen_2_5_omni/processing_omni.py new file mode 100644 index 000000000..bd6f4ae95 --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/processing_omni.py @@ -0,0 +1,373 @@ +from __future__ import annotations + +import base64 +import logging +import math +import os +import time +import warnings +from functools import lru_cache +from io import BytesIO +import requests +import torch +import torchvision +from packaging import version +from PIL import Image +from torchvision import io, transforms +from torchvision.transforms import InterpolationMode + + +logger = logging.getLogger(__name__) + +IMAGE_FACTOR = 28 +MIN_PIXELS = 4 * 28 * 28 +MAX_PIXELS = 16384 * 28 * 28 +MAX_RATIO = 200 + +VIDEO_MIN_PIXELS = 128 * 28 * 28 +VIDEO_MAX_PIXELS = 768 * 28 * 28 +VIDEO_TOTAL_PIXELS = 24576 * 28 * 28 +FRAME_FACTOR = 2 +FPS = 2.0 +FPS_MIN_FRAMES = 4 +FPS_MAX_FRAMES = 768 + +temporal_patch_size = 2 +spatial_patch_size = 14 +spatial_merge_size = 2 + + +def round_by_factor(number: int, factor: int) -> int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def ceil_by_factor(number: int, factor: int) -> int: + """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.""" + return math.ceil(number / factor) * factor + + +def floor_by_factor(number: int, factor: int) -> int: + """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.""" + return math.floor(number / factor) * factor + + +def smart_resize(height: int, + width: int, + factor: int = IMAGE_FACTOR, + min_pixels: int = MIN_PIXELS, + max_pixels: int = MAX_PIXELS) -> tuple[int, int]: + """ + Rescales the image so that the following conditions are met: + + 1. Both dimensions (height and width) are divisible by 'factor'. + + 2. The total number of pixels is within the range ['min_pixels', 'max_pixels']. + + 3. The aspect ratio of the image is maintained as closely as possible. + """ + if max(height, width) / min(height, width) > MAX_RATIO: + raise ValueError( + f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}" + ) + h_bar = max(factor, round_by_factor(height, factor)) + w_bar = max(factor, round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = floor_by_factor(height / beta, factor) + w_bar = floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = ceil_by_factor(height * beta, factor) + w_bar = ceil_by_factor(width * beta, factor) + return h_bar, w_bar + + +def fetch_image(ele: dict[str, str | Image.Image], + size_factor: int = IMAGE_FACTOR) -> Image.Image: + if "image" in ele: + image = ele["image"] + else: + image = ele["image_url"] + image_obj = None + if isinstance(image, Image.Image): + image_obj = image + elif image.startswith("http://") or image.startswith("https://"): + image_obj = Image.open(requests.get(image, stream=True).raw) + elif image.startswith("file://"): + image_obj = Image.open(image[7:]) + elif image.startswith("data:image"): + if "base64," in image: + _, base64_data = image.split("base64,", 1) + data = base64.b64decode(base64_data) + image_obj = Image.open(BytesIO(data)) + else: + image_obj = Image.open(image) + if image_obj is None: + raise ValueError( + f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}" + ) + image = image_obj.convert("RGB") + ## resize + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=size_factor, + ) + else: + width, height = image.size + min_pixels = ele.get("min_pixels", MIN_PIXELS) + max_pixels = ele.get("max_pixels", MAX_PIXELS) + resized_height, resized_width = smart_resize( + height, + width, + factor=size_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + image = image.resize((resized_width, resized_height)) + + return image + + +def smart_nframes( + ele: dict, + total_frames: int, + video_fps: int | float, +) -> int: + """calculate the number of frames for video used for model inputs. + + Args: + ele (dict): a dict contains the configuration of video. + support either `fps` or `nframes`: + - nframes: the number of frames to extract for model inputs. + - fps: the fps to extract frames for model inputs. + - min_frames: the minimum number of frames of the video, only used when fps is provided. + - max_frames: the maximum number of frames of the video, only used when fps is provided. + total_frames (int): the original total number of frames of the video. + video_fps (int | float): the original fps of the video. + + Raises: + ValueError: nframes should in interval [FRAME_FACTOR, total_frames]. + + Returns: + int: the number of frames for video used for model inputs. + """ + assert not ("fps" in ele + and "nframes" in ele), "Only accept either `fps` or `nframes`" + if "nframes" in ele: + nframes = round_by_factor(ele["nframes"], FRAME_FACTOR) + else: + fps = ele.get("fps", FPS) + min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), + FRAME_FACTOR) + max_frames = floor_by_factor( + ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), + FRAME_FACTOR) + nframes = total_frames / video_fps * fps + nframes = min(max(nframes, min_frames), max_frames) + nframes = round_by_factor(nframes, FRAME_FACTOR) + if not (FRAME_FACTOR <= nframes and nframes <= total_frames): + raise ValueError( + f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}." + ) + return nframes + + +def _read_video_torchvision(ele: dict, ) -> torch.Tensor: + """read video using torchvision.io.read_video + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + video_path = ele["video"] + if version.parse(torchvision.__version__) < version.parse("0.19.0"): + if "http://" in video_path or "https://" in video_path: + warnings.warn( + "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0." + ) + if "file://" in video_path: + video_path = video_path[7:] + st = time.time() + video, audio, info = io.read_video( + video_path, + start_pts=ele.get("video_start", 0.0), + end_pts=ele.get("video_end", None), + pts_unit="sec", + output_format="TCHW", + ) + total_frames, video_fps = video.size(0), info["video_fps"] + total_duration = round(total_frames / video_fps, 3) + logger.info( + f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s" + ) + nframes = smart_nframes(ele, + total_frames=total_frames, + video_fps=video_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long() + video = video[idx] + return video, total_duration, nframes + + +def is_decord_available() -> bool: + import importlib.util + + return importlib.util.find_spec("decord") is not None + + +def _read_video_decord(ele: dict, ) -> torch.Tensor: + """read video using decord.VideoReader + + Args: + ele (dict): a dict contains the configuration of video. + support keys: + - video: the path of video. support "file://", "http://", "https://" and local path. + - video_start: the start time of video. + - video_end: the end time of video. + Returns: + torch.Tensor: the video tensor with shape (T, C, H, W). + """ + import decord + video_path = ele["video"] + st = time.time() + vr = decord.VideoReader(video_path) + # TODO: support start_pts and end_pts + if 'video_start' in ele or 'video_end' in ele: + raise NotImplementedError( + "not support start_pts and end_pts in decord for now.") + total_frames, video_fps = len(vr), vr.get_avg_fps() + total_duration = round(total_frames / video_fps, 3) + logger.info( + f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s" + ) + nframes = smart_nframes(ele, + total_frames=total_frames, + video_fps=video_fps) + idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() + video = vr.get_batch(idx).asnumpy() + video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format + return video, total_duration, nframes + + +VIDEO_READER_BACKENDS = { + "decord": _read_video_decord, + "torchvision": _read_video_torchvision, +} + +FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None) + + +@lru_cache(maxsize=1) +def get_video_reader_backend() -> str: + if FORCE_QWENVL_VIDEO_READER is not None: + video_reader_backend = FORCE_QWENVL_VIDEO_READER + elif is_decord_available(): + video_reader_backend = "decord" + else: + video_reader_backend = "torchvision" + # print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr) + return video_reader_backend + + +def fetch_video( + ele: dict, + image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]: + if isinstance(ele["video"], str): + video_reader_backend = get_video_reader_backend() + video, total_dur, nframes = VIDEO_READER_BACKENDS[ + video_reader_backend](ele) + frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes + grid_timestamps = frame_timestamps[::FRAME_FACTOR] + second_per_grid = grid_timestamps[1] - grid_timestamps[0] + nframes, _, height, width = video.shape + factor = spatial_patch_size * spatial_merge_size + min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS) + total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS) + max_pixels = max( + min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), + int(min_pixels * 1.05)) + max_pixels = ele.get("max_pixels", max_pixels) + # min_pixels = (factor ** 2) * 52 + # max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size)) + if "resized_height" in ele and "resized_width" in ele: + resized_height, resized_width = smart_resize( + ele["resized_height"], + ele["resized_width"], + factor=image_factor, + ) + else: + resized_height, resized_width = smart_resize( + height, + width, + factor=image_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + video = transforms.functional.resize( + video, + [resized_height, resized_width], + interpolation=InterpolationMode.BICUBIC, + antialias=True, + ).float() + return video, total_dur, nframes, second_per_grid + else: + assert isinstance(ele["video"], (list, tuple)) + process_info = ele.copy() + process_info.pop("type", None) + process_info.pop("video", None) + images = [ + fetch_image({ + "image": video_element, + **process_info + }, + size_factor=image_factor) + for video_element in ele["video"] + ] + nframes = ceil_by_factor(len(images), FRAME_FACTOR) + if len(images) < nframes: + images.extend([images[-1]] * (nframes - len(images))) + return images, None, None, None + + +def extract_vision_info( + conversations: list[dict] | list[list[dict]]) -> list[dict]: + vision_infos = [] + if isinstance(conversations[0], dict): + conversations = [conversations] + for conversation in conversations: + for message in conversation: + if isinstance(message["content"], list): + for ele in message["content"]: + if ("image" in ele or "image_url" in ele or "video" in ele + or ele["type"] in ("image", "image_url", "video")): + vision_infos.append(ele) + return vision_infos + + +def process_vision_info( + conversations: list[dict] | list[list[dict]], +) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] + | None]: + vision_infos = extract_vision_info(conversations) + ## Read images or videos + image_inputs = [] + video_inputs = [] + for vision_info in vision_infos: + if "image" in vision_info or "image_url" in vision_info: + image_inputs.append(fetch_image(vision_info)) + elif "video" in vision_info: + video_inputs.append(fetch_video(vision_info)) + else: + raise ValueError("image, image_url or video should in content.") + if len(image_inputs) == 0: + image_inputs = None + if len(video_inputs) == 0: + video_inputs = None + return image_inputs, video_inputs \ No newline at end of file diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run.sh new file mode 100644 index 000000000..2b5db45cb --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/run.sh @@ -0,0 +1,9 @@ +export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH +export HF_ENDPOINT=https://hf-mirror.com +python end2end.py --model Qwen/Qwen2.5-Omni-7B \ + --prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." \ + --voice-type "m02" \ + --dit-ckpt none \ + --bigvgan-ckpt none \ + --output-wav output_audio \ + --prompt_type text \ No newline at end of file diff --git a/examples/offline_inference/qwen_2_5_omni/utils.py b/examples/offline_inference/qwen_2_5_omni/utils.py new file mode 100644 index 000000000..c8cf5392d --- /dev/null +++ b/examples/offline_inference/qwen_2_5_omni/utils.py @@ -0,0 +1,305 @@ +import tempfile +from urllib.request import urlopen +import librosa +import soundfile as sf +import resampy +from typing import Dict, Optional +import torch +import requests +import torchvision.io + +from typing import Union, List +from vllm.inputs import TextPrompt +from vllm_omni.inputs.data import OmniTokensPrompt +from processing_omni import fetch_image, fetch_video + + +def get_system_prompt(): + + return { + 'role': + 'system', + 'content': [{ + 'type': + 'text', + 'text': + 'You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.' + }] + } + + +def resample_wav_to_16khz(input_filepath): + data, original_sample_rate = sf.read(input_filepath) + # Only use the first channel + if len(data.shape) > 1: + data = data[:, 0] + # resample to 16kHz + data_resampled = resampy.resample(data, + sr_orig=original_sample_rate, + sr_new=16000) + return data_resampled + + +def fetch_and_read_video(args, video_url: str, fps=2): + + def read_video_with_torchvision(video_file_name: str): + video, audio, info = torchvision.io.read_video( + video_file_name, + start_pts=0.0, + end_pts=None, + pts_unit="sec", + output_format="TCHW", + ) + + total_frames, video_fps = video.size(0), info["video_fps"] + total_duration = round(total_frames / video_fps, 3) + nframes = int(total_frames / video_fps * fps) + + frame_timestamps = total_duration * torch.arange(1, + nframes + 1) / nframes + grid_timestamps = frame_timestamps[::2] + second_per_grid = grid_timestamps[1] - grid_timestamps[0] + + idx = torch.linspace(0, video.size(0) - 1, nframes).round().long() + video = video[idx] + + if args.legacy_omni_video: + return [video, total_duration, nframes, second_per_grid.item()] + else: + return video + + def read_video_with_transformers(video_file_name: Union[str, List[str]]): + video, total_duration, nframes, second_per_grid = fetch_video( + {'video': video_file_name}) + if total_duration is None and nframes is None: + nframes = len(video) + total_duration = 0.5 * nframes + second_per_grid = 1.0 + if args.legacy_omni_video: + return [video, total_duration, nframes, second_per_grid] + else: + return video + + def read_video(video_file_name: str): + if args.use_torchvision: + return read_video_with_torchvision(video_file_name) + else: + return read_video_with_transformers(video_file_name) + + if isinstance(video_url, str) and video_url.startswith("http"): + with tempfile.NamedTemporaryFile(delete=True) as temp_video_file: + resp = requests.get(video_url) + assert resp.status_code == requests.codes.ok, f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}" + + temp_video_file.write(urlopen(video_url).read()) + temp_video_file_path = temp_video_file.name + video_file_name = temp_video_file_path + return read_video(video_file_name) + else: + video_file_name = video_url + return read_video(video_file_name) + + +def make_inputs_qwen2_omni( + args, + messages: List[Dict[str, Union[str, List[Dict[str, str]]]]], + use_audio_in_video: Optional[bool] = False, + tokenize: bool = False, +) -> Union[OmniTokensPrompt, TextPrompt]: + + from transformers import AutoConfig, AutoProcessor, AutoTokenizer + processor = AutoProcessor.from_pretrained(args.model) + tokenizer = AutoTokenizer.from_pretrained(args.model) + + try: + config = AutoConfig.from_pretrained(args.model) + if 'Qwen2_5OmniModel' in config.architectures: + args.legacy_omni_video = False + else: + args.legacy_omni_video = True + except: + args.legacy_omni_video = True + + audios, images, videos = [], [], [] + for message in messages: + if not isinstance(message['content'], list): + message['content'] = [{ + 'type': 'text', + 'text': message['content'], + }] + index, num_contents = 0, len(message['content']) + while index < num_contents: + ele = message['content'][index] + if 'type' not in ele: + if 'text' in ele: + ele['type'] = 'text' + elif 'audio' in ele: + ele['type'] = 'audio' + elif 'audio_url' in ele: + ele['type'] = 'audio_url' + elif 'image' in ele: + ele['type'] = 'image' + elif 'image_url' in ele: + ele['type'] = 'image_url' + elif 'video' in ele: + ele['type'] = 'video' + elif 'video_url' in ele: + ele['type'] = 'video_url' + else: + raise ValueError(f'Unknown ele: {ele}') + + if ele['type'] == 'audio' or ele['type'] == 'audio_url': + if 'audio_url' in ele: + audio_key = 'audio_url' + with tempfile.NamedTemporaryFile( + delete=True) as temp_audio_file: + temp_audio_file.write(urlopen(ele[audio_key]).read()) + temp_audio_file_path = temp_audio_file.name + audios.append( + resample_wav_to_16khz(temp_audio_file_path)) + ele['audio'] = temp_audio_file_path + elif 'audio' in ele: + audio_key = 'audio' + audios.append(resample_wav_to_16khz(ele[audio_key])) + else: + raise ValueError(f'Unknown ele {ele}') + elif use_audio_in_video and (ele['type'] == 'video' + or ele['type'] == 'video_url'): + # use video as audio as well + if 'video_url' in ele: + audio_key = 'video_url' + with tempfile.NamedTemporaryFile( + delete=True) as temp_video_file: + temp_video_file.write(urlopen(ele[audio_key]).read()) + temp_video_file_path = temp_video_file.name + ele[audio_key] = temp_video_file_path + audios.append( + librosa.load(temp_video_file_path, sr=16000)[0]) + videos.append( + fetch_and_read_video(args, temp_video_file_path)) + ele['video'] = temp_video_file_path + elif 'video' in ele: + audio_key = 'video' + audios.append(librosa.load(ele[audio_key], sr=16000)[0]) + videos.append(fetch_and_read_video(args, audio_key)) + else: + raise ValueError("Unknown ele {}".format(ele)) + # insert a audio after the video + message['content'].insert(index + 1, { + "type": "audio", + "audio": ele[audio_key], + }) + # no need to load the added audio again + index += 1 + elif ele['type'] == 'video' or ele['type'] == 'video_url': + if 'video_url' in ele: + video_key = 'video_url' + with tempfile.NamedTemporaryFile( + delete=True) as temp_video_file: + temp_video_file.write(urlopen(ele['video_url']).read()) + temp_video_file_path = temp_video_file.name + videos.append(fetch_and_read_video(args, temp_video_file)) + ele['video'] = temp_video_file_path + else: + video_key = 'video' + videos.append(fetch_and_read_video(args, ele[video_key])) + elif ele['type'] == 'image' or ele['type'] == 'image_url': + images.append(fetch_image(ele)) + + # move to the next content + index += 1 + + prompt = processor.apply_chat_template( + messages, + tokenize=tokenize, + add_generation_prompt=True, + add_vision_id=True, + ) + + audios = audios if len(audios) > 0 else None + images = images if len(images) > 0 else None + videos = videos if len(videos) > 0 else None + + multi_modal_data = {} + if audios: + multi_modal_data["audio"] = audios + if images: + multi_modal_data["image"] = images + if videos: + multi_modal_data["video"] = videos + + if isinstance(prompt, list) and isinstance(prompt[0], (list, str)): + prompt = prompt[0] + + if tokenize: + return OmniTokensPrompt( + prompt_token_ids=prompt, + multi_modal_data=multi_modal_data, + ) + else: + return TextPrompt( + prompt=prompt, + multi_modal_data=multi_modal_data, + ) + + +def make_text_prompt(args, prompt): + messages = [ + get_system_prompt(), + { + "role": "user", + "content": [ + { + "type": "text", + "text": prompt + }, + ] + }, + ] + + prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize) + return prompt + + +def make_audio_in_video_v2_prompt(args): + messages = [ + { + 'role': + 'system', + 'content': [{ + 'type': + 'text', + 'text': + 'You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech.' + }] + }, + { + "role": + "user", + "content": [ + { + "type": + "video_url", + "video_url": + "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4" + }, + ] + }, + ] + prompt = make_inputs_qwen2_omni( + args, + messages, + use_audio_in_video=True, + tokenize=args.tokenize, + ) + return prompt + + +def make_omni_prompt(args, prompt = None) -> Union[OmniTokensPrompt, List[OmniTokensPrompt]]: + if args.prompt_type == 'text': + prompt = make_text_prompt(args, prompt) + elif args.prompt_type == 'audio-in-video-v2': + prompt = make_audio_in_video_v2_prompt(args) + else: + raise ValueError(f'Unsupported prompt type: {args.prompt_type}') + return prompt \ No newline at end of file diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index e48a874a7..a688faf1b 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -1,6 +1,14 @@ """ Configuration module for vLLM-omni. """ +from typing import Optional +from pydantic.dataclasses import dataclass +from pydantic import ConfigDict + +from vllm.config import ModelConfig +from vllm.config import config + +import vllm_omni.model_executor.models as me_models from .stage_config import ( OmniStageConfig, @@ -11,7 +19,28 @@ from .stage_config import ( create_dit_stage_config, ) + +@config +@dataclass(config=ConfigDict(arbitrary_types_allowed=True)) +class OmniModelConfig(ModelConfig): + """Configuration for Omni models, extending the base ModelConfig.""" + + stage_id: int = 0 + model_stage: str = "thinker" + model_arch: str = "Qwen2_5OmniForConditionalGeneration" + engine_output_type: Optional[str] = None + + @property + def registry(self): + return me_models.OmniModelRegistry + + @property + def architectures(self) -> list[str]: + return [self.model_arch] + + __all__ = [ + "OmniModelConfig", "OmniStageConfig", "DiTConfig", "DiTCacheConfig", diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 8f5fcb357..211db8545 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -1,215 +1,411 @@ -""" -Output processing for multimodal outputs in vLLM-omni. -""" +from typing import Dict, Callable, Optional, Any, Union -from typing import List, Dict, Any, Optional, Callable, Union -from vllm.outputs import RequestOutput, CompletionOutput -from vllm.v1.outputs import ModelRunnerOutput as EngineCoreOutput +import torch + +from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor +from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestState, RequestOutputCollector +from vllm.transformers_utils.tokenizer_group import TokenizerGroup +from vllm.v1.engine import FinishReason +from vllm.v1.metrics.stats import IterationStats +from vllm.sampling_params import RequestOutputKind +from vllm.transformers_utils.tokenizer import AnyTokenizer +from vllm.v1.engine import EngineCoreRequest +from vllm.v1.engine.detokenizer import IncrementalDetokenizer +from vllm.v1.engine.logprobs import LogprobsProcessor +from vllm.v1.engine.parallel_sampling import ParentRequest +from vllm.logger import init_logger +from vllm.v1.engine import EngineCoreOutput -class MultimodalOutputProcessor: - """Handles multimodal output processing for vLLM-omni.""" - - def __init__(self): - self.output_handlers: Dict[str, Callable] = { - "image": self._process_image_output, - "text+image": self._process_text_image_output, - "latents": self._process_latents_output, - "text": self._process_text_output, - "pooling": self._process_pooling_output, - } - - def process_output(self, engine_core_output: Any) -> List[RequestOutput]: - """Process engine core output and return formatted RequestOutput.""" - if engine_core_output is None: - return [] - - # If it's already a RequestOutput, return as is - if isinstance(engine_core_output, RequestOutput): - return [engine_core_output] - - # If it's a list of RequestOutputs, return as is - if isinstance(engine_core_output, list): - return engine_core_output - - # Otherwise, process based on output type - output_type = self._detect_output_type(engine_core_output) - handler = self.output_handlers.get(output_type, self._process_pooling_output) - - return handler(engine_core_output) +logger = init_logger(__name__) - def _build_request_output( + +class OmniRequestState(RequestState): + + def __init__( self, - source: Any, - completion_outputs: List[CompletionOutput], - ) -> RequestOutput: - """Helper to construct RequestOutput with safe defaults.""" + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.mm_type: Optional[str] = None + self.mm_accumulated: Optional[torch.Tensor] = None - return RequestOutput( - request_id=getattr(source, "request_id", "unknown"), - prompt=getattr(source, "prompt", ""), - prompt_token_ids=getattr(source, "prompt_token_ids", []), - prompt_logprobs=getattr(source, "prompt_logprobs", None), - outputs=completion_outputs, - finished=getattr(source, "finished", True), + @classmethod + def from_new_request( + cls, + tokenizer: AnyTokenizer, + request: EngineCoreRequest, + prompt: Optional[str], + parent_req: Optional[ParentRequest], + request_index: int, + queue: Optional[Any], + log_stats: bool, + ) -> "OmniRequestState": + + if sampling_params := request.sampling_params: + if not sampling_params.detokenize: + tokenizer = None + output_kind = sampling_params.output_kind + logprobs_processor = LogprobsProcessor.from_new_request( + tokenizer=tokenizer, + request=request, + ) + detokenizer = IncrementalDetokenizer.from_new_request( + tokenizer=tokenizer, + request=request, + ) + max_tokens_param = sampling_params.max_tokens + else: + logprobs_processor = None + detokenizer = None + max_tokens_param = None + assert request.pooling_params is not None + output_kind = request.pooling_params.output_kind + + return cls( + request_id=request.request_id, + parent_req=parent_req, + request_index=request_index, + lora_name=(request.lora_request.name + if request.lora_request is not None else None), + output_kind=output_kind, + prompt=prompt, + prompt_token_ids=request.prompt_token_ids, + logprobs_processor=logprobs_processor, + detokenizer=detokenizer, + max_tokens_param=max_tokens_param, + arrival_time=request.arrival_time, + queue=queue, + log_stats=log_stats, ) - def _detect_output_type(self, output: Any) -> str: - """Detect the type of output based on its content.""" - if hasattr(output, 'output_type'): - return output.output_type - - # Check for image-related attributes - if hasattr(output, 'image') or hasattr(output, 'images'): - if hasattr(output, 'text') or hasattr(output, 'texts'): - return "text+image" + def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], + mm_type: Optional[str]) -> None: + if tensor is None: + return + try: + if mm_type: + self.mm_type = (mm_type or "").lower() + t = tensor.detach() + try: + t = t.to("cpu") + except Exception: + # Best-effort CPU move; keep original device if conversion fails + logger.debug("Failed to move multimodal tensor to CPU", exc_info=True) + if self.mm_accumulated is None: + self.mm_accumulated = t else: - return "image" - - # Check for latent-related attributes - if hasattr(output, 'latents') or hasattr(output, 'latent_representation'): - return "latents" - - # Check for pooling output - if hasattr(output, 'pooler_output') and output.pooler_output is not None: - return "pooling" - - # Default to text - return "text" - - def _process_text_output(self, output: Any) -> List[RequestOutput]: - """Process text output.""" - if isinstance(output, RequestOutput): - return [output] - - # Create a mock RequestOutput for text - completion_output = CompletionOutput( - index=0, - text=getattr(output, 'text', ''), - token_ids=getattr(output, 'token_ids', []), - cumulative_logprob=getattr(output, 'cumulative_logprob', 0.0), - logprobs=getattr(output, 'logprobs', None), - finish_reason=getattr(output, 'finish_reason', 'length') - ) - - return [self._build_request_output(output, [completion_output])] + self.mm_accumulated = torch.cat([self.mm_accumulated, t], dim=0) + except Exception: + # Log and continue without crashing the output pipeline + logger.exception("Error accumulating multimodal tensor") - def _process_image_output(self, output: Any) -> List[RequestOutput]: - """Process image output.""" - # For image outputs, we need to create a special RequestOutput - # that can handle image data + # Override: do not route to pooling-only path; always create completion + # outputs, and attach pooling_result into the CompletionOutput. + def make_request_output( + self, + new_token_ids: list[int], + pooling_output: Optional[torch.Tensor], + finish_reason: Optional[FinishReason], + stop_reason: Optional[Union[int, str]], + kv_transfer_params: Optional[dict[str, Any]] = None, + num_cached_tokens: Optional[int] = None, + ) -> Optional[Any]: + finished = finish_reason is not None + final_only = self.output_kind == RequestOutputKind.FINAL_ONLY + + if not finished and final_only: + return None + + if num_cached_tokens is not None: + # Keep num_cached_tokens in RequestOutput for compatibility + try: + self.num_cached_tokens = num_cached_tokens # type: ignore[attr-defined] + except Exception: + pass + + request_id = self.request_id + output = self._new_completion_output(new_token_ids, finish_reason, + stop_reason) + + if self.parent_req is None: + outputs = [output] + else: + request_id, outputs, finished = self.parent_req.get_outputs( + request_id, output) + if not outputs: + return None + + return self._new_request_output(request_id, outputs, finished, + kv_transfer_params) + + def _new_completion_output( + self, + token_ids: list[int], + finish_reason: Optional[FinishReason], + stop_reason: Optional[Union[int, str]] + ) -> Any: + # Reuse base text/logprobs logic, then annotate with pooling_result. + base_output = super()._new_completion_output(token_ids, finish_reason, + stop_reason) + try: + if self.mm_accumulated is not None: + tensor = self.mm_accumulated + try: + tensor = tensor.detach().to("cpu") + except Exception: + logger.debug("Failed to move accumulated multimodal tensor to CPU", exc_info=True) + # Attach on the completion output for downstream consumers. + if not hasattr(base_output, "multimodal_output"): + setattr(base_output, "multimodal_output", {}) + setattr(base_output, "multimodal_output", {self.mm_type: tensor}) + except Exception: + logger.exception("Error in _new_completion_output") + return base_output + + +class MultimodalOutputProcessor(VLLMOutputProcessor): + """Handles multimodal output processing by normalizing EngineCoreOutput + before delegating to the base vLLM OutputProcessor. + + Strategy: + - Route by EngineCoreOutput.output_type when present + ("image", "text+image", "latents", "text"). + - Fallback to pooling/text heuristics when output_type is absent. + - Mutate EngineCoreOutput in-place to ensure vLLM's base processor can + produce the correct RequestOutput/PoolingRequestOutput. + - Allow custom per-modality handlers via register_handler(). + """ + def __init__(self, tokenizer: TokenizerGroup, log_stats: bool, engine_core_output_type: Optional[str] = None): + super().__init__(tokenizer=tokenizer, log_stats=log_stats) + self.output_handlers: Dict[str, Callable[[EngineCoreOutput], None]] = {} + self._reqid_to_mm_type: Dict[str, str] = {} + self.request_states: dict[str, OmniRequestState] = {} + self.engine_core_output_type = engine_core_output_type - # Extract image data - image_data = getattr(output, 'image', None) - if image_data is None: - image_data = getattr(output, 'images', [None])[0] - - # Create a completion output with image data - completion_output = CompletionOutput( - index=0, - text="", # No text for pure image output - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason="stop" + def register_handler(self, modality: str, + handler: Callable[[EngineCoreOutput], None]) -> None: + self.output_handlers[modality.lower()] = handler + + def add_request( + self, + request: EngineCoreRequest, + prompt: Optional[str], + parent_req: Optional[ParentRequest] = None, + request_index: int = 0, + queue: Optional[RequestOutputCollector] = None, + ) -> None: + request_id = request.request_id + if request_id in self.request_states: + raise ValueError(f"Request id {request_id} already running.") + + tokenizer = None if not self.tokenizer else \ + self.tokenizer.get_lora_tokenizer(request.lora_request) + + req_state = OmniRequestState.from_new_request(tokenizer=tokenizer, + request=request, + prompt=prompt, + parent_req=parent_req, + request_index=request_index, + queue=queue, + log_stats=self.log_stats) + self.request_states[request_id] = req_state + self.lora_states.add_request(req_state) + if parent_req: + self.parent_requests[parent_req.request_id] = parent_req + + def process_outputs( + self, + engine_core_outputs: list[EngineCoreOutput], + engine_core_timestamp: Optional[float] = None, + iteration_stats: Optional[IterationStats] = None, + ) -> OutputProcessorOutput: + self._reqid_to_mm_type.clear() + for eco in engine_core_outputs: + mm_type = (getattr(eco, "output_type", None) or "").lower() + if mm_type: + self._reqid_to_mm_type[eco.request_id] = mm_type + self._route_and_normalize(eco) + + # Build RequestOutputs without delegating to base, so we can keep ids + request_outputs: list[Any] = [] + reqs_to_abort: list[str] = [] + for eco in engine_core_outputs: + req_id = eco.request_id + req_state = self.request_states.get(req_id) + if req_state is None: + continue + + # 1) Stats + self._update_stats_from_output(req_state, eco, + engine_core_timestamp, + iteration_stats) + + new_token_ids = eco.new_token_ids + pooling_output = eco.pooling_output + finish_reason = eco.finish_reason + stop_reason = eco.stop_reason + kv_transfer_params = eco.kv_transfer_params + num_cached_tokens = eco.num_cached_tokens + req_state.is_prefilling = False + + # 2) Detokenize and logprobs when text path + assert req_state.detokenizer is not None + assert req_state.logprobs_processor is not None + stop_string = req_state.detokenizer.update( + new_token_ids, finish_reason == FinishReason.STOP) + if stop_string: + finish_reason = FinishReason.STOP + stop_reason = stop_string + req_state.logprobs_processor.update_from_output(eco) + + # 2.5) Accumulate multimodal tensors in RequestState + try: + mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower() + if pooling_output is not None and isinstance(req_state, OmniRequestState): + req_state.add_multimodal_tensor(pooling_output, mm_type) + except Exception: + logger.debug("Failed to accumulate multimodal tensor for request %s", req_id, exc_info=True) + + # 3) Create RequestOutput objects, forcing combined mode to keep ids + pooling_for_make = pooling_output + if pooling_output is not None and new_token_ids: + # Do not consume pooling path now; keep ids and attach mm later + pooling_for_make = None + + ro = req_state.make_request_output(new_token_ids, pooling_for_make, + finish_reason, stop_reason, + kv_transfer_params, + num_cached_tokens) + if ro: + # Attach accumulated multimodal payload if any + try: + if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None: + mm_key = req_state.mm_type or "latents" + if not hasattr(ro, "multimodal_output"): + setattr(ro, "multimodal_output", {}) + ro.multimodal_output[mm_key] = req_state.mm_accumulated + except Exception: + logger.exception("Error attaching multimodal payload in process_outputs") + if req_state.queue is not None: + req_state.queue.put(ro) + else: + request_outputs.append(ro) + + # 4) Free completed + if finish_reason is not None: + self.request_states.pop(req_id) + parent_req = req_state.parent_req + if parent_req and not parent_req.child_requests: + self.parent_requests.pop(parent_req.request_id, None) + if not eco.finished: + reqs_to_abort.append(req_id) + self._update_stats_from_finished(req_state, finish_reason, + iteration_stats) + # Cleanup per-request mm state + if isinstance(req_state, OmniRequestState): + req_state.mm_accumulated = None + req_state.mm_type = None + + return OutputProcessorOutput( + request_outputs=request_outputs, + reqs_to_abort=reqs_to_abort, ) - - # Add image data to the completion output - completion_output.image = image_data - - return [ - self._build_request_output(output, [completion_output]) - ] - def _process_text_image_output(self, output: Any) -> List[RequestOutput]: - """Process combined text and image output.""" - # Extract text and image data - text_data = getattr(output, 'text', '') - image_data = getattr(output, 'image', None) - - if image_data is None: - image_data = getattr(output, 'images', [None])[0] - - # Create a completion output with both text and image - completion_output = CompletionOutput( - index=0, - text=text_data, - token_ids=getattr(output, 'token_ids', []), - cumulative_logprob=getattr(output, 'cumulative_logprob', 0.0), - logprobs=getattr(output, 'logprobs', None), - finish_reason="stop" - ) - - # Add image data to the completion output - completion_output.image = image_data - - return [ - self._build_request_output(output, [completion_output]) - ] - - def _process_latents_output(self, output: Any) -> List[RequestOutput]: - """Process latent representation output.""" - # Extract latent data - latent_data = getattr(output, 'latents', None) - if latent_data is None: - latent_data = getattr(output, 'latent_representation', None) - - # Create a completion output with latent data - completion_output = CompletionOutput( - index=0, - text="", # No text for latent output - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason="stop" - ) - - # Add latent data to the completion output - completion_output.latents = latent_data - - return [ - self._build_request_output(output, [completion_output]) - ] - - def _process_pooling_output(self, output: Any) -> List[RequestOutput]: - """Process pooling output (hidden states, embeddings, etc.).""" - # Extract pooling data - pooling_data = getattr(output, 'pooler_output', None) - if pooling_data is None: - pooling_data = getattr(output, 'hidden_states', None) - - # Create a completion output with pooling data - completion_output = CompletionOutput( - index=0, - text="", # No text for pooling output - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason="stop" - ) - - # Add pooling data to the completion output - completion_output.pooler_output = pooling_data - - return [ - self._build_request_output(output, [completion_output]) - ] - - def process_outputs(self, engine_core_outputs: List[EngineCoreOutput], **kwargs) -> List[RequestOutput]: - """Process multiple engine core outputs.""" - all_outputs = [] - - for engine_core_output in engine_core_outputs: - outputs = self.process_output(engine_core_output) - all_outputs.extend(outputs) - - return all_outputs - - def add_output_handler(self, output_type: str, handler: Callable) -> None: - """Add a custom output handler for a specific output type.""" - self.output_handlers[output_type] = handler - - def remove_output_handler(self, output_type: str) -> None: - """Remove an output handler for a specific output type.""" + # ---- routing helpers ---- + def _route_and_normalize(self, eco: EngineCoreOutput) -> None: + output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower() + + # Custom handler first (if registered) if output_type in self.output_handlers: - del self.output_handlers[output_type] + try: + self.output_handlers[output_type](eco) + # Fall through to default fixups in case the handler left gaps + except Exception: + logger.exception("Error in custom output handler for %s", output_type) + + if output_type == "image": + self._process_image_output(eco) + elif output_type in ("text+image", "text,image", "image+text"): + self._process_text_image_output(eco) + elif output_type in ("latents", "latent"): + self._process_latents_output(eco) + elif output_type in ("audio", "speech"): + self._process_audio_output(eco) + elif output_type == "text": + self._process_text_output(eco) + else: + # Fallback heuristic + if eco.pooling_output is not None: + self._process_pooling_output(eco) + else: + self._process_text_output(eco) + + # ---- modality processors ---- + def _process_image_output(self, eco: EngineCoreOutput) -> None: + """Ensure image tensors are surfaced via pooling_output for vLLM.""" + if eco.pooling_output is None: + tensor = self._extract_from_multimodal_outputs( + eco, keys=("image", "images", "pixel_values", "pixels")) + if tensor is not None: + eco.pooling_output = tensor + + def _process_text_image_output(self, eco: EngineCoreOutput) -> None: + """Allow text+image outputs. Text path stays as new_token_ids; + image/latents route via pooling_output.""" + # Preserve text tokens as-is; ensure pooling_output carries image/latents + if eco.pooling_output is None: + tensor = self._extract_from_multimodal_outputs( + eco, keys=("image", "images", "pixel_values", "pixels", + "latent", "latents", "z")) + if tensor is not None: + eco.pooling_output = tensor + + def _process_latents_output(self, eco: EngineCoreOutput) -> None: + """Ensure latent tensors are surfaced via pooling_output.""" + if eco.pooling_output is None: + tensor = self._extract_from_multimodal_outputs( + eco, keys=("latent", "latents", "z", "posterior")) + if tensor is not None: + eco.pooling_output = tensor + + def _process_audio_output(self, eco: EngineCoreOutput) -> None: + """Ensure audio tensors are surfaced via pooling_output.""" + if eco.pooling_output is None: + tensor = self._extract_from_multimodal_outputs( + eco, keys=("audio", "audios", "wav", "waveform", + "audio_pcm", "pcm")) + if tensor is not None: + eco.pooling_output = tensor + + def _process_text_output(self, eco: EngineCoreOutput) -> None: + """No-op; base processor will detokenize new_token_ids → text.""" + return + + def _process_pooling_output(self, eco: EngineCoreOutput) -> None: + """Optional sanity checks for pooling tensor.""" + if eco.pooling_output is None: + return + if not isinstance(eco.pooling_output, torch.Tensor): + # Best-effort: convert to tensor if it's a list/ndarray-like + try: + eco.pooling_output = torch.as_tensor(eco.pooling_output) + except Exception: + pass + + def _extract_from_multimodal_outputs( + self, eco: EngineCoreOutput, keys: tuple[str, ...] + ) -> Optional[torch.Tensor]: + mm = getattr(eco, "multimodal_outputs", None) + if not isinstance(mm, dict): + return None + for k in keys: + v = mm.get(k) + if isinstance(v, torch.Tensor): + return v + # Try the first tensor in the dict as a fallback + for v in mm.values(): + if isinstance(v, torch.Tensor): + return v + return None \ No newline at end of file diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 8223c6f20..cdfe5d53c 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,576 +1,157 @@ -""" -Core OmniLLM and AsyncOmniLLM classes for multi-stage processing. -""" +from typing import Union, Sequence, Optional, Any +import cloudpickle +from pydantic import ValidationError -from typing import List, Dict, Any, Optional, Union, Callable +from vllm.inputs import PromptType +from vllm.sampling_params import SamplingParams from vllm.entrypoints.llm import LLM -from vllm.v1.engine.async_llm import AsyncLLM + from vllm.v1.engine.llm_engine import LLMEngine -from vllm.outputs import RequestOutput, LoRARequest +from vllm.engine.arg_utils import HfOverrides +from vllm.usage.usage_lib import UsageContext +from vllm.config import CompilationConfig, is_init_field +from vllm.utils import Counter +from vllm.logger import init_logger +import vllm.envs as envs -from ..config import OmniStageConfig -from .stage_manager import StageManager -from ..engine.output_processor import MultimodalOutputProcessor -from ..engine.diffusion_engine import DiffusersPipelineEngine +from vllm_omni.entrypoints.utils import load_stage_configs_from_model +from vllm_omni.entrypoints.omni_stage import OmniStage +from vllm_omni.engine.arg_utils import OmniEngineArgs +from vllm_omni.engine.output_processor import MultimodalOutputProcessor +from vllm_omni.engine.processor import OmniProcessor +from vllm_omni.outputs import OmniRequestOutput +logger = init_logger(__name__) -class OmniLLM(LLM): - """Extended LLM supporting multiple engines and stage-based processing.""" - - def __init__( - self, - stage_configs: List[OmniStageConfig], - log_stats: bool = False, - **kwargs - ): - # Use the first stage's model as the default model for LLM - default_model = stage_configs[0].model_path if stage_configs else "test-model" - # Track whether we should launch engines in multiprocess mode. - self.multiprocess_mode = kwargs.pop("multiprocess_mode", False) - - # Fix configuration validation issues - # Ensure max_num_batched_tokens is at least as large as max_model_len - if 'max_model_len' in kwargs and 'max_num_batched_tokens' in kwargs: - if kwargs['max_num_batched_tokens'] < kwargs['max_model_len']: - kwargs['max_num_batched_tokens'] = kwargs['max_model_len'] - elif 'max_model_len' in kwargs: - # If max_model_len is set but max_num_batched_tokens is not, set it to max_model_len - kwargs['max_num_batched_tokens'] = kwargs['max_model_len'] +class OmniLLM: + def __init__(self, model: str, stage_configs = None, log_stats: bool = False, **kwargs): + if stage_configs is None: + self.initialize_stage_configs(model) else: - # Set reasonable defaults to avoid validation errors - kwargs['max_model_len'] = 2048 - kwargs['max_num_batched_tokens'] = 2048 + self.stage_configs = stage_configs - super().__init__(model=default_model, **kwargs) - self.stage_configs = stage_configs - self.log_stats = log_stats - self.stage_manager = StageManager(stage_configs, log_stats) - self.output_processor = MultimodalOutputProcessor() - self._initialize_stage_engines() + self.stage_list = [] + self.initialize_stages(model) + + def initialize_stage_configs(self, model: str): + self.stage_configs = load_stage_configs_from_model(model) - def _initialize_stage_engines(self) -> None: - """Initialize LLMEngine instances for each stage.""" - self.stage_manager.initialize_engines() + def initialize_stages(self, model: str): + for stage_config in self.stage_configs: + stage = OmniStage(stage_config) + stage_llm = OmniStageLLM(model=model, **stage_config.engine_args) + stage.set_engine(stage_llm) + self.stage_list.append(stage) def generate( self, - stage_args_list: Optional[List[Dict[str, Any]]] = None, - use_tqdm: Union[bool, Callable[..., Any]] = True, - lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None, - priority: Optional[List[int]] = None, - *, - prompt: Optional[str] = None, - stage_overrides: Optional[Dict[int, Dict[str, Any]]] = None, - ) -> List[RequestOutput]: - """Main generation interface - orchestrates multi-stage processing.""" - if stage_args_list is None: - if prompt is None: + prompts: Union[PromptType, Sequence[PromptType]], + sampling_params_list: Optional[Union[SamplingParams, + Sequence[SamplingParams]]] = None, + ) -> list[OmniRequestOutput]: + """Generate text outputs for the given prompts.""" + final_outputs: list[OmniRequestOutput] = [] + if len(sampling_params_list) != len(self.stage_list): + raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}") + for stage_id, stage in enumerate(self.stage_list): + if stage_id > 0: + engine_inputs = stage.process_engine_inputs(self.stage_list, prompts) + else: + engine_inputs = prompts + engine_outputs = self._run_generation(stage, sampling_params_list[stage_id], engine_inputs) + stage.set_engine_outputs(engine_outputs) + if hasattr(stage, 'final_output') and stage.final_output: + final_outputs.append(OmniRequestOutput( + stage_id=stage_id, + final_output_type=stage.final_output_type, + request_output=engine_outputs)) + return final_outputs + + def _run_generation(self, stage: OmniStage, sampling_params: SamplingParams, prompts: Union[PromptType, Sequence[PromptType]]): + engine_outputs = [] + for ro in stage.engine.generate(prompts, sampling_params): + engine_outputs.append(ro) + return engine_outputs + + +class OmniStageLLM(LLM): + def __init__(self, + model: str, + compilation_config: Optional[Union[int, dict[str, Any], + CompilationConfig]] = None, + hf_overrides: Optional[HfOverrides] = None, + **kwargs): + """LLM constructor.""" + if "disable_log_stats" not in kwargs: + kwargs["disable_log_stats"] = True + + if "worker_cls" in kwargs: + worker_cls = kwargs["worker_cls"] + # if the worker_cls is not qualified string name, + # we serialize it using cloudpickle to avoid pickling issues + if isinstance(worker_cls, type): + kwargs["worker_cls"] = cloudpickle.dumps(worker_cls) + + if "kv_transfer_config" in kwargs and isinstance( + kwargs["kv_transfer_config"], dict): + from vllm.config import KVTransferConfig + raw_config_dict = kwargs["kv_transfer_config"] + try: + kwargs["kv_transfer_config"] = KVTransferConfig( + **raw_config_dict) + except ValidationError as e: + logger.error( + "Failed to convert 'kv_transfer_config' dict to " + "KVTransferConfig object. Dict: %s. Error: %s", + raw_config_dict, e) + # Consider re-raising a more specific vLLM error or ValueError + # to provide better context to the user. raise ValueError( - "prompt must be provided when stage_args_list is not supplied" - ) - stage_args_list = self._build_stage_args_from_config( - prompt, stage_overrides or {} - ) + f"Invalid 'kv_transfer_config' provided: {e}") from e - if len(stage_args_list) != len(self.stage_configs): - raise ValueError( - f"Number of stage arguments ({len(stage_args_list)}) must match " - f"number of stage configs ({len(self.stage_configs)})" - ) - - # Process through each stage sequentially - current_output = None - - for i, (stage_config, stage_args) in enumerate(zip(self.stage_configs, stage_args_list)): - stage_engine = self.stage_manager.get_engine(i) - - # Prepare input for this stage - processed_input = self._process_stage_inputs( - stage_config, stage_args or {}, current_output - ) - - # Execute stage - stage_output = self._execute_stage( - stage_engine, processed_input, lora_request, priority, stage_config - ) - - # Update for next stage - current_output = stage_output - stage_config.stage_output = stage_output - - # Process final output - final_output = self.output_processor.process_output(current_output) - return final_output - - def _build_stage_args_from_config( - self, - prompt: str, - stage_overrides: Dict[int, Dict[str, Any]], - ) -> List[Dict[str, Any]]: - """Derive per-stage argument dictionaries from configuration defaults.""" - stage_args: List[Dict[str, Any]] = [] - for stage_config in self.stage_configs: - combined: Dict[str, Any] = dict(stage_config.default_stage_args or {}) - override = stage_overrides.get(stage_config.stage_id) - if override: - combined.update(override) - if stage_config.engine_type == "AR": - combined["prompt"] = prompt - stage_args.append(combined) - return stage_args + if hf_overrides is None: + hf_overrides = {} - def _process_stage_inputs( - self, - stage_config: OmniStageConfig, - stage_args: Dict[str, Any], - previous_output: Optional[Any] - ) -> Dict[str, Any]: - """Prepare input for specific stage.""" - if stage_config.engine_type == "AR": - return self._process_ar_inputs(stage_config, stage_args, previous_output) - elif stage_config.engine_type == "DiT": - return self._process_dit_inputs(stage_config, stage_args, previous_output) + if compilation_config is not None: + if isinstance(compilation_config, int): + compilation_config_instance = CompilationConfig( + level=compilation_config) + elif isinstance(compilation_config, dict): + predicate = lambda x: is_init_field(CompilationConfig, x[0]) + compilation_config_instance = CompilationConfig( + **dict(filter(predicate, compilation_config.items()))) + else: + compilation_config_instance = compilation_config else: - raise NotImplementedError(f"Unknown engine type: {stage_config.engine_type}") + compilation_config_instance = CompilationConfig() - def _process_ar_inputs( - self, - stage_config: OmniStageConfig, - stage_args: Dict[str, Any], - previous_output: Optional[Any] - ) -> Dict[str, Any]: - """Process inputs for AR stage.""" - combined = dict(stage_config.default_stage_args or {}) - combined.update(stage_args) - combined.setdefault("prompt", "") - combined.setdefault("max_tokens", 100) - combined.setdefault("temperature", 0.7) - - # If we have previous output (e.g., from a previous AR stage), - # we might want to use it as context - if previous_output is not None: - # Extract text from previous output if available - if hasattr(previous_output, 'outputs') and previous_output.outputs: - last_output = previous_output.outputs[-1] - if hasattr(last_output, 'text'): - combined["prompt"] = last_output.text + " " + combined["prompt"] - - return combined - - def _process_dit_inputs( - self, - stage_config: OmniStageConfig, - stage_args: Dict[str, Any], - previous_output: Optional[Any] - ) -> Dict[str, Any]: - """Process inputs for DiT stage.""" - combined = dict(stage_config.default_stage_args or {}) - combined.update(stage_args) - - dit = stage_config.dit_config - if dit is not None: - combined.setdefault("height", getattr(dit, "height", 512)) - combined.setdefault("width", getattr(dit, "width", 512)) - combined.setdefault( - "num_inference_steps", getattr(dit, "num_inference_steps", 50) - ) - combined.setdefault( - "guidance_scale", getattr(dit, "guidance_scale", 7.5) - ) - else: - combined.setdefault("height", 512) - combined.setdefault("width", 512) - combined.setdefault("num_inference_steps", 50) - combined.setdefault("guidance_scale", 7.5) - - # Handle image inputs if present - if "image" in stage_args: - # For now, we'll pass the image path directly - # In a full implementation, this would involve VAE encoding - combined["image"] = stage_args["image"] - - # If we have previous output from an AR stage, we might want to use it - if previous_output is not None: - # Extract text from previous AR output - if hasattr(previous_output, 'outputs') and previous_output.outputs: - last_output = previous_output.outputs[-1] - if hasattr(last_output, 'text'): - combined["prompt"] = last_output.text - - combined.setdefault("prompt", stage_args.get("prompt", "")) - - return combined - - def _execute_stage( - self, - stage_engine: Optional[Union[LLMEngine, DiffusersPipelineEngine]], - processed_input: Dict[str, Any], - lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None, - priority: Optional[List[int]] = None, - stage_config: Optional[OmniStageConfig] = None, - ) -> Any: - """Execute a single stage.""" - # DiT via diffusers backend - if stage_config and stage_config.engine_type == "DiT": - dit = stage_config.dit_config - if dit and getattr(dit, "use_diffusers", False): - # Lazy-init executor per stage - if not hasattr(self, "_dit_engines"): - self._dit_engines = {} - exec_inst = self._dit_engines.get(stage_config.stage_id) - if exec_inst is None: - exec_inst = DiffusersPipelineEngine( - dit_config=dit, - model_path=stage_config.model_path, - log_stats=self.log_stats, - multiprocess_mode=self.multiprocess_mode, - ) - - self._dit_engines[stage_config.stage_id] = exec_inst - - return exec_inst.generate( - prompt=processed_input.get("prompt", ""), - height=processed_input.get("height", getattr(dit, "height", 512)), - width=processed_input.get("width", getattr(dit, "width", 512)), - num_inference_steps=processed_input.get( - "num_inference_steps", getattr(dit, "num_inference_steps", 30) - ), - guidance_scale=processed_input.get( - "guidance_scale", getattr(dit, "guidance_scale", 5.0) - ), - negative_prompt=processed_input.get("negative_prompt"), - seed=processed_input.get("seed"), - image=processed_input.get("image"), - ) - - # Use the parent LLM's generate method for AR text generation - prompt = processed_input.get("prompt", "") - max_tokens = processed_input.get("max_tokens", 100) - temperature = processed_input.get("temperature", 0.7) - - # Generate using the base LLM class - from vllm.sampling_params import SamplingParams - - sampling_params = SamplingParams( - max_tokens=max_tokens, - temperature=temperature, - top_p=processed_input.get("top_p", 1.0), - frequency_penalty=processed_input.get("frequency_penalty", 0.0), - presence_penalty=processed_input.get("presence_penalty", 0.0), - stop=processed_input.get("stop", None) + engine_args = OmniEngineArgs( + model=model, + hf_overrides=hf_overrides, + compilation_config=compilation_config_instance, + **kwargs, ) - - # Use the parent class's generate method - outputs = super().generate([prompt], sampling_params) - - # Return the first output (we're processing one prompt at a time) - if outputs: - return outputs[0] + + # Create the Engine (autoselects V0 vs V1) + self.llm_engine = LLMEngine.from_engine_args( + engine_args=engine_args, usage_context=UsageContext.LLM_CLASS) + self.llm_engine.output_processor = MultimodalOutputProcessor(tokenizer=self.llm_engine.tokenizer, + log_stats=self.llm_engine.log_stats) + self.llm_engine.processor = OmniProcessor(vllm_config=self.llm_engine.vllm_config, + tokenizer=self.llm_engine.tokenizer) + self.engine_class = type(self.llm_engine) + + self.request_counter = Counter() + self.default_sampling_params: Union[dict[str, Any], None] = None + + if envs.VLLM_USE_V1: + supported_tasks = self.llm_engine \ + .get_supported_tasks() # type: ignore else: - # Fallback to mock output if generation fails - from vllm.outputs import RequestOutput, CompletionOutput - - mock_output = CompletionOutput( - index=0, - text="Generation failed", - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason="error" - ) - - return RequestOutput( - request_id="fallback_request", - prompt=prompt, - prompt_token_ids=[], - prompt_logprobs=None, - outputs=[mock_output], - finished=True - ) + supported_tasks = self.llm_engine.model_config.supported_tasks + logger.info("Supported_tasks: %s", supported_tasks) -class AsyncOmniLLM(LLM): - """Extended LLM class supporting multiple engines and stage-based processing.""" - - def __init__( - self, - stage_configs: List[OmniStageConfig], - log_stats: bool = False, - **kwargs - ): - # Use the first stage's model for the base LLM - if stage_configs and stage_configs[0].model_path: - model = stage_configs[0].model_path - else: - model = "Qwen/Qwen3-0.6B" - - # Fix configuration validation issues - # Ensure max_num_batched_tokens is at least as large as max_model_len - if 'max_model_len' in kwargs and 'max_num_batched_tokens' in kwargs: - if kwargs['max_num_batched_tokens'] < kwargs['max_model_len']: - kwargs['max_num_batched_tokens'] = kwargs['max_model_len'] - elif 'max_model_len' in kwargs: - # If max_model_len is set but max_num_batched_tokens is not, set it to max_model_len - kwargs['max_num_batched_tokens'] = kwargs['max_model_len'] - else: - # Set reasonable defaults to avoid validation errors - kwargs['max_model_len'] = 2048 - kwargs['max_num_batched_tokens'] = 2048 - - super().__init__(model=model, **kwargs) - self.stage_configs = stage_configs - self.log_stats = log_stats - self.stage_manager = StageManager(stage_configs, log_stats) - self.output_processor = MultimodalOutputProcessor() - - def _initialize_async_stage_engines(self) -> None: - """Initialize AsyncLLM instances for each stage.""" - self.stage_manager.initialize_async_engines() - - async def generate_async( - self, - stage_args_list: Optional[List[Dict[str, Any]]] = None, - use_tqdm: Union[bool, Callable[..., Any]] = True, - lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None, - priority: Optional[List[int]] = None, - *, - prompt: Optional[str] = None, - stage_overrides: Optional[Dict[int, Dict[str, Any]]] = None, - ) -> List[RequestOutput]: - """Async generation interface - orchestrates multi-stage processing.""" - - if stage_args_list is None: - if prompt is None: - raise ValueError( - "prompt must be provided when stage_args_list is not supplied" - ) - stage_args_list = self._build_stage_args_from_config( - prompt, stage_overrides or {} - ) - - if len(stage_args_list) != len(self.stage_configs): - raise ValueError( - f"Number of stage arguments ({len(stage_args_list)}) must match " - f"number of stage configs ({len(self.stage_configs)})" - ) - - # Process through each stage sequentially - current_output = None - - for i, (stage_config, stage_args) in enumerate(zip(self.stage_configs, stage_args_list)): - stage_engine = self.stage_manager.get_async_engine(i) - - # Prepare input for this stage - processed_input = self._process_stage_inputs( - stage_config, stage_args or {}, current_output - ) - - # Execute stage asynchronously - stage_output = await self._execute_stage_async( - stage_engine, processed_input, lora_request, priority, stage_config - ) - - # Update for next stage - current_output = stage_output - stage_config.stage_output = stage_output - - # Process final output - final_output = self.output_processor.process_output(current_output) - return final_output - - def _build_stage_args_from_config( - self, - prompt: str, - stage_overrides: Dict[int, Dict[str, Any]], - ) -> List[Dict[str, Any]]: - stage_args: List[Dict[str, Any]] = [] - for stage_config in self.stage_configs: - combined: Dict[str, Any] = dict(stage_config.default_stage_args or {}) - override = stage_overrides.get(stage_config.stage_id) - if override: - combined.update(override) - if stage_config.engine_type == "AR": - combined["prompt"] = prompt - stage_args.append(combined) - return stage_args - - def _process_stage_inputs( - self, - stage_config: OmniStageConfig, - stage_args: Dict[str, Any], - previous_output: Optional[Any] - ) -> Dict[str, Any]: - """Prepare input for specific stage (same as OmniLLM).""" - if stage_config.engine_type == "AR": - return self._process_ar_inputs(stage_config, stage_args, previous_output) - elif stage_config.engine_type == "DiT": - return self._process_dit_inputs(stage_config, stage_args, previous_output) - else: - raise NotImplementedError(f"Unknown engine type: {stage_config.engine_type}") - - def _process_ar_inputs( - self, - stage_config: OmniStageConfig, - stage_args: Dict[str, Any], - previous_output: Optional[Any] - ) -> Dict[str, Any]: - """Process inputs for AR stage (same as OmniLLM).""" - combined = dict(stage_config.default_stage_args or {}) - combined.update(stage_args) - combined.setdefault("prompt", "") - combined.setdefault("max_tokens", 100) - combined.setdefault("temperature", 0.7) - - if previous_output is not None: - if hasattr(previous_output, 'outputs') and previous_output.outputs: - last_output = previous_output.outputs[-1] - if hasattr(last_output, 'text'): - combined["prompt"] = last_output.text + " " + combined["prompt"] - - return combined - - def _process_dit_inputs( - self, - stage_config: OmniStageConfig, - stage_args: Dict[str, Any], - previous_output: Optional[Any] - ) -> Dict[str, Any]: - """Process inputs for DiT stage (same as OmniLLM).""" - combined = dict(stage_config.default_stage_args or {}) - combined.update(stage_args) - - dit = stage_config.dit_config - if dit is not None: - combined.setdefault("height", getattr(dit, "height", 512)) - combined.setdefault("width", getattr(dit, "width", 512)) - combined.setdefault( - "num_inference_steps", getattr(dit, "num_inference_steps", 50) - ) - combined.setdefault( - "guidance_scale", getattr(dit, "guidance_scale", 7.5) - ) - else: - combined.setdefault("height", 512) - combined.setdefault("width", 512) - combined.setdefault("num_inference_steps", 50) - combined.setdefault("guidance_scale", 7.5) - - if "image" in stage_args: - combined["image"] = stage_args["image"] - - if previous_output is not None: - if hasattr(previous_output, 'outputs') and previous_output.outputs: - last_output = previous_output.outputs[-1] - if hasattr(last_output, 'text'): - combined["prompt"] = last_output.text - - combined.setdefault("prompt", stage_args.get("prompt", "")) - - return combined - - async def _execute_stage_async( - self, - stage_engine: AsyncLLM, - processed_input: Dict[str, Any], - lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None, - priority: Optional[List[int]] = None, - stage_config: Optional[OmniStageConfig] = None, - ) -> Any: - """Execute a single stage asynchronously.""" - # DiT via diffusers backend (sync call inside async for MVP) - if stage_config and stage_config.engine_type == "DiT": - dit = stage_config.dit_config - if dit and getattr(dit, "use_diffusers", False): - if not hasattr(self, "_dit_engines"): - self._dit_engines = {} - exec_inst = self._dit_engines.get(stage_config.stage_id) - if exec_inst is None: - from vllm_omni.engine.diffusion_engine import ( - DiffusersPipelineEngine, - ) - - pipeline_name = getattr(dit, "diffusers_pipeline", None) - device_cfg = getattr(dit, "device_config", None) - model_cfg = getattr(dit, "model_config", None) - if isinstance(device_cfg, dict): - device = device_cfg.get("device") - dtype = device_cfg.get("dtype") - else: - device = getattr(device_cfg, "device", None) - dtype = getattr(device_cfg, "dtype", None) - - if dtype is None: - if isinstance(model_cfg, dict): - dtype = model_cfg.get("dtype") - else: - dtype = getattr(model_cfg, "dtype", None) - - exec_inst = DiffusersPipelineEngine( - model_path=stage_config.model_path, - pipeline_name=pipeline_name, - device=device, - dtype=dtype, - ) - self._dit_engines[stage_config.stage_id] = exec_inst - - return exec_inst.generate( - prompt=processed_input.get("prompt", ""), - height=processed_input.get("height", getattr(dit, "height", 512)), - width=processed_input.get("width", getattr(dit, "width", 512)), - num_inference_steps=processed_input.get( - "num_inference_steps", getattr(dit, "num_inference_steps", 30) - ), - guidance_scale=processed_input.get( - "guidance_scale", getattr(dit, "guidance_scale", 5.0) - ), - negative_prompt=processed_input.get("negative_prompt"), - seed=processed_input.get("seed"), - image=processed_input.get("image"), - ) - - # Use the parent LLM's generate method for AR text generation - prompt = processed_input.get("prompt", "") - max_tokens = processed_input.get("max_tokens", 100) - temperature = processed_input.get("temperature", 0.7) - - # Generate using the base LLM class - from vllm.sampling_params import SamplingParams - - sampling_params = SamplingParams( - max_tokens=max_tokens, - temperature=temperature, - top_p=processed_input.get("top_p", 1.0), - frequency_penalty=processed_input.get("frequency_penalty", 0.0), - presence_penalty=processed_input.get("presence_penalty", 0.0), - stop=processed_input.get("stop", None) - ) - - # Use the parent class's generate method - outputs = super().generate([prompt], sampling_params) - - # Return the first output (we're processing one prompt at a time) - if outputs: - return outputs[0] - else: - # Fallback to mock output if generation fails - from vllm.outputs import RequestOutput, CompletionOutput - - mock_output = CompletionOutput( - index=0, - text="Generation failed", - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason="error" - ) - - return RequestOutput( - request_id="fallback_request", - prompt=prompt, - prompt_token_ids=[], - prompt_logprobs=None, - outputs=[mock_output], - finished=True - ) + self.supported_tasks = supported_tasks \ No newline at end of file diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py new file mode 100644 index 000000000..ef2249ac3 --- /dev/null +++ b/vllm_omni/entrypoints/omni_stage.py @@ -0,0 +1,80 @@ +""" +Stage manager for orchestrating multiple engines in vLLM-omni. +""" + +import importlib +from typing import List, Union +from vllm.v1.engine.llm_engine import LLMEngine +from vllm.v1.engine.async_llm import AsyncLLM +from vllm.inputs import TextPrompt +from vllm.v1.engine import EngineCoreOutput + +from vllm_omni.inputs.data import OmniTokensPrompt + + +class OmniStage: + def __init__(self, stage_config): + self.stage_config = stage_config + self.engine = None + self.async_engine = None + self.stage_id = stage_config.stage_id + self.engine_args = stage_config.engine_args + self.model_stage = stage_config.engine_args.model_stage + if hasattr(stage_config, 'engine_input_source'): + self.engine_input_source = stage_config.engine_input_source + else: + self.engine_input_source = [] + self.engine_output_type = stage_config.engine_args.engine_output_type + self.engine_outputs = None + if hasattr(stage_config, 'custom_process_input_func'): + # Import the module specified in the config (already a full module path) + module_path, func_name = stage_config.custom_process_input_func.rsplit('.', 1) + module = importlib.import_module(module_path) + self.custom_process_input_func = getattr(module, func_name) + else: + self.custom_process_input_func = None + + if hasattr(stage_config, 'final_output'): + self.final_output = stage_config.final_output + else: + self.final_output = False + + if hasattr(stage_config, 'final_output_type'): + self.final_output_type = stage_config.final_output_type + else: + self.final_output_type = None + + def set_engine(self, engine: LLMEngine) -> None: + """Initialize the engine for the stage.""" + self.engine = engine + + def set_async_engine(self, async_engine: AsyncLLM) -> None: + """Initialize the async engine for the stage.""" + self.async_engine = async_engine + + def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None: + """Set the engine output for the stage.""" + self.engine_outputs = engine_outputs + + def process_engine_inputs(self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None) -> List[Union[OmniTokensPrompt, TextPrompt]]: + """Process the engine input for the stage.""" + if self.custom_process_input_func is None: + engine_inputs = [] + if len(self.engine_input_source) == 0: + raise ValueError("engine_input_source is empty") + source_stage_id = self.engine_input_source[0] + source_outputs = stage_list[source_stage_id].engine_outputs + multi_modal_data = {source_output.request_id: + prompt.get('multi_modal_data', None) for source_output, prompt in zip(source_outputs, prompt)} + + for source_output in source_outputs: + engine_input = OmniTokensPrompt( + prompt_token_ids = source_output.outputs[0].token_ids, + multi_modal_data=multi_modal_data[source_output.request_id] if multi_modal_data else None, + ) + engine_inputs.append(engine_input) + return engine_inputs + + else: + engine_input_source = self.engine_input_source + return self.custom_process_input_func(stage_list, engine_input_source, prompt) \ No newline at end of file diff --git a/vllm_omni/entrypoints/stage_manager.py b/vllm_omni/entrypoints/stage_manager.py deleted file mode 100644 index 22a03d51c..000000000 --- a/vllm_omni/entrypoints/stage_manager.py +++ /dev/null @@ -1,87 +0,0 @@ -""" -Stage manager for orchestrating multiple engines in vLLM-omni. -""" - -from typing import List, Optional, Union -from vllm.v1.engine.llm_engine import LLMEngine -from vllm.v1.engine.async_llm import AsyncLLM -from ..engine.diffusion_engine import DiffusionEngine - -from ..config import OmniStageConfig - - -class StageManager: - """Manages multiple stage engines for multi-stage processing.""" - - def __init__(self, stage_configs: List[OmniStageConfig], log_stats: bool = False): - self.stage_configs = stage_configs - self.log_stats = log_stats - self.engine_list: List[Union[LLMEngine, DiffusionEngine]] = [] - self.async_engine_list: List[AsyncLLM] = [] - self._initialized = False - self._async_initialized = False - - def initialize_engines(self) -> None: - """Initialize LLMEngine instances for each stage.""" - if self._initialized: - return - - # For now, create placeholder engines - # In a full implementation, this would create actual engines - for stage_config in self.stage_configs: - # Placeholder - would create actual engine here - self.engine_list.append(None) - - self._initialized = True - - def initialize_async_engines(self) -> None: - """Initialize AsyncLLM instances for each stage.""" - if self._async_initialized: - return - - # For now, create placeholder engines - # In a full implementation, this would create actual engines - for stage_config in self.stage_configs: - # Placeholder - would create actual engine here - self.async_engine_list.append(None) - - self._async_initialized = True - - def get_engine(self, stage_id: int) -> LLMEngine: - """Get the engine for a specific stage.""" - if not self._initialized: - self.initialize_engines() - - if stage_id >= len(self.engine_list): - raise IndexError(f"Stage {stage_id} not found. Available stages: 0-{len(self.engine_list)-1}") - - return self.engine_list[stage_id] - - def get_async_engine(self, stage_id: int) -> AsyncLLM: - """Get the async engine for a specific stage.""" - if not self._async_initialized: - self.initialize_async_engines() - - if stage_id >= len(self.async_engine_list): - raise IndexError(f"Async stage {stage_id} not found. Available stages: 0-{len(self.async_engine_list)-1}") - - return self.async_engine_list[stage_id] - - def get_stage_config(self, stage_id: int) -> OmniStageConfig: - """Get the configuration for a specific stage.""" - if stage_id >= len(self.stage_configs): - raise IndexError(f"Stage config {stage_id} not found. Available stages: 0-{len(self.stage_configs)-1}") - - return self.stage_configs[stage_id] - - def get_num_stages(self) -> int: - """Get the number of stages.""" - return len(self.stage_configs) - - def cleanup(self) -> None: - """Clean up resources.""" - # Clean up engines if needed - self.engine_list.clear() - self.async_engine_list.clear() - self._initialized = False - self._async_initialized = False diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py new file mode 100644 index 000000000..42a98f164 --- /dev/null +++ b/vllm_omni/entrypoints/utils.py @@ -0,0 +1,25 @@ +import os +from pathlib import Path +from omegaconf import OmegaConf +from vllm.transformers_utils.config import get_config + +# Get the project root directory (2 levels up from this file) +PROJECT_ROOT = Path(__file__).parent.parent.parent + + +def load_stage_configs_from_model(model: str): + """Load stage configs from model.""" + hf_config = get_config(model, trust_remote_code=True) + model_type = hf_config.model_type + stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml" + stage_config_path = PROJECT_ROOT / stage_config_file + if not os.path.exists(stage_config_path): + raise FileNotFoundError(f"Stage config file {stage_config_path} not found") + stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path)) + return stage_configs + + +def load_stage_configs_from_yaml(config_path: str): + """Load stage configs from yaml file.""" + config_data = OmegaConf.load(config_path) + return config_data.stage_args \ No newline at end of file diff --git a/vllm_omni/model_executor/layers/__init__.py b/vllm_omni/model_executor/layers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/layers/mrope.py b/vllm_omni/model_executor/layers/mrope.py new file mode 100644 index 000000000..eb266d641 --- /dev/null +++ b/vllm_omni/model_executor/layers/mrope.py @@ -0,0 +1,699 @@ +import itertools +from typing import Optional, Union + +import numpy as np +import torch +from transformers import PretrainedConfig + +from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding +from vllm.model_executor.layers.rotary_embedding.common import apply_rotary_emb_dispatch + +def _apply_rotary_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + is_neox_style: bool, +) -> torch.Tensor: + """ + Args: + x: [num_tokens, num_heads, head_size] + cos: [num_tokens, head_size // 2] + sin: [num_tokens, head_size // 2] + is_neox_style: Whether to use the Neox-style or GPT-J-style rotary + positional embeddings. + """ + cos = cos.unsqueeze(-2).to(x.dtype) + sin = sin.unsqueeze(-2).to(x.dtype) + if is_neox_style: + x1, x2 = torch.chunk(x, 2, dim=-1) + else: + x1 = x[..., ::2] + x2 = x[..., 1::2] + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + if is_neox_style: + return torch.cat((o1, o2), dim=-1) + else: + return torch.stack((o1, o2), dim=-1).flatten(-2) + +class MRotaryEmbedding(RotaryEmbedding): + """Rotary Embedding with Multimodal Sections.""" + + def __init__( + self, + head_size: int, + rotary_dim: int, + max_position_embeddings: int, + base: float, + is_neox_style: bool, + dtype: torch.dtype, + mrope_section: Optional[list[int]] = None, + ) -> None: + # In Qwen2.5-VL, the maximum index value is related to the duration of + # the input video. We enlarge max_position_embeddings to 4 times to get + # a larger the cos and sin cache. + self.cache_max_position_num = max_position_embeddings * 4 + super().__init__(head_size, rotary_dim, self.cache_max_position_num, + base, is_neox_style, dtype) + + self.mrope_section = mrope_section + if self.mrope_section: + print("Warning: mrope_section check is disabled in Qwen2.5-Omni, this may cause errors, and should be restored in the future.") + # assert sum(self.mrope_section) == rotary_dim // 2 + + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """PyTorch-native implementation equivalent to forward(). + + Args: + positions: + [num_tokens,] (text only) or + [3, num_tokens] (T/H/W positions with multimodal inputs) + query: [num_tokens, num_heads * head_size] + key: [num_tokens, num_kv_heads * head_size] + """ + assert positions.ndim == 1 or positions.ndim == 2 + assert key is not None + + num_tokens = positions.shape[-1] + cos_sin = self.cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + if positions.ndim == 2: + assert self.mrope_section + + cos = torch.cat([ + m[i] + for i, m in enumerate(cos.split(self.mrope_section, dim=-1)) + ], + dim=-1) + sin = torch.cat([ + m[i] + for i, m in enumerate(sin.split(self.mrope_section, dim=-1)) + ], + dim=-1) + + query_shape = query.shape + query = query.view(num_tokens, -1, self.head_size) + query_rot = query[..., :self.rotary_dim] + query_pass = query[..., self.rotary_dim:] + # query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin, + # self.is_neox_style) + query_rot = _apply_rotary_emb(query_rot, cos, sin, + self.is_neox_style) + query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape) + + key_shape = key.shape + key = key.view(num_tokens, -1, self.head_size) + key_rot = key[..., :self.rotary_dim] + key_pass = key[..., self.rotary_dim:] + # key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin, + # self.is_neox_style) + key_rot = _apply_rotary_emb(key_rot, cos, sin, + self.is_neox_style) + key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) + return query, key + + @classmethod + def get_input_positions( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]], + video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]], + second_per_grid_ts: Optional[list[float]], + context_len: int = 0, + seq_len: Optional[int] = None, + audio_feature_lengths: Optional[torch.Tensor] = None, + use_audio_in_video: bool = False, + ) -> tuple[list[list[int]], int]: + """Get mrope input positions and delta value.""" + + image_grid_thw = [] if image_grid_thw is None else image_grid_thw + video_grid_thw = [] if video_grid_thw is None else video_grid_thw + second_per_grid_ts = [] if second_per_grid_ts is None else \ + second_per_grid_ts + + llm_positions, mrope_position_delta = \ + cls.get_input_positions_tensor( + input_tokens=input_tokens, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + context_len=context_len, + seq_len=seq_len, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) + + return llm_positions.tolist(), mrope_position_delta + + @classmethod + def get_input_positions_tensor( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: Union[list[list[int]], torch.Tensor], + video_grid_thw: Union[list[list[int]], torch.Tensor], + second_per_grid_ts: list[float], + context_len: int = 0, + seq_len: Optional[int] = None, + audio_feature_lengths: Optional[torch.Tensor] = None, + use_audio_in_video: bool = False, + ) -> tuple[torch.Tensor, int]: + from vllm.transformers_utils.config import thinker_uses_mrope + if thinker_uses_mrope(hf_config): + return cls._omni_get_input_positions_tensor( + input_tokens=input_tokens, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + context_len=context_len, + seq_len=seq_len, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) + elif hf_config.model_type in ["glm4v", "glm4v_moe"]: + return cls._glm4v_get_input_positions_tensor( + input_tokens=input_tokens, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + context_len=context_len, + seq_len=seq_len, + ) + else: + return cls._vl_get_input_positions_tensor( + input_tokens=input_tokens, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + context_len=context_len, + seq_len=seq_len, + ) + + @classmethod + def _glm4v_get_input_positions_tensor( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: Union[list[list[int]], torch.Tensor], + video_grid_thw: Union[list[list[int]], torch.Tensor], + context_len: int = 0, + seq_len: Optional[int] = None, + ) -> tuple[torch.Tensor, int]: + """Get mrope input positions and delta value for GLM4V.""" + + image_token_id = hf_config.image_token_id + video_start_token_id = hf_config.video_start_token_id + video_end_token_id = hf_config.video_end_token_id + spatial_merge_size = hf_config.vision_config.spatial_merge_size + llm_pos_ids_list: list = [] + + if not (image_grid_thw is None and video_grid_thw is None): + if isinstance(image_grid_thw, torch.Tensor): + image_grid_thw = image_grid_thw.tolist() + + input_token_type: list[str] = [] + video_check_flg = False + for token in input_tokens: + if token == video_start_token_id: + video_check_flg = True + elif token == video_end_token_id: + video_check_flg = False + + if (token == image_token_id) and (video_check_flg is False): + input_token_type.append("image") + elif (token == image_token_id) and (video_check_flg is True): + input_token_type.append("video") + else: + input_token_type.append("text") + + input_type_group: list[tuple[str, int, int]] = [] + for key, group_iter in itertools.groupby( + enumerate(input_token_type), lambda x: x[1]): + group_list = list(group_iter) + start_index = group_list[0][0] + end_index = group_list[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + video_frame_num = 1 + mm_data_idx = 0 + for modality_type, start_idx, end_idx in input_type_group: + st_idx = llm_pos_ids_list[-1].max() + 1 if len( + llm_pos_ids_list) > 0 else 0 + if modality_type == "image": + t, h, w = ( + image_grid_thw[mm_data_idx][0], + image_grid_thw[mm_data_idx][1], + image_grid_thw[mm_data_idx][2], + ) + llm_grid_t, llm_grid_h, llm_grid_w = \ + t, h // spatial_merge_size, w // spatial_merge_size + + t_index = torch.arange(llm_grid_t).view(-1, 1).expand( + -1, llm_grid_h * llm_grid_w).flatten() + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand( + llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand( + llm_grid_t, llm_grid_h, -1).flatten() + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + st_idx) + mm_data_idx += 1 + + elif modality_type == "video": + t, h, w = ( + video_frame_num, + image_grid_thw[mm_data_idx][1], + image_grid_thw[mm_data_idx][2], + ) + llm_grid_t, llm_grid_h, llm_grid_w = \ + t, h // spatial_merge_size, w // spatial_merge_size + + for t_idx in range(llm_grid_t): + t_index = torch.tensor(t_idx).view(-1, 1).expand( + -1, llm_grid_h * llm_grid_w).flatten() + h_index = torch.arange(llm_grid_h).view( + 1, -1, 1).expand(1, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view( + 1, 1, -1).expand(1, llm_grid_h, -1).flatten() + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + st_idx) + + mm_data_idx += 1 + video_frame_num += 1 + + else: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + + st_idx) + video_frame_num = 1 + + else: + text_len = len(input_tokens) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1)) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + llm_positions = llm_positions[:, context_len:seq_len] + mrope_position_delta = (llm_positions.max() + 1 - + len(input_tokens)).item() + return llm_positions, mrope_position_delta + + @classmethod + def _vl_get_input_positions_tensor( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: Union[list[list[int]], torch.Tensor], + video_grid_thw: Union[list[list[int]], torch.Tensor], + second_per_grid_ts: list[float], + context_len: int = 0, + seq_len: Optional[int] = None, + ) -> tuple[torch.Tensor, int]: + """Get mrope input positions and delta value.""" + + image_token_id = hf_config.image_token_id + video_token_id = hf_config.video_token_id + vision_start_token_id = hf_config.vision_start_token_id + spatial_merge_size = hf_config.vision_config.spatial_merge_size + tokens_per_second = getattr(hf_config.vision_config, + "tokens_per_second", 1.0) + + input_tokens_tensor = torch.tensor(input_tokens) + vision_start_indices = torch.argwhere( + input_tokens_tensor == vision_start_token_id).squeeze(1) + vision_tokens = input_tokens_tensor[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + llm_pos_ids_list: list = [] + + st = 0 + remain_images, remain_videos = image_nums, video_nums + + image_index, video_index = 0, 0 + for _ in range(image_nums + video_nums): + video_second_per_grid_t = 0.0 + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + if ed_image < ed_video: + t, h, w = ( + image_grid_thw[image_index][0], + image_grid_thw[image_index][1], + image_grid_thw[image_index][2], + ) + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = ( + video_grid_thw[video_index][0], + video_grid_thw[video_index][1], + video_grid_thw[video_index][2], + ) + video_second_per_grid_t = 1.0 + if second_per_grid_ts: + video_second_per_grid_t = second_per_grid_ts[video_index] + video_index += 1 + remain_videos -= 1 + ed = ed_video + + llm_grid_t, llm_grid_h, llm_grid_w = \ + t, h // spatial_merge_size, w // spatial_merge_size + text_len = ed - st + + st_idx = llm_pos_ids_list[-1].max() + 1 if len( + llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + t_index = (torch.arange(llm_grid_t).view(-1, 1).expand( + -1, llm_grid_h * llm_grid_w) * video_second_per_grid_t * + tokens_per_second).long().flatten() + + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand( + llm_grid_t, -1, llm_grid_w).flatten() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand( + llm_grid_t, llm_grid_h, -1).flatten() + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len( + llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + mrope_position_delta = (llm_positions.max() + 1 - + len(input_tokens)).item() + llm_positions = llm_positions[:, context_len:seq_len] + + return llm_positions, mrope_position_delta + + @classmethod + def _omni_get_input_positions_tensor( + cls, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: Union[list[list[int]], torch.Tensor], + video_grid_thw: Union[list[list[int]], torch.Tensor], + second_per_grid_ts: Optional[list[float]] = None, + context_len: int = 0, + seq_len: Optional[int] = None, + audio_feature_lengths: Optional[torch.Tensor] = None, + use_audio_in_video: bool = False, + ) -> tuple[torch.Tensor, int]: + """Get mrope input positions and delta value (Qwen2.5-Omni version). + + Differences from MRotaryEmbedding: + 1. Add audio support (and related `audio_feature_lengths`). + 2. Add `use_audio_in_video` option to read audio from video inputs. + In this case, audio and vision position ids will be split into + chunks and interleaved. + + Example: + + (V_i are vision position ids, A_i are audio position ids) + + |V_1 ... V_n|A_1 ... A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|... + |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |... + """ + + # TODO(fyabc): refactor and share more code with + # _vl_get_input_positions_tensor. + + thinker_config = hf_config.thinker_config + audio_token_id = thinker_config.audio_token_index + image_token_id = thinker_config.image_token_index + video_token_id = thinker_config.video_token_index + audio_start_token_id = thinker_config.audio_start_token_id + audio_end_token_id = thinker_config.audio_end_token_id + vision_start_token_id = thinker_config.vision_start_token_id + vision_end_token_id = thinker_config.vision_end_token_id + seconds_per_chunk = thinker_config.seconds_per_chunk + spatial_merge_size = thinker_config.vision_config.spatial_merge_size + tokens_per_second = getattr(thinker_config.vision_config, + "tokens_per_second", 25) + + if isinstance(image_grid_thw, list): + image_grid_thw = torch.tensor(image_grid_thw) + if isinstance(video_grid_thw, list): + video_grid_thw = torch.tensor(video_grid_thw) + + src_item = input_tokens + audio_seqlens = audio_feature_lengths + if not second_per_grid_ts: + second_per_grid_ts = [1] * video_grid_thw.shape[0] + audio_idx = 0 + video_idx = 0 + image_idx = 0 + new_src_item: list[int] = [] + llm_pos_ids_list: list[torch.Tensor] = [] + + idx = 0 + while idx < len(src_item): + new_src_item_len = len(new_src_item) + start_idx = llm_pos_ids_list[-1].max() + 1 if len( + llm_pos_ids_list) > 0 else 0 + if src_item[idx] not in [ + audio_token_id, video_token_id, image_token_id + ]: + if use_audio_in_video and idx > 0: + if src_item[idx] == vision_end_token_id and \ + src_item[idx - 1] == audio_end_token_id: + # processing the <|audio_eos|> before <|vision_eos|> + start_idx -= 1 + elif src_item[idx] == audio_start_token_id and \ + src_item[idx - 1] == vision_start_token_id: + # processing the <|audio_bos|> after <|vision_eos|> + start_idx -= 1 + new_src_item.append(src_item[idx]) + llm_pos_ids = torch.tensor([start_idx], + dtype=torch.long).expand(3, -1) + llm_pos_ids_list.append(llm_pos_ids) + elif src_item[idx] == audio_token_id: + assert audio_seqlens is not None + audio_seqlen = audio_seqlens[audio_idx] + place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + new_src_item.extend([audio_token_id] * place_num) + llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx + llm_pos_ids_list.append(llm_pos_ids) + audio_idx += 1 + elif src_item[idx] == image_token_id: + grid_t = image_grid_thw[image_idx][0] + grid_hs = image_grid_thw[:, 1] + grid_ws = image_grid_thw[:, 2] + t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long() + llm_pos_ids = cls._get_llm_pos_ids_for_vision( + start_idx, image_idx, spatial_merge_size, t_index, grid_hs, + grid_ws) + llm_pos_ids_list.append(llm_pos_ids) + vision_seqlen = image_grid_thw[image_idx].prod() // ( + spatial_merge_size**2) + new_src_item.extend([image_token_id] * vision_seqlen) + image_idx += 1 + elif src_item[idx] == video_token_id and not use_audio_in_video: + grid_t = video_grid_thw[video_idx][0] + grid_hs = video_grid_thw[:, 1] + grid_ws = video_grid_thw[:, 2] + t_index = (torch.arange(grid_t) * + second_per_grid_ts[video_idx] * + tokens_per_second).long() + llm_pos_ids = cls._get_llm_pos_ids_for_vision( + start_idx, video_idx, spatial_merge_size, t_index, grid_hs, + grid_ws) + llm_pos_ids_list.append(llm_pos_ids) + vision_seqlen = video_grid_thw[video_idx].prod() // ( + spatial_merge_size**2) + new_src_item.extend([video_token_id] * vision_seqlen) + video_idx += 1 + else: + # read audio from video + assert audio_seqlens is not None + audio_seqlen = audio_seqlens[audio_idx] + vision_seqlen = video_grid_thw[video_idx].prod() // ( + spatial_merge_size**2) + grid_t = video_grid_thw[video_idx][0] + grid_h = video_grid_thw[video_idx][1] + grid_w = video_grid_thw[video_idx][2] + grid_hs = video_grid_thw[:, 1] + grid_ws = video_grid_thw[:, 2] + t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk) + t_index = (torch.arange(grid_t) * + second_per_grid_ts[video_idx] * + tokens_per_second).long() + t_index_split_chunk = cls._split_list_into_ranges( + t_index, t_ntoken_per_chunk) + place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2 + pure_audio_len = place_num - 2 + added_audio_len = 0 + audio_llm_pos_ids_list: list[torch.Tensor] = [] + for t_chunk in t_index_split_chunk: + vision_ntoken_per_chunk = len( + t_chunk) * grid_h * grid_w // (spatial_merge_size**2) + new_src_item.extend([video_token_id] * + vision_ntoken_per_chunk) + vision_llm_pos_ids_list = cls._get_llm_pos_ids_for_vision( + start_idx, video_idx, spatial_merge_size, t_chunk, + grid_hs, grid_ws).split(1, dim=1) + llm_pos_ids_list.extend(vision_llm_pos_ids_list) + new_src_item.extend( + min(t_ntoken_per_chunk, pure_audio_len - + added_audio_len) * [audio_token_id]) + audio_start_idx = start_idx if len( + audio_llm_pos_ids_list + ) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1 + if min(t_ntoken_per_chunk, + pure_audio_len - added_audio_len) > 0: + audio_llm_pos_ids_list = (torch.arange( + min(t_ntoken_per_chunk, pure_audio_len - + added_audio_len)).expand(3, -1) + + audio_start_idx).split(1, + dim=1) + else: + audio_llm_pos_ids_list = [] + added_audio_len += min(t_ntoken_per_chunk, + pure_audio_len - added_audio_len) + llm_pos_ids_list.extend(audio_llm_pos_ids_list) + if added_audio_len < pure_audio_len: + new_src_item.extend( + (pure_audio_len - added_audio_len) * [audio_token_id]) + audio_llm_pos_ids_list = ( + torch.arange(pure_audio_len - added_audio_len).expand( + 3, -1) + llm_pos_ids_list[-1].max() + 1).split( + 1, dim=1) + llm_pos_ids_list.extend(audio_llm_pos_ids_list) + audio_idx += 1 + video_idx += 1 + # move to the next token + idx += len(new_src_item) - new_src_item_len + + llm_positions = torch.cat(llm_pos_ids_list, dim=1) + mrope_position_delta = torch.cat(llm_pos_ids_list, + dim=1).max() + 1 - len(src_item) + llm_positions = llm_positions[:, context_len:seq_len] + + return llm_positions, mrope_position_delta + + @staticmethod + def _get_llm_pos_ids_for_vision( + start_idx: int, + vision_idx: int, + spatial_merge_size: int, + t_index: list[int], + grid_hs: torch.Tensor, + grid_ws: torch.Tensor, + ) -> torch.Tensor: + llm_pos_ids_list = [] + llm_grid_h = grid_hs[vision_idx] // spatial_merge_size + llm_grid_w = grid_ws[vision_idx] // spatial_merge_size + h_index = (torch.arange(llm_grid_h).view(1, -1, 1).expand( + len(t_index), -1, llm_grid_w).flatten()) + w_index = (torch.arange(llm_grid_w).view(1, 1, -1).expand( + len(t_index), llm_grid_h, -1).flatten()) + t_index_tensor = torch.Tensor(t_index).to(llm_grid_h.device).view( + -1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten() + _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index]) + llm_pos_ids_list.append(_llm_pos_ids + start_idx) + llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1) + return llm_pos_ids + + @staticmethod + def _split_list_into_ranges(lst: torch.Tensor, + interval: int) -> list[list[int]]: + ranges: list[list[int]] = [[] + for _ in range((max(lst) // interval) + 1)] + for num in lst: + index = num // interval + ranges[index].append(num) + return ranges + + @staticmethod + def get_next_input_positions( + mrope_position_delta: int, + context_len: int, + seq_len: int, + ) -> list[list[int]]: + return [ + list( + range(context_len + mrope_position_delta, + seq_len + mrope_position_delta)) for _ in range(3) + ] + + @staticmethod + def get_next_input_positions_tensor(out: np.ndarray, out_offset: int, + mrope_position_delta: int, + context_len: int, num_new_tokens: int): + + values = np.arange(mrope_position_delta + context_len, + mrope_position_delta + context_len + num_new_tokens, + dtype=out.dtype) + out[:, out_offset:out_offset + num_new_tokens] = values + + @classmethod + def omni_get_updates_use_audio_in_video( + cls, + thinker_config: PretrainedConfig, + audio_len: int, + video_grid_thw: Union[list[int], torch.Tensor], + video_second_per_grid_t: float, + ) -> list[int]: + """Get video prompt updates when `use_audio_in_video` is True. + + In this case, audio and vision update ids will be split into + chunks and interleaved (details in `_omni_get_input_positions_tensor`). + + <|video_bos|><|VIDEO|><|video_eos|> => + <|video_bos|><|audio_bos|>(... chunks ...)<|audio_eos|><|video_eos|> + """ + + audio_token_id = thinker_config.audio_token_index + video_token_id = thinker_config.video_token_index + audio_start_token_id = thinker_config.audio_start_token_id + audio_end_token_id = thinker_config.audio_end_token_id + seconds_per_chunk = thinker_config.seconds_per_chunk + spatial_merge_size = thinker_config.vision_config.spatial_merge_size + tokens_per_second = getattr(thinker_config.vision_config, + "tokens_per_second", 25) + + grid_t = video_grid_thw[0] + grid_h = video_grid_thw[1] + grid_w = video_grid_thw[2] + t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk) + t_index = (torch.arange(grid_t) * video_second_per_grid_t * + tokens_per_second).long() + t_index_split_chunk = cls._split_list_into_ranges( + t_index, t_ntoken_per_chunk) + + updates = [audio_start_token_id] + added_audio_len = 0 + for t_chunk in t_index_split_chunk: + vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // ( + spatial_merge_size**2) + updates.extend([video_token_id] * vision_ntoken_per_chunk) + + audio_chunk_size = min(t_ntoken_per_chunk, + audio_len - added_audio_len) + updates.extend(audio_chunk_size * [audio_token_id]) + added_audio_len += audio_chunk_size + if added_audio_len < audio_len: + updates.extend((audio_len - added_audio_len) * [audio_token_id]) + updates.extend([audio_end_token_id]) + + return updates diff --git a/vllm_omni/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py new file mode 100644 index 000000000..2a2ba700e --- /dev/null +++ b/vllm_omni/model_executor/models/qwen2_5_omni.py @@ -0,0 +1,710 @@ +from functools import cached_property +from typing import Iterable, List, Optional, Set, Tuple, Union, NamedTuple, Dict + +import os +import glob +import numpy as np +import torch +import torch.nn as nn +from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( + Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig, Qwen2_5OmniTalkerConfig) + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler +from vllm.model_executor.sampling_metadata import SamplingMetadata +from vllm.model_executor.models.qwen2_5_omni_thinker import ( + Qwen2_5OmniConditionalGenerationMixin, + Qwen2_5OmniThinkerMultiModalProcessor, + Qwen2_5OmniThinkerProcessingInfo, + Qwen2_5OmniThinkerDummyInputsBuilder) +# from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors + +from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP +from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper, + init_vllm_registered_model, + maybe_prefix) +from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf +from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights + +TALKER_CODEC_EOS_TOKEN_ID = 8294 +TALKER_CODEC_BOS_TOKEN_ID = 8293 + +class OmniOutput(NamedTuple): + """Output from the merged Omni model containing both text and audio.""" + text_hidden_states: torch.Tensor + multimodal_outputs: dict = {} + intermediate_tensors: Optional[IntermediateTensors] = None + +logger = init_logger(__name__) + +@MULTIMODAL_REGISTRY.register_processor( + Qwen2_5OmniThinkerMultiModalProcessor, + info=Qwen2_5OmniThinkerProcessingInfo, + dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder, +) +class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal, + SupportsPP, + Qwen2_5OmniConditionalGenerationMixin): + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.have_multimodal_outputs = True + config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + # keep vllm_config for later submodule init + self.vllm_config = vllm_config + + # Initialize thinker components + thinker_config: Qwen2_5OmniThinkerConfig = config.thinker_config + self.thinker_config = thinker_config + self.multimodal_config = multimodal_config + + # Initialize talker components + talker_config: Qwen2_5OmniTalkerConfig = config.talker_config + self.talker_config = talker_config + + + + self.model_stage = vllm_config.model_config.model_stage + if self.model_stage=="thinker": + # Initialize thinker model (multimodal processing) + self.thinker = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "thinker"), + hf_config=thinker_config, + # Use registry architecture key + architectures=["Qwen2_5OmniThinkerModel"], + ) + self.model = self.thinker + self.talker = None + self.token2wav = None + + elif self.model_stage=="talker": + self.thinker = None + # Initialize talker model wrapper (handles projection + LM) + self.talker = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "talker"), + hf_config=talker_config, + # Use registry architecture key + architectures=["Qwen2_5OmniTalkerModel"], + ) + self.talker.init_multi_modal(thinker_config) + self.model=self.talker + self.token2wav = None + + elif self.model_stage=="code2wav": + self.thinker = None + self.talker = None + # Initialize token2wav (code->mel->wav) like thinker/talker + self.token2wav_config = getattr(config, 'token2wav_config', None) + self.token2wav = None + if self.token2wav_config is not None: + self.token2wav = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "token2wav"), + hf_config=self.token2wav_config, + architectures=["Qwen2_5OmniToken2WavModel"], + ) + # voice resources (loaded on demand) + self._token2wav_conds: Dict[str, torch.Tensor] = {} + self._token2wav_ref_mels: Dict[str, torch.Tensor] = {} + self.model = self.token2wav + else: + raise ValueError("Invalid model stage") + + # Set up intermediate tensors + self.make_empty_intermediate_tensors = ( + self.thinker.make_empty_intermediate_tensors) if self.model_stage=="thinker" else lambda: None + + # -------------------- Device utilities -------------------- + @staticmethod + def _module_device(module: nn.Module) -> torch.device: + try: + return next(module.parameters()).device + except StopIteration: + # No parameters; fall back to buffers or cpu + for _, buf in module.named_buffers(recurse=True): + return buf.device + return torch.device("cpu") + + def move_submodules_to_devices( + self, + *, + thinker_device: Optional[Union[str, torch.device]] = None, + talker_device: Optional[Union[str, torch.device]] = None, + token2wav_device: Optional[Union[str, torch.device]] = None, + ) -> None: + """Optionally move thinker/talker/token2wav to different devices. + + Example: + model.move_submodules_to_devices( + thinker_device='cuda:0', + talker_device='cuda:1', + token2wav_device='cpu', + ) + """ + if thinker_device is not None and self.thinker is not None: + self.thinker.to(thinker_device) + if talker_device is not None and self.talker is not None: + self.talker.to(talker_device) + if token2wav_device is not None and self.token2wav is not None: + self.token2wav.to(token2wav_device) + + @cached_property + def sampler(self): + if hasattr(self.model, "sampler"): + return self.model.sampler + return get_sampler() + + def get_input_embeddings( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + ) -> torch.Tensor: + if self.model_stage == "code2wav": + return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size()) + return self.model.get_input_embeddings( + input_ids, multimodal_embeddings) + + + def get_multimodal_embeddings(self, **kwargs): + # Delegate to thinker model for multimodal processing + return self.model.get_multimodal_embeddings(**kwargs) + + def last_index_of(self, list, value): + return len(list) - 1 - list[::-1].index(value) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + generate_audio: bool = True, + voice_type: str = "Chelsie", + codec: Optional[torch.Tensor] = None, + sampling_metadata: Optional[SamplingMetadata] = None, + logits_index: Optional[int] = None, + sampler = None, + additional_information: Optional[dict[str, object]] = None, + **kwargs: object, + ) -> Union[torch.Tensor, IntermediateTensors, OmniOutput]: + """ + Workflow: + 1) Thinker: multimodal understanding → text hidden states. + 2) If audio requested and codec not provided, use talker to derive codec. + 3) If audio requested (or codec provided), use token2wav to synthesize waveform. + 4) Return text hidden states (and audio when applicable). + """ + if self.model_stage=="thinker": + # Normalize to batched inputs if caller provides 1D/2D unbatched tensors + added_batch_dim = False + if input_ids is not None and input_ids.ndim == 1: + input_ids = input_ids.unsqueeze(0) + added_batch_dim = True + if positions is not None and positions.ndim == 1: + positions = positions.unsqueeze(0) + added_batch_dim = True + if inputs_embeds is not None and inputs_embeds.ndim == 2: + inputs_embeds = inputs_embeds.unsqueeze(0) + added_batch_dim = True + thinker_dev = self._module_device(self.thinker) + + #if input_ids is None, set it to a zero tensor, in the length of the same as the embedding seq length + if input_ids is None: + input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev).unsqueeze(0) #(1, 0) + added_batch_dim = True + + # 1) Thinker (ensure inputs on thinker's device) + if input_ids is not None and input_ids.device != thinker_dev: + input_ids = input_ids.to(thinker_dev) + if positions is not None and positions.device != thinker_dev: + positions = positions.to(thinker_dev) + if inputs_embeds is not None and inputs_embeds.device != thinker_dev: + inputs_embeds = inputs_embeds.to(thinker_dev) + # Run thinker + thinker_output = self.thinker( + input_ids=input_ids, + positions=positions[0], + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + if isinstance(thinker_output, tuple): + embeds, text_hidden_states = thinker_output + else: + text_hidden_states = thinker_output + + # Text-only path + return OmniOutput( + text_hidden_states=text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states, + multimodal_outputs=None + ) + + # 2) Talker (if codec not provided) + if self.model_stage=="talker": + + if input_ids is None and additional_information is None: + input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device) + additional_information = {} + self.thinker_reply_part = torch.zeros_like(inputs_embeds) + is_profile = True + else: + is_profile = False + + if input_ids is not None and additional_information is not None and not is_profile: + # read from additional_information dict + thinker_result = None + if additional_information is not None and isinstance(additional_information, dict): + thinker_result = additional_information.get("thinker_result") + prompt_embeds = additional_information.get("prompt_embeds") + prompt_token_ids = additional_information.get("prompt_token_ids") + thinker_output_token_ids = additional_information.get("thinker_output_token_ids") + else: + thinker_result = torch.zeros_like(inputs_embeds) + prompt_embeds = torch.zeros_like(inputs_embeds) + prompt_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device) + thinker_output_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device) + + if thinker_result is None: + thinker_result = torch.zeros_like(inputs_embeds) + self.thinker_reply_part = thinker_result.squeeze(0) + if self.thinker_reply_part.shape[1] > 1: + self.thinker_reply_part = self.thinker_reply_part[1:, :] + input_ids, inputs_embeds = self._thinker_to_talker_prefill( + voice_type=voice_type, + output_prompt_embeds=thinker_result, + output_token_ids = thinker_output_token_ids, + thinker_prompt_embeds=prompt_embeds, + prompt_token_ids = prompt_token_ids, + ) + elif not is_profile: + input_ids, inputs_embeds = \ + self._thinker_to_talker_decode_one_step( + output_prompt_embeds=self.thinker_reply_part[:1] if self.thinker_reply_part.shape[0]>=1 \ + else torch.zeros(1, self.thinker_reply_part.shape[1]).to(self._module_device(self.model)).to(torch.bfloat16)+(-1.25*2**(-123)), + output_token_ids=input_ids, + ) + + if self.thinker_reply_part.shape[0] >=1: + self.thinker_reply_part = self.thinker_reply_part[1:, :] + + with torch.inference_mode(): + talker_hidden = self.talker( + input_ids=input_ids, + positions=positions[0], + inputs_embeds=inputs_embeds, + + ) + + return OmniOutput( + text_hidden_states=talker_hidden, + multimodal_outputs=None + ) + + if self.model_stage=="code2wav": + code = input_ids if input_ids is not None else torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device) + audio_tensor = self.generate_audio(code[:-1] if code[-1]==TALKER_CODEC_EOS_TOKEN_ID else code, voice_type) + return OmniOutput( + text_hidden_states = None, + multimodal_outputs = { + "audio": audio_tensor + } + ) + + return OmniOutput( + text_hidden_states=torch.cat( + [ + torch.zeros([inputs_embeds.shape[0],self.talker.config.hidden_size], dtype=torch.bfloat16).to(self._module_device(self.model)), + self.talker.thinker_to_talker_proj(self.talker.get_input_embeddings(torch.tensor([TALKER_CODEC_BOS_TOKEN_ID,TALKER_CODEC_EOS_TOKEN_ID]).to(torch.bfloat16).to(self._module_device(self.model))))[0] + ], + dim=0), + multimodal_outputs=None + ) + + def generate_audio(self, code, voice_type): + token2wav_dev = self._module_device(self.token2wav) + if isinstance(code, torch.Tensor): + code_tensor = code.to(dtype=torch.long, device=token2wav_dev) + else: + code_tensor = torch.as_tensor(code, dtype=torch.long, device=token2wav_dev) + if code_tensor.ndim == 2 and code_tensor.shape[0] == 1: + code_tensor = code_tensor.squeeze(0) + + audio_tensor = self._codec_to_audio(code_tensor, voice_type) + + return audio_tensor + + def _load_talker_embedding( + self, + ) -> torch.nn.Embedding: + return self.talker.language_model.model.embed_tokens + + def _init_special_tokens_embeddings( + self, + ): + # talker embeddings + self.talker_embedding = self._load_talker_embedding() + + # embed_text_bos_token + self.tts_text_spk_token_ids = { + # M02: Male voice with standard Mandarin and a slight northern accent + 'm02': 151870, + 'Ethan': 151870, + + # F030: Your anime-styled virtual girlfriend + 'f030': 151872, + 'Chelsie': 151872, + } + self.default_tts_text_spk_type = list( + self.tts_text_spk_token_ids.keys())[0] + self.tts_text_spk_token_ids['prefix_caching'] = 151870 + + talker_hf_config = self.talker_config + if hasattr(talker_hf_config, 'talker_config'): + talker_hf_config = talker_hf_config.talker_config + + self.embed_text_bos_token = self.thinker_embedding( + torch.tensor( + [talker_hf_config.tts_text_start_token_id], + dtype=torch.long, + device=self._module_device(self.talker), + )) + self.embed_text_spk_tokens = { + key: + self.thinker_embedding( + torch.tensor( + [value], + dtype=torch.long, + device=self._module_device(self.talker), + )) + for key, value in self.tts_text_spk_token_ids.items() + } + self.embed_text_eos_token = self.thinker_embedding( + torch.tensor( + [talker_hf_config.tts_text_end_token_id], + dtype=torch.long, + device=self._module_device(self.talker), + )) + self.embed_text_pad_token = self.thinker_embedding( + torch.tensor( + [talker_hf_config.tts_text_pad_token_id], + dtype=torch.long, + device=self._module_device(self.talker), + )) + self.embed_codec_bos_token = self.talker_embedding( + torch.tensor( + [talker_hf_config.tts_codec_start_token_id], + dtype=torch.long, + device=self._module_device(self.talker), + )) + self.embed_codec_eos_token = self.talker_embedding( + torch.tensor( + [talker_hf_config.tts_codec_end_token_id], + dtype=torch.long, + device=self._module_device(self.talker), + )) + self.embed_codec_pad_token = self.talker_embedding( + torch.tensor( + [talker_hf_config.tts_codec_pad_token_id], + dtype=torch.long, + device=self._module_device(self.talker), + )) + return set(["thinker_embedding.weight", "talker_embedding.weight"]) + + def _get_embed_text_spk_token(self, voice_type: str): + if voice_type not in self.embed_text_spk_tokens: + return self.embed_text_bos_token + return self.embed_text_spk_tokens[voice_type] + + def _get_text_spk_token_id(self, voice_type: str): + talker_hf_config = self.talker_config + if hasattr(talker_hf_config, 'talker_config'): + talker_hf_config = talker_hf_config.talker_config + + if voice_type not in self.tts_text_spk_token_ids: + return talker_hf_config.tts_text_start_token_id + return self.tts_text_spk_token_ids[voice_type] + + def _thinker_to_talker_prefill( + self, + voice_type: str, + output_prompt_embeds, + output_token_ids, + thinker_prompt_embeds, + prompt_token_ids + ): + + talker_hf_config = self.talker_config + if hasattr(talker_hf_config, 'talker_config'): + talker_hf_config = talker_hf_config.talker_config + + # if len(output.outputs[0].token_ids) == 2: + # issue request + prompt_embeds = torch.cat([ + thinker_prompt_embeds, + self._get_embed_text_spk_token(voice_type) + + self.embed_codec_pad_token, + output_prompt_embeds[:1] + self.embed_codec_bos_token, + ], + dim=0) + + prompt_token_ids_processed = prompt_token_ids + [ + talker_hf_config.tts_codec_pad_token_id, + output_token_ids[0], + + ] + input_tokens_len = len(prompt_token_ids_processed) + # the code below is from model runner in Qwen, may need to further discuss later + if input_tokens_len > 2: + prompt_token_ids_processed = ( + [self.talker_config.tts_codec_mask_token_id] * + (input_tokens_len - 2) + [ + self.talker_config.tts_codec_pad_token_id, + self.talker_config.tts_codec_start_token_id + ]) + else: + prompt_token_ids_processed = [ + self.talker_config.tts_codec_pad_token_id, + self.talker_config.tts_codec_start_token_id, + ][-input_tokens_len:] + if isinstance(prompt_token_ids_processed,list): + prompt_token_ids_processed = torch.Tensor(prompt_token_ids_processed).to(torch.int64).to(self._module_device(self.talker)) + return prompt_token_ids_processed, prompt_embeds + + def _thinker_to_talker_decode_one_step( + self, + output_prompt_embeds, + output_token_ids, + ): + processed_output_token_embeds = output_prompt_embeds + self.talker.get_input_embeddings(output_token_ids) #for decode + return output_token_ids, processed_output_token_embeds + + def compute_logits( + self, + hidden_states: Union[torch.Tensor, OmniOutput], + sampling_metadata: SamplingMetadata, + ) -> Optional[torch.Tensor]: + # Handle OmniOutput type + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + + # Use thinker model for logits computation + return self.model.compute_logits(hidden_states, sampling_metadata) + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[SamplerOutput]: + # Use thinker model for sampling + return self.model.sample(logits, sampling_metadata) + + def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default"): + """ + Generate speech from text tokens using the talker and token2wav models. + This method is kept for backward compatibility and direct speech generation. + + Args: + text_tokens: Text tokens from thinker model + voice_type: Voice type for speech generation + + Returns: + Audio tensor + """ + # Generate codec tokens using talker model + talker_output = self.talker( + input_ids=None, + positions=None, + inputs_embeds=text_tokens + ) + + # Convert talker output to codec tokens + codec_tokens = self._convert_to_codec_tokens(talker_output) + + # Generate audio using token2wav model + return self._codec_to_audio(codec_tokens, voice_type=voice_type) + + + def _convert_to_codec_tokens(self, talker_output: torch.Tensor, sampling_metadata: SamplingMetadata) -> torch.Tensor: + """ + Reference (HF): use the talker's codec head to obtain logits, suppress BOS, + then greedily select the next codec token for the current step. + """ + with torch.inference_mode(): + logits = self.talker.compute_logits(talker_output, None) + if logits is None: + return torch.zeros((talker_output.size(0), 0), dtype=torch.long, device=talker_output.device) + + # Suppress only codec_bos, consistent with HF generate's suppress_tokens behavior + bos_id = None + if hasattr(self, 'talker_config') and hasattr(self.talker_config, 'tts_codec_start_token_id'): + bos_id = int(getattr(self.talker_config, 'tts_codec_start_token_id')) + if bos_id is not None: + logits[..., bos_id] = -1e9 + + # Take the distribution at the last step and select greedily + next_id = self.talker.sample(logits, sampling_metadata).sampled_token_ids + return next_id.to(dtype=torch.long) + + def _init_token2wav_model(self): + """Initialize speaker resources if provided; model is constructed in __init__.""" + if self.token2wav is None or self.token2wav_config is None: + return + device = 'cuda' if torch.cuda.is_available() else 'cpu' + # optional speaker resources + conds = getattr(self.token2wav_config, 'conds', None) + ref_mels = getattr(self.token2wav_config, 'ref_mels', None) + if isinstance(conds, dict) and isinstance(ref_mels, dict): + self._token2wav_conds = {k: torch.as_tensor(v, device=device) for k, v in conds.items()} + self._token2wav_ref_mels = {k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()} + # legacy: load from directory if provided + model_path = self.vllm_config.model_config.model + if isinstance(model_path, str) and os.path.isdir(model_path): + spk_pt = os.path.join(model_path, 'spk_dict.pt') + if os.path.exists(spk_pt): + data = torch.load(spk_pt, map_location=device) + for key, value in data.items(): + self._token2wav_conds[key] = value["cond"].to(device) + self._token2wav_ref_mels[key] = value["ref_mel"].to(device) + else: + # legacy npy inputs + for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*spk_emb.npy'))): + key = os.path.basename(f).split('_')[0].lower() + self._token2wav_conds[key] = torch.as_tensor(np.load(f), device=device) + for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*ref_mel.npy'))): + key = os.path.basename(f).split('_')[0].lower() + self._token2wav_ref_mels[key] = torch.as_tensor(np.load(f), device=device) + + def _codec_to_audio(self, codec_tokens: torch.Tensor, voice_type: str = "default") -> Optional[torch.Tensor]: + if self.token2wav is None: + self._init_token2wav_model() + if self.token2wav is None: + return None + # Normalize voice type + voice = (voice_type or 'default') + # Resolve cond / ref_mel if provided + cond = None + ref_mel = None + if voice in self._token2wav_conds and voice in self._token2wav_ref_mels: + cond = self._token2wav_conds[voice] + ref_mel = self._token2wav_ref_mels[voice] + # Fallback: create dummy cond/ref_mel if not provided + token2wav_dev = self._module_device(self.token2wav) + if cond is None: + cond = torch.zeros((1, self.token2wav_config.dit_config.enc_emb_dim), device=token2wav_dev, dtype=torch.float32) + if ref_mel is None: + ref_mel = torch.zeros((1, 300, self.token2wav_config.dit_config.mel_dim), device=token2wav_dev, dtype=torch.float32) + + # Ensure codec is (1, T) long tensor on correct device + if isinstance(codec_tokens, torch.Tensor): + codec = codec_tokens.to(dtype=torch.long, device=token2wav_dev) + if codec.ndim == 1: + codec = codec.unsqueeze(0) + else: + codec = torch.as_tensor(codec_tokens, dtype=torch.long, device=token2wav_dev).unsqueeze(0) + + # Streaming with chunked process and boundary alignment (rely on token2wav.process_chunk) + factor = getattr(self.token2wav.token2wav.factor, 'factor', 2) + chunk_size = 48 + mel_dim = getattr(self.token2wav.token2wav.code2wav_dit_model, 'mel_dim', self.token2wav_config.dit_config.mel_dim) + total_mel = int(codec.shape[1] * factor) + steps = 10 + + # Prepare initial noise for the whole sequence + y_all = torch.randn((1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev) + + logger.info(f"Currently, we do not use the chunked process, we only use the token2wav.process_chunk for the whole sequence.\ + The stream mode will be implemented in the future.") + + chunk_ends = [] + for i in range(codec.shape[1]): + chunk_code_length = i * 2 - 24 + finished = i==(codec.shape[1]-1) + if (chunk_code_length > 0 and + chunk_code_length % chunk_size == 0) or finished: + chunk_ends.append(i) + + # Number of chunks in mel domain + prev_generated = None + wav_chunks: list = [] + prev_id = 0 + + with torch.inference_mode(): + for n,i in enumerate([0]): + finished = (i == codec.shape[1] - 1) + _, audio_chunk = self.token2wav.process_chunk( + conditioning=cond, + reference_mel=ref_mel, + codec_all=codec, + y_all=y_all, + i=n, + steps=steps, + prev_generated=prev_generated if prev_generated is not None else [], + finished=True, + ) + prev_generated = audio_chunk + wav_chunks.append(audio_chunk.detach().cpu().numpy()) + prev_id = i + + if len(wav_chunks) == 0: + return torch.zeros(0, device=token2wav_dev) + + waveform = np.concatenate(wav_chunks) + return torch.as_tensor(waveform, device=token2wav_dev) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]: + """Load weights for all components of the omni model.""" + loaded_weights = set() + thinker_weights = [] + talker_weights = [] + token2wav_weights = [] + for k, v in weights: + if k.startswith('thinker.'): + thinker_weights.append((k, v)) + elif k.startswith('talker.'): + talker_weights.append((k, v)) + elif k.startswith('token2wav.'): + token2wav_weights.append((k, v)) + else: + raise ValueError(f"Unknown weight prefix: {k}") + + # Load thinker weights + if self.thinker: + if thinker_weights: + thinker_loaded = self.thinker.load_weights(thinker_weights) + else: + thinker_loaded = set([k for k,v in thinker_weights]) + thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, 'thinker') + loaded_weights.update(thinker_loaded) + + + # Load talker weights + if talker_weights and self.talker is not None: + # Map talker weights to appropriate components + if self.thinker is None: + thinker_embedding_weights = [w for n, w in thinker_weights if n == "thinker.model.embed_tokens.weight"] + if thinker_embedding_weights: + self.thinker_embedding = nn.Embedding(thinker_embedding_weights[0].shape[0], thinker_embedding_weights[0].shape[1]) + self.thinker_embedding.weight = nn.Parameter(thinker_embedding_weights[0].to(self._module_device(self.talker))) + talker_loaded = self.talker.load_weights(talker_weights) + talker_loaded = add_prefix_to_loaded_weights(talker_loaded, 'talker') + loaded_weights.update(talker_loaded) + loaded_weights.update(self._init_special_tokens_embeddings()) + + + # Load token2wav weights (if any) + if token2wav_weights and self.token2wav is not None: + self._init_token2wav_model() + t2w_loaded = self.token2wav.load_weights(token2wav_weights, os.path.join(self.vllm_config.model_config.model, "spk_dict.pt")) + t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, 'token2wav') + loaded_weights.update(t2w_loaded) + + return loaded_weights \ No newline at end of file diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni_talker.py new file mode 100644 index 000000000..8e6f95d6f --- /dev/null +++ b/vllm_omni/model_executor/models/qwen2_5_omni_talker.py @@ -0,0 +1,239 @@ +from functools import cached_property +from typing import Iterable, List, Optional, Set, Tuple, Union + +import torch +import torch.nn as nn +from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniTalkerConfig +from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder + +from vllm.attention import AttentionMetadata +from vllm.config import VllmConfig +from vllm.forward_context import ForwardContext, get_forward_context +from vllm.model_executor.layers.linear import ColumnParallelLinear +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler +from vllm.model_executor.sampling_metadata import SamplingMetadata +from vllm.model_executor.models.qwen2_5_omni_thinker import ( + Qwen2_5OmniConditionalGenerationMixin, + Qwen2_5OmniThinkerMultiModalProcessor, + Qwen2_5OmniThinkerProcessingInfo, + Qwen2_5OmniThinkerDummyInputsBuilder) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.logger import init_logger +from vllm.sequence import IntermediateTensors + +from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP, MultiModalEmbeddings +from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper, + init_vllm_registered_model, + maybe_prefix, + merge_multimodal_embeddings) +from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen2_5OmniThinkerMultiModalProcessor, + info=Qwen2_5OmniThinkerProcessingInfo, + dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder, +) +class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal, + SupportsPP, + Qwen2_5OmniConditionalGenerationMixin): + logger = init_logger(__name__) + # Align to thinker-style static mapper for clarity + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + # text LM head/body in talker + "talker.codec_head.": "language_model.lm_head.", + "talker.model.": "language_model.model.", + # projection weights + "talker.thinker_to_talker_proj.": "thinker_to_talker_proj.", + # fallback root + "talker.": "", + }) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config: Qwen2_5OmniTalkerConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + + self.prefix = prefix + self.quant_config = quant_config + + if hasattr(config, "talker_config"): + self.config = config.talker_config + vllm_config.model_config.hf_text_config = \ + vllm_config.model_config.hf_config.talker_config + else: + self.config = config + + self.thinker_to_talker_proj = ColumnParallelLinear( + self.config.embedding_size, + self.config.hidden_size, + bias=True, + gather_output=True, + skip_bias_add=False, + quant_config=quant_config, + ) + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + hf_config=getattr(self.config, 'text_config', self.config), + architectures=["Qwen2ForCausalLM_old"], + ) + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors) + + def init_multi_modal(self, thinker_config): + self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config) + self.visual = Qwen2_5_VisionTransformer( + vision_config=thinker_config.vision_config, + norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), + quant_config=self.quant_config, + prefix=maybe_prefix(self.prefix, "visual"), + ) + + def get_language_model(self) -> torch.nn.Module: + return self.language_model + + @cached_property + def sampler(self): + if hasattr(self.language_model, "sampler"): + return self.language_model.sampler + + return get_sampler() + + def get_input_embeddings( + self, + input_ids: torch.Tensor, + multimodal_embeddings: Optional[MultiModalEmbeddings] = None, + ) -> torch.Tensor: + inputs_embeds = self.language_model.get_input_embeddings(input_ids) + if multimodal_embeddings is not None \ + and len(multimodal_embeddings) != 0: + + # TODO (ywang96): support overlapping modalitiy embeddings so that + # `use_audio_in_video` will work on V1. + inputs_embeds = merge_multimodal_embeddings( + input_ids, inputs_embeds, multimodal_embeddings, [ + self.config.image_token_index, + self.config.video_token_index, + self.config.audio_token_index + ]) + return inputs_embeds + + def forward( + self, + input_ids: torch.Tensor = None, + positions: torch.Tensor = None, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + **kwargs: object, + ) -> Union[torch.Tensor, IntermediateTensors]: + assert input_ids is not None or inputs_embeds is not None, "input_ids or inputs_embeds must be provided" + forward_context: ForwardContext = get_forward_context() + attn_metadata: AttentionMetadata = forward_context.attn_metadata + + if intermediate_tensors is not None: + inputs_embeds = None + elif inputs_embeds is None: + # for profile_run: + inputs_embeds = self.get_input_embeddings(input_ids) + + input_ids = None + + # projection + inputs_embeds, _ = self.thinker_to_talker_proj(inputs_embeds) + + hidden_states = self.language_model.model(input_ids, + positions, + intermediate_tensors, + inputs_embeds=inputs_embeds) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[torch.Tensor]: + return self.language_model.compute_logits(hidden_states, + sampling_metadata) + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[SamplerOutput]: + return self.language_model.sample(logits, sampling_metadata) + + def load_weights(self, weights: Iterable[Tuple[str, + torch.Tensor]]) -> Set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=['thinker.', 'token2wav.'], + ) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + # Log load summary + try: + total_bytes = 0 + for name, param in self.named_parameters(): + if param is not None and param.data is not None: + total_bytes += param.data.numel() * param.data.element_size() + device = next(self.parameters()).device + self.logger.info( + "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s", + self.__class__.__name__, True, total_bytes / (1024**2), str(device)) + except Exception: + pass + multi_model_weights = set() + for name, param in self.visual.named_parameters(): + multi_model_weights.add("visual."+name) + for name, param in self.audio_tower.named_parameters(): + multi_model_weights.add("audio_tower."+name) + loaded.update(multi_model_weights) + return loaded + + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: + mm_input_by_modality = {} + + # Preserve the order of modalities if there are multiple of them + # from the order of kwargs. + for input_key in kwargs: + if input_key in ("pixel_values", "image_embeds" + ) and "image" not in mm_input_by_modality: + mm_input_by_modality[ + "image"] = self._parse_and_validate_image_input(**kwargs) + if input_key in ("pixel_values_videos", "video_embeds" + ) and "video" not in mm_input_by_modality: + mm_input_by_modality[ + "video"] = self._parse_and_validate_video_input(**kwargs) + if input_key in ("input_audio_features" + ) and "audio" not in mm_input_by_modality: + mm_input_by_modality[ + "audio"] = self._parse_and_validate_audio_input(**kwargs) + return mm_input_by_modality + + def get_multimodal_embeddings(self, + **kwargs: object) -> MultiModalEmbeddings: + + mm_input_by_modality = self._parse_and_validate_multimodal_inputs( + **kwargs) + if not mm_input_by_modality: + return [] + + # The result multimodal_embeddings is tuple of tensors, with each + # tensor correspoending to a multimodal data item (image or video). + multimodal_embeddings: tuple[torch.Tensor, ...] = () + + # NOTE: It is important to iterate over the keys in this dictionary + # to preserve the order of the modalities. + for modality in mm_input_by_modality: + multimodal_input = mm_input_by_modality[modality] + if modality == "image": + vision_embeddings = self._process_image_input(multimodal_input) + multimodal_embeddings += vision_embeddings + if modality == "video": + video_embeddings = self._process_video_input(multimodal_input) + multimodal_embeddings += video_embeddings + if modality == "audio": + audio_embeddings = self._process_audio_input(multimodal_input) + multimodal_embeddings += audio_embeddings + return multimodal_embeddings \ No newline at end of file diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py new file mode 100644 index 000000000..f1eb2c005 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py @@ -0,0 +1,889 @@ +from collections.abc import Iterable, Mapping, Sequence +from copy import copy +from functools import partial +from typing import Any, Optional, Union + +import torch +import torch.nn as nn +from transformers.feature_extraction_utils import BatchFeature +from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( + Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig) +from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import ( + Qwen2_5OmniAudioEncoder) +from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import ( + Qwen2_5OmniProcessor) +from transformers.models.whisper import WhisperFeatureExtractor + +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from vllm.model_executor.models.qwen2_5_vl import ( + Qwen2_5_VisionTransformer, Qwen2_5_VLImageEmbeddingInputs, + Qwen2_5_VLImageInputs, Qwen2_5_VLImagePixelInputs, + Qwen2_5_VLProcessingInfo, Qwen2_5_VLVideoEmbeddingInputs, + Qwen2_5_VLVideoInputs, Qwen2_5_VLVideoPixelInputs) +from vllm.model_executor.models.qwen2_audio import ( + Qwen2AudioInputs, Qwen2AudioProcessingInfo, + _get_feat_extract_output_lengths) +from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser +from vllm.model_executor.sampling_metadata import SamplingMetadata +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import (ImageItem, ModalityData, + MultiModalDataDict, MultiModalFieldConfig, + MultiModalKwargs, NestedTensors) +from vllm.multimodal.parse import (AudioProcessorItems, DictEmbeddingItems, + ModalityDataItems, MultiModalDataItems, + MultiModalDataParser) +from vllm.multimodal.processing import (BaseMultiModalProcessor, + PlaceholderFeaturesInfo, + PromptReplacement, PromptUpdate) +from vllm.multimodal.profiling import BaseDummyInputsBuilder +from vllm.sequence import IntermediateTensors +from vllm.transformers_utils.tokenizer import decode_tokens, encode_tokens + +from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP +from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper, + init_vllm_registered_model, maybe_prefix, + merge_multimodal_embeddings) + +from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding + +try: + import flash_attn +except (ImportError, ModuleNotFoundError): + flash_attn = None + +logger = init_logger(__name__) + + +def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]): + audio_feature_lengths = hf_inputs.get("audio_feature_lengths", + torch.empty((0, ))) + + image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3))) + image_grid_sizes = image_grid_thw.prod(-1) + + video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3))) + video_grid_sizes = video_grid_thw.prod(-1) + + return dict( + input_audio_features=MultiModalFieldConfig.flat_from_sizes( + "audio", audio_feature_lengths, dim=1), + feature_attention_mask=MultiModalFieldConfig.batched("audio"), + audio_feature_lengths=MultiModalFieldConfig.batched("audio"), + pixel_values=MultiModalFieldConfig.flat_from_sizes( + "image", image_grid_sizes), + image_embeds=MultiModalFieldConfig.flat_from_sizes( + "image", image_grid_sizes), + image_grid_thw=MultiModalFieldConfig.batched("image"), + pixel_values_videos=MultiModalFieldConfig.flat_from_sizes( + "video", video_grid_sizes), + video_embeds=MultiModalFieldConfig.flat_from_sizes( + "video", video_grid_sizes), + video_grid_thw=MultiModalFieldConfig.batched("video"), + second_per_grid_ts=MultiModalFieldConfig.batched("video"), + ) + + +class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser): + + def _parse_audio_data( + self, + data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]], + ) -> ModalityDataItems[Any, Any]: + if isinstance(data, dict): + return DictEmbeddingItems( + data, + modality="audio", + required_fields={ + "input_audio_features", "audio_feature_lengths" + }, + fields_factory=_qwen2_5_omni_thinker_field_config, + ) + + return super()._parse_audio_data(data) + + +class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo, + Qwen2_5_VLProcessingInfo): + + def get_hf_config(self): + return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config + + def get_hf_processor(self, **kwargs: object) -> Qwen2_5OmniProcessor: + return self.ctx.get_hf_processor( + Qwen2_5OmniProcessor, + use_fast=kwargs.pop("use_fast", True), + **kwargs, + ) + + def get_feature_extractor(self, **kwargs: object): + hf_processor = self.get_hf_processor(**kwargs) + feature_extractor = hf_processor.feature_extractor # type: ignore + assert isinstance(feature_extractor, WhisperFeatureExtractor) + return feature_extractor + + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]: + return {"audio": None, "image": None, "video": None} + + +class Qwen2_5OmniThinkerDummyInputsBuilder( + BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]): + + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_audios = mm_counts.get("audio", 0) + num_images = mm_counts.get("image", 0) + num_videos = mm_counts.get("video", 0) + + hf_processor = self.info.get_hf_processor() + + audio_token: str = hf_processor.audio_token + image_token: str = hf_processor.image_token + video_token: str = hf_processor.video_token + + return (audio_token * num_audios + image_token * num_images + + video_token * num_videos) + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> MultiModalDataDict: + num_audios = mm_counts.get("audio", 0) + num_images = mm_counts.get("image", 0) + num_videos = mm_counts.get("video", 0) + + feature_extractor = self.info.get_feature_extractor() + + target_audio_length = min( + feature_extractor.chunk_length, + 30, + ) * feature_extractor.sampling_rate + target_width, target_height = \ + self.info.get_image_size_with_most_features() + target_num_frames = \ + self.info.get_num_frames_with_most_features(seq_len, mm_counts) + + mm_data = { + "audio": + self._get_dummy_audios(length=target_audio_length, + num_audios=num_audios), + "image": + self._get_dummy_images(width=target_width, + height=target_height, + num_images=num_images), + "video": + self._get_dummy_videos(width=target_width, + height=target_height, + num_frames=target_num_frames, + num_videos=num_videos), + } + + return mm_data + + +class Qwen2_5OmniThinkerMultiModalProcessor( + BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]): + + def _get_data_parser(self) -> MultiModalDataParser: + feature_extractor = self.info.get_feature_extractor() + return Qwen2_5OmniThinkerMultiModalDataParser( + target_sr=feature_extractor.sampling_rate) + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + mm_data = dict(mm_data) + audios = mm_data.pop("audios", []) + + # NOTE: WhisperFeatureExtractor cannot handle empty list of audios + if audios: + # NOTE: Qwen2.5-Omni processor accept "audio" + mm_data["audio"] = audios + mm_kwargs = dict(**mm_kwargs, ) + + hf_inputs = super()._call_hf_processor( + prompt=prompt, + mm_data=mm_data, + mm_kwargs=mm_kwargs, + tok_kwargs=tok_kwargs, + ) + + input_features = hf_inputs.pop('input_features', None) + feature_attention_mask = hf_inputs.get('feature_attention_mask', None) + if ('input_audio_features' not in hf_inputs + and input_features is not None): + if feature_attention_mask is not None: + input_features = input_features.permute( + 0, 2, 1)[feature_attention_mask.bool()].permute(1, 0) + hf_inputs['input_audio_features'] = input_features + if ('audio_feature_lengths' not in hf_inputs + and feature_attention_mask is not None): + hf_inputs['audio_feature_lengths'] = feature_attention_mask.sum(-1) + return hf_inputs + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return _qwen2_5_omni_thinker_field_config(hf_inputs) + + def _maybe_apply_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + prompt_ids: list[int], + mm_kwargs: MultiModalKwargs, + is_update_applied: bool, + ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]: + """ + Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`. + """ + unbound_prompt_updates = self._get_prompt_updates( + mm_items, + hf_processor_mm_kwargs, + mm_kwargs, + ) + mm_prompt_updates = self._bind_and_group_updates( + unbound_prompt_updates) + + mm_item_counts = mm_items.get_all_counts() + self._validate_mm_kwargs(mm_kwargs, mm_item_counts) + + use_audio_in_video = hf_processor_mm_kwargs.get( + "use_audio_in_video", False) + + if is_update_applied: + mm_placeholders = self._find_mm_placeholders( + mm_prompt_updates, + prompt_ids, + mm_item_counts, + ) + self._validate_mm_placeholders( + mm_placeholders, + mm_item_counts, + use_audio_in_video=use_audio_in_video) + + tokenizer = self.info.get_tokenizer() + prompt = decode_tokens(tokenizer, prompt_ids) + else: + ( + prompt_ids, + prompt, + mm_placeholders, + ) = self._apply_prompt_updates( + prompt_ids, + mm_prompt_updates, + mm_item_counts, + ) + self._validate_mm_placeholders( + mm_placeholders, + mm_item_counts, + use_audio_in_video=use_audio_in_video) + + tokenizer = self.info.get_tokenizer() + prompt = decode_tokens(tokenizer, prompt_ids) + + if use_audio_in_video: + mm_kwargs["use_audio_in_video"] = True + + return prompt_ids, prompt, mm_placeholders + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargs, + ) -> Sequence[PromptUpdate]: + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + tokenizer = self.info.get_tokenizer() + image_processor = self.info.get_image_processor( + **hf_processor_mm_kwargs) + vocab = tokenizer.get_vocab() + + audio_token = processor.audio_token + image_token = processor.image_token + video_token = processor.video_token + audio_token_id = vocab[audio_token] + image_token_id = vocab[image_token] + video_token_id = vocab[video_token] + + audio_feature_lengths = out_mm_kwargs.get("audio_feature_lengths") + feature_attention_mask = out_mm_kwargs.get("feature_attention_mask") + if audio_feature_lengths is None and feature_attention_mask is None: + audio_output_lengths = [] + elif audio_feature_lengths is not None: + _, audio_output_lens = _get_feat_extract_output_lengths( + audio_feature_lengths) + audio_output_lengths = audio_output_lens.tolist() + elif feature_attention_mask is not None: + assert isinstance(feature_attention_mask, torch.Tensor) + _, audio_output_lens = _get_feat_extract_output_lengths( + feature_attention_mask.sum(-1)) + audio_output_lengths = audio_output_lens.tolist() + + # number of audios read from video. + audio_in_video_item_idx = 0 + + def get_replacement_qwen2_audio(item_idx: int): + item_idx += audio_in_video_item_idx + + num_features = audio_output_lengths[item_idx] + if num_features == 0: + audios = mm_items.get_items("audio", AudioProcessorItems) + audio = audios.get(item_idx) + raise ValueError( + f"The audio {audio} (len={len(audio)}) is too short " + "to be represented inside the model") + + return [audio_token_id] * num_features + + def get_replacement_qwen2_vision(item_idx: int, modality: str): + grid_thw = out_mm_kwargs[f"{modality}_grid_thw"][item_idx] + assert isinstance(grid_thw, torch.Tensor) + merge_length = image_processor.merge_size**2 + + token_id = image_token_id if modality == "image" else video_token_id + return [token_id] * (int(grid_thw.prod()) // merge_length) + + use_audio_in_video = hf_processor_mm_kwargs.get( + "use_audio_in_video", False) + thinker_config = self.info.get_hf_config() + + def get_replacement_qwen2_use_audio_in_video(item_idx: int): + nonlocal audio_in_video_item_idx + + audio_num_features = audio_output_lengths[audio_in_video_item_idx + + item_idx] + video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx] + + audio_in_video_item_idx += 1 + + second_per_grid_ts = hf_processor_mm_kwargs.get( + "second_per_grid_ts", None) + if second_per_grid_ts: + video_second_per_grid_t = second_per_grid_ts[item_idx] + else: + video_second_per_grid_t = 1.0 + + return MRotaryEmbedding.omni_get_updates_use_audio_in_video( + thinker_config=thinker_config, + audio_len=audio_num_features, + video_grid_thw=video_grid_thw, + video_second_per_grid_t=video_second_per_grid_t, + ) + + video_replacement_fn = ( + get_replacement_qwen2_use_audio_in_video if use_audio_in_video else + partial(get_replacement_qwen2_vision, modality="video")) + + return [ + PromptReplacement( + modality="audio", + target=audio_token, + replacement=get_replacement_qwen2_audio, + ), + PromptReplacement( + modality="image", + target=image_token, + replacement=partial(get_replacement_qwen2_vision, + modality="image"), + ), + PromptReplacement( + modality="video", + target=video_token, + replacement=video_replacement_fn, + ), + ] + + def _apply_hf_processor_main( + self, + prompt: Union[str, list[int]], + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + *, + enable_hf_prompt_update: bool, + ) -> tuple[list[int], BatchFeature, bool]: + """ + Qwen2.5-Omni reimplements this function to handle text only. + """ + if isinstance(prompt, str): + if enable_hf_prompt_update: + return self._apply_hf_processor_text_mm( + prompt_text=prompt, + mm_items=mm_items, + hf_processor_mm_kwargs=hf_processor_mm_kwargs, + tokenization_kwargs=tokenization_kwargs, + ) + tokenizer = self.info.get_tokenizer() + prompt_ids = encode_tokens(tokenizer, prompt) + else: + prompt_ids = self._apply_hf_processor_tokens_only(prompt) + + mm_processed_data = self._apply_hf_processor_mm_only( + mm_items=mm_items, + hf_processor_mm_kwargs=hf_processor_mm_kwargs, + tokenization_kwargs=tokenization_kwargs, + ) + + return prompt_ids, mm_processed_data, False + + def _apply_hf_processor_mm_only( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> BatchFeature: + """ + Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`. + """ + mm_counts = mm_items.get_all_counts() + + use_audio_in_video = hf_processor_mm_kwargs.get( + "use_audio_in_video", False) + if use_audio_in_video and "video" in mm_counts: + assert "audio" in mm_counts + mm_counts["audio"] -= mm_counts["video"] + + _, mm_processed_data, _ = self._apply_hf_processor_text_mm( + prompt_text=self.dummy_inputs.get_dummy_text(mm_counts), + mm_items=mm_items, + hf_processor_mm_kwargs=hf_processor_mm_kwargs, + tokenization_kwargs=tokenization_kwargs, + ) + + return mm_processed_data + + def _validate_mm_placeholders( + self, + mm_placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], + mm_item_counts: Mapping[str, int], + use_audio_in_video: bool = False, + ) -> None: + if use_audio_in_video: + mm_item_counts = copy(mm_item_counts) + if "video" in mm_item_counts: + assert "audio" in mm_item_counts + mm_item_counts["audio"] -= mm_item_counts["video"] + super()._validate_mm_placeholders(mm_placeholders, mm_item_counts) + + +class Qwen2_5OmniConditionalGenerationMixin: + + def _validate_and_reshape_mm_tensor(self, + mm_input: object, + name: str, + dim: int = 0) -> torch.Tensor: + if not isinstance(mm_input, (torch.Tensor, list)): + raise ValueError(f"Incorrect type of {name}. " + f"Got type: {type(mm_input)}") + if isinstance(mm_input, torch.Tensor): + return torch.concat(list(mm_input), dim=dim) + else: + return torch.concat(mm_input, dim=dim) + + def _parse_and_validate_audio_input( + self, **kwargs: object) -> Optional[Qwen2AudioInputs]: + input_audio_features = kwargs.pop('input_audio_features', None) + audio_feature_lengths = kwargs.pop('audio_feature_lengths', None) + feature_attention_mask = kwargs.pop('feature_attention_mask', None) + if input_audio_features is None: + return None + input_audio_features = self._validate_and_reshape_mm_tensor( + input_audio_features, 'input_audio_features', dim=1) + if feature_attention_mask is not None: + feature_attention_mask = self._validate_and_reshape_mm_tensor( + feature_attention_mask, 'feature_attention_mask') + if not isinstance(input_audio_features, (torch.Tensor, list)): + raise ValueError("Incorrect type of audio input features. " + f"Got type: {type(input_audio_features)}") + return Qwen2AudioInputs(input_features=input_audio_features, + audio_feature_lengths=audio_feature_lengths, + feature_attention_mask=feature_attention_mask) + + def _parse_and_validate_image_input( + self, + **kwargs: dict[str, Any], + ) -> Optional[Qwen2_5_VLImageInputs]: + pixel_values = kwargs.pop("pixel_values", None) + image_embeds = kwargs.pop("image_embeds", None) + image_grid_thw = kwargs.pop("image_grid_thw", None) + + if pixel_values is None and image_embeds is None: + return None + + if pixel_values is not None: + pixel_values = self._validate_and_reshape_mm_tensor( + pixel_values, "image pixel values") + image_grid_thw = self._validate_and_reshape_mm_tensor( + image_grid_thw, "image grid_thw") + + if not isinstance(pixel_values, (torch.Tensor, list)): + raise ValueError("Incorrect type of image pixel values. " + f"Got type: {type(pixel_values)}") + + return Qwen2_5_VLImagePixelInputs(type="pixel_values", + pixel_values=pixel_values, + image_grid_thw=image_grid_thw) + + if image_embeds is not None: + image_embeds = self._validate_and_reshape_mm_tensor( + image_embeds, "image embeds") + image_grid_thw = self._validate_and_reshape_mm_tensor( + image_grid_thw, "image grid_thw") + + if not isinstance(image_embeds, torch.Tensor): + raise ValueError("Incorrect type of image embeddings. " + f"Got type: {type(image_embeds)}") + return Qwen2_5_VLImageEmbeddingInputs( + type="image_embeds", + image_embeds=image_embeds, + image_grid_thw=image_grid_thw) + + def _parse_and_validate_video_input( + self, + **kwargs: dict[str, Any], + ) -> Optional[Qwen2_5_VLVideoInputs]: + pixel_values_videos = kwargs.pop("pixel_values_videos", None) + video_embeds = kwargs.pop("video_embeds", None) + video_grid_thw = kwargs.pop("video_grid_thw", None) + + if pixel_values_videos is None and video_embeds is None: + return None + + if pixel_values_videos is not None: + pixel_values_videos = self._validate_and_reshape_mm_tensor( + pixel_values_videos, "video pixel values") + video_grid_thw = self._validate_and_reshape_mm_tensor( + video_grid_thw, "video grid_thw") + + return Qwen2_5_VLVideoPixelInputs( + type="pixel_values_videos", + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + ) + + if video_embeds is not None: + video_embeds = self._validate_and_reshape_mm_tensor( + video_embeds, "video embeds") + video_grid_thw = self._validate_and_reshape_mm_tensor( + video_grid_thw, "video grid_thw") + + if not isinstance(video_embeds, torch.Tensor): + raise ValueError("Incorrect type of video embeddings. " + f"Got type: {type(video_embeds)}") + return Qwen2_5_VLVideoEmbeddingInputs( + type="video_embeds", + video_embeds=video_embeds, + video_grid_thw=video_grid_thw) + + def _process_audio_input( + self, + audio_input: Qwen2AudioInputs, + audio_hashes: list[str] = None, + cached_audio_features: torch.Tensor = None, + ) -> torch.Tensor: + + input_features = audio_input["input_features"] + audio_feature_lengths = audio_input["audio_feature_lengths"] + if input_features.ndim == 3: + assert input_features.shape[0] == 1 + input_features = input_features.squeeze(0) + if audio_feature_lengths.ndim == 2: + assert audio_feature_lengths.shape[ + 0] == 1 or audio_feature_lengths.shape[1] == 1 + if audio_feature_lengths.shape[0] == 1: + audio_feature_lengths = audio_feature_lengths.squeeze(0) + else: + audio_feature_lengths = audio_feature_lengths.squeeze(1) + + audio_feat_lengths, audio_output_lengths = ( + self.audio_tower._get_feat_extract_output_lengths( + audio_feature_lengths)) + + audio_outputs = self.audio_tower( + input_features.to(self.audio_tower.dtype), + feature_lens=audio_feature_lengths, + aftercnn_lens=audio_feat_lengths, + ) + audio_features = audio_outputs.last_hidden_state + return audio_features.split(audio_output_lengths.tolist()) + + def _process_image_input( + self, + image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]: + if image_input["type"] == "image_embeds": + return image_input["image_embeds"].type(self.visual.dtype) + + grid_thw = image_input["image_grid_thw"] + assert grid_thw.ndim == 2 + + pixel_values = image_input["pixel_values"].type(self.visual.dtype) + image_embeds = self.visual(pixel_values, grid_thw=grid_thw) + # Split concatenated embeddings for each image item. + merge_size = self.visual.spatial_merge_size + sizes = grid_thw.prod(-1) // merge_size // merge_size + + return image_embeds.split(sizes.tolist()) + + def _process_video_input( + self, + video_input: Qwen2_5_VLVideoInputs, + video_hashes: list[str] = None, + cached_video_embeds: torch.Tensor = None) -> torch.Tensor: + if video_input["type"] == "video_embeds": + return video_input["video_embeds"].type(self.visual.dtype) + + grid_thw = video_input["video_grid_thw"] + assert grid_thw.ndim == 2 + + pixel_values_videos = video_input["pixel_values_videos"].type( + self.visual.dtype) + video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw) + # Split concatenated embeddings for each video item. + merge_size = self.visual.spatial_merge_size + sizes = grid_thw.prod(-1) // merge_size // merge_size + + return video_embeds.split(sizes.tolist()) + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen2_5OmniThinkerMultiModalProcessor, + info=Qwen2_5OmniThinkerProcessingInfo, + dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder, +) +class Qwen2_5OmniThinkerForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsPP, + Qwen2_5OmniConditionalGenerationMixin): + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "thinker.lm_head.": "language_model.lm_head.", + "thinker.model.": "language_model.model.", + "thinker.": "", + }) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]: + if modality.startswith("image"): + return "<|vision_start|><|IMAGE|><|vision_end|>" + if modality.startswith("video"): + return "<|vision_start|><|VIDEO|><|vision_end|>" + if modality.startswith("audio"): + return f"Audio {i}: <|audio_bos|><|AUDIO|><|audio_eos|>" + + raise ValueError("Only image, video or audio modality is supported") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + thinker_config: Qwen2_5OmniThinkerConfig = ( + vllm_config.model_config.hf_config) + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + self.config = thinker_config + self.multimodal_config = multimodal_config + + # force "use_flash_attention_2=True" to audio tower to align + # the results. + if flash_attn is not None: + audio_config = thinker_config.audio_config + audio_config._attn_implementation_autoset = True + audio_config._attn_implementation = "flash_attention_2" + else: + logger.warning( + "flash_attn is not available, the model may not yield the " + "exactly same result as the transformers implementation " + "in the audio tower part.") + + self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config) + self.visual = Qwen2_5_VisionTransformer( + vision_config=thinker_config.vision_config, + norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + prefix=maybe_prefix(prefix, "visual"), + ) + self.quant_config = quant_config + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + hf_config=thinker_config.text_config, + architectures=["Qwen2ForCausalLM"], + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors) + + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: + mm_input_by_modality = {} + + # Preserve the order of modalities if there are multiple of them + # from the order of kwargs. + for input_key in kwargs: + if input_key in ("pixel_values", "image_embeds" + ) and "image" not in mm_input_by_modality: + mm_input_by_modality[ + "image"] = self._parse_and_validate_image_input(**kwargs) + if input_key in ("pixel_values_videos", "video_embeds" + ) and "video" not in mm_input_by_modality: + mm_input_by_modality[ + "video"] = self._parse_and_validate_video_input(**kwargs) + if input_key in ("input_audio_features" + ) and "audio" not in mm_input_by_modality: + mm_input_by_modality[ + "audio"] = self._parse_and_validate_audio_input(**kwargs) + return mm_input_by_modality + + def get_language_model(self) -> torch.nn.Module: + return self.language_model + + def get_multimodal_embeddings(self, + **kwargs: object) -> MultiModalEmbeddings: + + mm_input_by_modality = self._parse_and_validate_multimodal_inputs( + **kwargs) + if not mm_input_by_modality: + return [] + + # The result multimodal_embeddings is tuple of tensors, with each + # tensor correspoending to a multimodal data item (image or video). + multimodal_embeddings: tuple[torch.Tensor, ...] = () + + # NOTE: It is important to iterate over the keys in this dictionary + # to preserve the order of the modalities. + for modality in mm_input_by_modality: + multimodal_input = mm_input_by_modality[modality] + if modality == "image": + vision_embeddings = self._process_image_input(multimodal_input) + multimodal_embeddings += vision_embeddings + if modality == "video": + video_embeddings = self._process_video_input(multimodal_input) + multimodal_embeddings += video_embeddings + if modality == "audio": + audio_embeddings = self._process_audio_input(multimodal_input) + multimodal_embeddings += audio_embeddings + return multimodal_embeddings + + def get_input_embeddings( + self, + input_ids: torch.Tensor, + multimodal_embeddings: Optional[MultiModalEmbeddings] = None, + ) -> torch.Tensor: + inputs_embeds = self.language_model.get_input_embeddings(input_ids) + if multimodal_embeddings is not None \ + and len(multimodal_embeddings) != 0: + + # TODO (ywang96): support overlapping modalitiy embeddings so that + # `use_audio_in_video` will work on V1. + inputs_embeds = merge_multimodal_embeddings( + input_ids, inputs_embeds, multimodal_embeddings, [ + self.config.image_token_index, + self.config.video_token_index, + self.config.audio_token_index + ]) + return inputs_embeds + + def get_multimodal_embeddings_v0( + self, **kwargs: object) -> Optional[NestedTensors]: + audio_input = self._parse_and_validate_audio_input(**kwargs) + image_input = self._parse_and_validate_image_input(**kwargs) + video_input = self._parse_and_validate_video_input(**kwargs) + + if audio_input is None and image_input is None and video_input is None: + return None + + multimodal_embeddings: list[tuple[NestedTensors, str]] = [] + + if audio_input is not None: + audio_embeds = self._process_audio_input(audio_input) + multimodal_embeddings.append((audio_embeds, "audio")) + if image_input is not None: + image_embeds = self._process_image_input(image_input) + multimodal_embeddings.append((image_embeds, "image")) + if video_input is not None: + video_embeds = self._process_video_input(video_input) + multimodal_embeddings.append((video_embeds, "video")) + return multimodal_embeddings + + def get_input_embeddings_v0( + self, + input_ids: torch.Tensor, + multimodal_embeddings: Optional[NestedTensors] = None, + ) -> torch.Tensor: + inputs_embeds = self.language_model.get_input_embeddings(input_ids) + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + for embeddings, modality in multimodal_embeddings: + if modality == "audio": + placeholder_token_id = self.config.audio_token_index + if modality == "image": + placeholder_token_id = self.config.image_token_index + if modality == "video": + placeholder_token_id = self.config.video_token_index + inputs_embeds = merge_multimodal_embeddings( + input_ids, inputs_embeds, embeddings, placeholder_token_id) + return inputs_embeds + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + **kwargs: object, + ) -> Union[torch.Tensor, IntermediateTensors]: + if intermediate_tensors is not None: + inputs_embeds = None + + # NOTE: In v1, inputs_embeds is always generated at model runner, this + # condition is for v0 compatibility. + elif inputs_embeds is None: + multimodal_embeddings = self.get_multimodal_embeddings(**kwargs) + inputs_embeds = self.get_input_embeddings( + input_ids, multimodal_embeddings) + text_inputs_embeds = self.get_input_embeddings(input_ids, [(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings] if multimodal_embeddings is not None else None) + input_ids = None + else: + text_inputs_embeds = inputs_embeds + + hidden_states = self.language_model.model(input_ids, + positions, + intermediate_tensors, + inputs_embeds=inputs_embeds) + return text_inputs_embeds, hidden_states.unsqueeze(0) #(1, S, D) + + def compute_logits( + self, + hidden_states: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[torch.Tensor]: + return self.language_model.compute_logits(hidden_states, + sampling_metadata) + + def load_weights(self, weights: Iterable[tuple[str, + torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=["talker.", "token2wav."], + ) + loaded_weights = loader.load_weights(weights, + mapper=self.hf_to_vllm_mapper) + # Log load summary + try: + total_bytes = 0 + for name, param in self.named_parameters(): + if param is not None and param.data is not None: + total_bytes += param.data.numel() * param.data.element_size() + device = next(self.parameters()).device + logger.info( + "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s", + self.__class__.__name__, True, total_bytes / (1024**2), str(device)) + except Exception: + pass + return loaded_weights + diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py new file mode 100644 index 000000000..3937c2942 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py @@ -0,0 +1,1690 @@ +############################ +# Start Token2Wav # +############################ + +from typing import Iterable, Optional, Set, Tuple, List, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +import math +from torch.nn import Parameter + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.sequence import IntermediateTensors +from vllm.model_executor.sampling_metadata import SamplingMetadata +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler +from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader as _Vllm_AutoWeightsLoader, + WeightsMapper as _Vllm_WeightsMapper, + init_vllm_registered_model as _Vllm_init_vllm_registered_model, + maybe_prefix as _Vllm_maybe_prefix, +) + +# Bring in HF base classes, configs and utilities used below +from transformers.utils import ModelOutput +from transformers.utils.logging import get_logger as _hf_get_logger +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS +from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import ( + Qwen2_5OmniPreTrainedModel, + Qwen2_5OmniPreTrainedModelForConditionalGeneration, +) +from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( + Qwen2_5OmniDiTConfig, + Qwen2_5OmniBigVGANConfig, + Qwen2_5OmniToken2WavConfig, +) + +# Provide a no-op auto_docstring decorator to satisfy annotations if missing +def auto_docstring(func=None, **_kwargs): + if func is None: + def wrapper(f): + return f + return wrapper + return func + +# HF logger alias +logger = _hf_get_logger(__name__) + + +# Using custom RoPE, will use LlamaRotaryEmbedding next version +class Qwen2_5OmniDiTRotaryEmbedding(nn.Module): + def __init__(self, dim, base=10000): + super().__init__() + + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) + self.register_buffer("inv_freq", inv_freq) + + def forward(self, x): + batch_size, seq_len = x.shape[0], x.shape[1] + t = torch.arange(seq_len, device=x.device) + device_type = x.device.type + device_type = device_type if device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = t.unsqueeze(1).float() @ self.inv_freq.unsqueeze(0).float() + freqs = torch.stack((freqs, freqs), dim=-1) + freqs = freqs.reshape(*freqs.shape[:-2], -1) + freqs = freqs.repeat(batch_size, *([1] * freqs.dim())) + cos = freqs.cos() + sin = freqs.sin() + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class TimeDelayNetBlock(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + dilation, + ): + super().__init__() + self.conv = nn.Conv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + dilation=dilation, + padding="same", + padding_mode="reflect", + ) + self.activation = nn.ReLU() + + def forward(self, hidden_states: torch.Tensor): + return self.activation(self.conv(hidden_states)) + + +class Res2NetBlock(torch.nn.Module): + def __init__(self, in_channels, out_channels, scale=8, kernel_size=3, dilation=1): + super().__init__() + + in_channel = in_channels // scale + hidden_channel = out_channels // scale + + self.blocks = nn.ModuleList( + [ + TimeDelayNetBlock( + in_channel, + hidden_channel, + kernel_size=kernel_size, + dilation=dilation, + ) + for i in range(scale - 1) + ] + ) + self.scale = scale + + def forward(self, hidden_states): + outputs = [] + for i, hidden_part in enumerate(torch.chunk(hidden_states, self.scale, dim=1)): + if i == 0: + output_part = hidden_part + elif i == 1: + output_part = self.blocks[i - 1](hidden_part) + else: + output_part = self.blocks[i - 1](hidden_part + output_part) + outputs.append(output_part) + output = torch.cat(outputs, dim=1) + return output + + +class SqueezeExcitationBlock(nn.Module): + def __init__(self, in_channels, se_channels, out_channels): + super().__init__() + + self.conv1 = nn.Conv1d( + in_channels=in_channels, + out_channels=se_channels, + kernel_size=1, + padding="same", + padding_mode="reflect", + ) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv1d( + in_channels=se_channels, + out_channels=out_channels, + kernel_size=1, + padding="same", + padding_mode="reflect", + ) + self.sigmoid = nn.Sigmoid() + + def forward(self, hidden_states): + hidden_states_mean = hidden_states.mean(dim=2, keepdim=True) + + hidden_states_mean = self.relu(self.conv1(hidden_states_mean)) + hidden_states_mean = self.sigmoid(self.conv2(hidden_states_mean)) + + return hidden_states * hidden_states_mean + + +class AttentiveStatisticsPooling(nn.Module): + """This class implements an attentive statistic pooling layer for each channel. + It returns the concatenated mean and std of the input tensor. + """ + + def __init__(self, channels, attention_channels=128): + super().__init__() + + self.eps = 1e-12 + self.tdnn = TimeDelayNetBlock(channels * 3, attention_channels, 1, 1) + self.tanh = nn.Tanh() + self.conv = nn.Conv1d( + in_channels=attention_channels, + out_channels=channels, + kernel_size=1, + padding="same", + padding_mode="reflect", + ) + + def _length_to_mask(self, length, max_len=None, dtype=None, device=None): + """Creates a binary mask for each sequence. + + Reference: https://discuss.pytorch.org/t/how-to-generate-variable-length-mask/23397/3 + + Arguments + --------- + length : torch.LongTensor + Containing the length of each sequence in the batch. Must be 1D. + max_len : int + Max length for the mask, also the size of the second dimension. + dtype : torch.dtype, default: None + The dtype of the generated mask. + device: torch.device, default: None + The device to put the mask variable. + + Returns + ------- + mask : tensor + The binary mask. + """ + + if max_len is None: + max_len = length.max().long().item() # using arange to generate mask + mask = torch.arange(max_len, device=length.device, dtype=length.dtype).expand( + len(length), max_len + ) < length.unsqueeze(1) + + mask = torch.as_tensor(mask, dtype=dtype, device=device) + return mask + + def _compute_statistics(self, x, m, dim=2): + mean = (m * x).sum(dim) + std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps)) + return mean, std + + def forward(self, hidden_states): + seq_length = hidden_states.shape[-1] + lengths = torch.ones(hidden_states.shape[0], device=hidden_states.device) + + # Make binary mask of shape [N, 1, L] + mask = self._length_to_mask( + lengths * seq_length, max_len=seq_length, dtype=hidden_states.dtype, device=hidden_states.device + ) + mask = mask.unsqueeze(1) + + # Expand the temporal context of the pooling layer by allowing the + # self-attention to look at global properties of the utterance. + total = mask.sum(dim=2, keepdim=True) + + mean, std = self._compute_statistics(hidden_states, mask / total) + mean = mean.unsqueeze(2).repeat(1, 1, seq_length) + std = std.unsqueeze(2).repeat(1, 1, seq_length) + attention = torch.cat([hidden_states, mean, std], dim=1) + + # Apply layers + attention = self.conv(self.tanh(self.tdnn(attention))) + + # Filter out zero-paddings + attention = attention.masked_fill(mask == 0, float("-inf")) + + attention = F.softmax(attention, dim=2) + mean, std = self._compute_statistics(hidden_states, attention) + # Append mean and std of the batch + pooled_stats = torch.cat((mean, std), dim=1) + pooled_stats = pooled_stats.unsqueeze(2) + + return pooled_stats + + +class SqueezeExcitationRes2NetBlock(nn.Module): + """An implementation of building block in ECAPA-TDNN, i.e., + TDNN-Res2Net-TDNN-SqueezeExcitationBlock. + """ + + def __init__( + self, + in_channels, + out_channels, + res2net_scale=8, + se_channels=128, + kernel_size=1, + dilation=1, + ): + super().__init__() + self.out_channels = out_channels + self.tdnn1 = TimeDelayNetBlock( + in_channels, + out_channels, + kernel_size=1, + dilation=1, + ) + self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation) + self.tdnn2 = TimeDelayNetBlock( + out_channels, + out_channels, + kernel_size=1, + dilation=1, + ) + self.se_block = SqueezeExcitationBlock(out_channels, se_channels, out_channels) + + def forward(self, hidden_state): + residual = hidden_state + + hidden_state = self.tdnn1(hidden_state) + hidden_state = self.res2net_block(hidden_state) + hidden_state = self.tdnn2(hidden_state) + hidden_state = self.se_block(hidden_state) + + return hidden_state + residual + + +class ECAPA_TimeDelayNet(torch.nn.Module): + """An implementation of the speaker embedding model in a paper. + "ECAPA-TDNN: Emphasized Channel Attention, Propagation and Aggregation in + TDNN Based Speaker Verification" (https://huggingface.co/papers/2005.07143). + """ + + def __init__(self, config: Qwen2_5OmniDiTConfig): + super().__init__() + if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(config.enc_channels) != len( + config.enc_dilations + ): + raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations should have same length") + self.channels = config.enc_channels + self.blocks = nn.ModuleList() + + # The initial TDNN layer + self.blocks.append( + TimeDelayNetBlock( + config.mel_dim, + config.enc_channels[0], + config.enc_kernel_sizes[0], + config.enc_dilations[0], + ) + ) + + # SE-Res2Net layers + for i in range(1, len(config.enc_channels) - 1): + self.blocks.append( + SqueezeExcitationRes2NetBlock( + config.enc_channels[i - 1], + config.enc_channels[i], + res2net_scale=config.enc_res2net_scale, + se_channels=config.enc_se_channels, + kernel_size=config.enc_kernel_sizes[i], + dilation=config.enc_dilations[i], + ) + ) + + # Multi-layer feature aggregation + self.mfa = TimeDelayNetBlock( + config.enc_channels[-1], + config.enc_channels[-1], + config.enc_kernel_sizes[-1], + config.enc_dilations[-1], + ) + + # Attentive Statistical Pooling + self.asp = AttentiveStatisticsPooling( + config.enc_channels[-1], + attention_channels=config.enc_attention_channels, + ) + + # Final linear transformation + self.fc = nn.Conv1d( + in_channels=config.enc_channels[-1] * 2, + out_channels=config.enc_dim, + kernel_size=1, + padding="same", + padding_mode="reflect", + ) + + def forward(self, hidden_states): + # Minimize transpose for efficiency + hidden_states = hidden_states.transpose(1, 2) + + hidden_states_list = [] + for layer in self.blocks: + hidden_states = layer(hidden_states) + hidden_states_list.append(hidden_states) + + # Multi-layer feature aggregation + hidden_states = torch.cat(hidden_states_list[1:], dim=1) + hidden_states = self.mfa(hidden_states) + + # Attentive Statistical Pooling + hidden_states = self.asp(hidden_states) + + # Final linear transformation + hidden_states = self.fc(hidden_states) + + hidden_states = hidden_states.squeeze(-1) + return hidden_states + + +class DiTInputEmbedding(nn.Module): + def __init__(self, config: Qwen2_5OmniDiTConfig): + super().__init__() + self.proj = nn.Linear( + config.mel_dim + config.enc_dim + config.enc_emb_dim + config.emb_dim, + config.hidden_size, + ) + self.spk_encoder = ECAPA_TimeDelayNet(config) + + def forward( + self, + hidden_states: torch.Tensor, + speaker_embedding: torch.Tensor, + condition_vector: torch.Tensor, + code_embed: torch.Tensor, + drop_audio_cond: Optional[bool] = False, + code_embed_uncond: Optional[bool] = None, + apply_cfg: Optional[bool] = True, + ): + if apply_cfg: + hidden_states = torch.cat([hidden_states, hidden_states], dim=0) + speaker_embedding = torch.cat([speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0) + condition_vector = torch.cat([condition_vector, torch.zeros_like(condition_vector)], dim=0) + code_embed = torch.cat([code_embed, code_embed_uncond], dim=0) + elif drop_audio_cond: # cfg for cond audio + condition_vector = torch.zeros_like(condition_vector) + speaker_embedding = torch.zeros_like(speaker_embedding) + condition_vector = self.spk_encoder(condition_vector).unsqueeze(1).repeat(1, hidden_states.size(1), 1) + hidden_states = self.proj(torch.cat((hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1)) + + return hidden_states + + +# Transformer backbone using DiT blocks +class DiTCodecEmbedding(nn.Module): + def __init__(self, codec_num_embeds, codec_dim, repeats): + super().__init__() + self.repeats = repeats + self.codec_embed = nn.Embedding(codec_num_embeds + 1, codec_dim) + + def forward(self, code, drop_code=False): + if drop_code: + code = torch.zeros_like(code) + code_embed = self.codec_embed(code) + + code_embed = torch.repeat_interleave(code_embed, repeats=self.repeats, dim=1) + return code_embed + + +# AdaLayerNormZero +# return with modulated x for attn input, and params for later mlp modulation +class Qwen2_5_OmniAdaLayerNormZero(nn.Module): + def __init__(self, dim): + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 6) + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, hidden_states, emb=None): + emb = self.linear(self.silu(emb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1) + + hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +# AdaLayerNormZero for final layer +# return only with modulated x for attn input, cuz no more mlp modulation +class Qwen2_5_OmniAdaLayerNormZero_Final(nn.Module): + def __init__(self, dim): + super().__init__() + + self.silu = nn.SiLU() + self.linear = nn.Linear(dim, dim * 2) + + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) + + def forward(self, hidden_states, emb): + emb = self.linear(self.silu(emb)) + scale, shift = torch.chunk(emb, 2, dim=1) + + hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :] + return hidden_states + + +# FeedForward +class DiTMLP(nn.Module): + def __init__(self, dim, mult=4, dropout=0.0): + super().__init__() + inner_dim = int(dim * mult) + + self.ff = nn.ModuleList( + [ + nn.Linear(dim, inner_dim), + nn.GELU(approximate="tanh"), + nn.Dropout(dropout), + nn.Linear(inner_dim, dim), + ] + ) + + def forward(self, hidden_states): + for layer in self.ff: + hidden_states = layer(hidden_states) + return hidden_states + + +# Modified from Llama with a different rotate function, will fixed in next release +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + + def rotate_half_codec(x): + # x = rearrange(x, "... (d r) -> ... d r", r=2) + x = x.reshape(*x.shape[:-1], -1, 2) + x1, x2 = x.unbind(dim=-1) + x = torch.stack((-x2, x1), dim=-1) + return x.reshape(*x.shape[:-2], -1) + + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half_codec(q) * sin) + k_embed = (k * cos) + (rotate_half_codec(k) * sin) + return q_embed, k_embed + + +class DiTAttention(nn.Module): + def __init__(self, config: Qwen2_5OmniDiTConfig): + super().__init__() + + self.config = config + self.dim = config.hidden_size + self.heads = config.num_attention_heads + self.inner_dim = config.head_dim * config.num_attention_heads + self.dropout = config.dropout + self.is_causal = False + + self.to_q = nn.Linear(config.hidden_size, self.inner_dim) + self.to_k = nn.Linear(config.hidden_size, self.inner_dim) + self.to_v = nn.Linear(config.hidden_size, self.inner_dim) + + self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)]) + + def forward( + self, + hidden_states, # noised input x + position_embeddings=None, # rotary position embedding for x + attention_mask=None, + ) -> torch.Tensor: + batch_size = hidden_states.shape[0] + + # `sample` projections. + query = self.to_q(hidden_states) + key = self.to_k(hidden_states) + value = self.to_v(hidden_states) + + # attention + inner_dim = key.shape[-1] + head_dim = inner_dim // self.heads + query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + # apply rotary position embedding + # Due to training process, only first head is applied with RoPE, will be fixed at next release + cos, sin = position_embeddings + query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin) + + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attention_weights, _ = attention_interface( + self, + query, + key, + value, + attention_mask=attention_mask, + is_causal=False, + ) + + # mask. e.g. inference got a batch with different target durations, mask out the padding + attention_weights = attention_weights.reshape(batch_size, -1, self.heads * head_dim) + attention_weights = attention_weights.to(query.dtype) + + # linear proj + attention_output = self.to_out[0](attention_weights) + attention_output = self.to_out[1](attention_output) + + return attention_output + + +# time step conditioning embedding +class SinusPositionEmbedding(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, hidden_states, scale=1000): + device = hidden_states.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * hidden_states.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb.type_as(hidden_states) + + +class DiTTimestepEmbedding(nn.Module): + def __init__(self, dim, freq_embed_dim=256): + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)]) + + def forward(self, timestep): # noqa: F821 + time_hidden = self.time_embed(timestep) + time_hidden = time_hidden.to(timestep.dtype) + for layer in self.time_mlp: + time_hidden = layer(time_hidden) # b d + return time_hidden + + +class DiTDecoderLayer(nn.Module): + def __init__(self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0): + super().__init__() + self.attn_norm = Qwen2_5_OmniAdaLayerNormZero(config.hidden_size) + + self.attn = DiTAttention(config) + self.look_ahead_block = look_ahead_block + self.look_backward_block = look_backward_block + self.ff_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6) + self.ff = DiTMLP(dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout) + + def forward( + self, hidden_states, timestep, position_embeddings=None, block_diff=None + ): # x: noised input, t: time embedding + # pre-norm & modulation for attention input + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, emb=timestep) + + # attention + attn_output = self.attn( + hidden_states=norm, + position_embeddings=position_embeddings, + attention_mask=(block_diff >= -float(self.look_backward_block)) + & (block_diff <= float(self.look_ahead_block)), + ) + + # process attention output for input x + hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output + + norm = self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ff_output = self.ff(norm) + hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output + + return hidden_states + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude of the periodic components + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + References: + - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://huggingface.co/papers/2006.08195 + """ + + def __init__(self, in_features, alpha=1.0): + super().__init__() + self.in_features = in_features + + # initialize alpha + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + + self.no_div_by_zero = 0.000000001 + + def forward(self, hidden_states): + """ + Forward pass of the function. + Applies the function to the input elementwise. + SnakeBeta ∶= x + 1/b * sin^2 (xa) + """ + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T] + beta = self.beta.unsqueeze(0).unsqueeze(-1) + alpha = torch.exp(alpha) + beta = torch.exp(beta) + hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow( + torch.sin(hidden_states * alpha), 2 + ) + + return hidden_states + + +def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): + """Generates a 1D Kaiser-windowed sinc filter. + + Args: + cutoff (float): Normalized cutoff frequency (0 to 0.5). + half_width (float): Transition bandwidth. + kernel_size (int): Number of filter taps. + + Returns: + torch.Tensor: A tensor of shape (1, 1, kernel_size) representing the filter. + """ + is_even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + + # Compute Kaiser window parameters + delta_f = 4 * half_width + attenuation = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + + if attenuation > 50.0: + beta = 0.1102 * (attenuation - 8.7) + elif attenuation >= 21.0: + beta = 0.5842 * (attenuation - 21) ** 0.4 + 0.07886 * (attenuation - 21.0) + else: + beta = 0.0 + + kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32) + + # Compute time indices + if is_even: + time_indices = torch.arange(-half_size, half_size) + 0.5 + else: + time_indices = torch.arange(kernel_size) - half_size + + # Compute sinc filter + if cutoff == 0: + return torch.zeros((1, 1, kernel_size), dtype=torch.float32) # Ensures correct shape + + sinc_filter = torch.sinc(2 * cutoff * time_indices) + normalized_filter = 2 * cutoff * kaiser_window * sinc_filter + + # Normalize to ensure sum = 1 (avoid leakage of constant component) + normalized_filter /= normalized_filter.sum() + + return normalized_filter.view(1, 1, kernel_size) + + +class UpSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + self.ratio = ratio + self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + self.stride = ratio + self.pad = self.kernel_size // ratio - 1 + self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2 + self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + + filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size) + self.register_buffer("filter", filter, persistent=False) + + def forward(self, hidden_states): + channels = hidden_states.shape[1] + hidden_states_dtype = hidden_states.dtype + hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype) + hidden_states = self.ratio * F.conv_transpose1d( + hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels + ).to(hidden_states_dtype) + hidden_states = hidden_states[..., self.pad_left : -self.pad_right] + + return hidden_states + + +class DownSample1d(nn.Module): + def __init__(self, ratio=2, kernel_size=None): + super().__init__() + cutoff = 0.5 / ratio + half_width = 0.6 / ratio + + if cutoff < 0.0: + raise ValueError("Minimum cutoff must be larger than zero.") + if cutoff > 0.5: + raise ValueError("A cutoff above 0.5 does not make sense.") + + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = ratio + filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + self.register_buffer("filter", filter, persistent=False) + + def forward(self, hidden_states): + channels = hidden_states.shape[1] + hidden_states_dtype = hidden_states.dtype + hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(self.filter.dtype) + out = F.conv1d(hidden_states, self.filter.expand(channels, -1, -1), stride=self.stride, groups=channels).to(hidden_states_dtype) + return out + + +class TorchActivation1d(nn.Module): + def __init__( + self, + activation, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + ): + super().__init__() + if not callable(activation): + raise TypeError("Activation function must be callable") + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + def forward(self, hidden_states): + hidden_states = self.upsample(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states = self.downsample(hidden_states) + + return hidden_states + + +class AMPBlock(torch.nn.Module): + def __init__( + self, + channels, + kernel_size=3, + dilation=(1, 3, 5), + ): + super().__init__() + + self.convs1 = nn.ModuleList( + [ + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=self._get_padding(kernel_size, dilation[0]), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=self._get_padding(kernel_size, dilation[1]), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + padding=self._get_padding(kernel_size, dilation[2]), + ), + ] + ) + + self.convs2 = nn.ModuleList( + [ + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=self._get_padding(kernel_size, 1), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=self._get_padding(kernel_size, 1), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=self._get_padding(kernel_size, 1), + ), + ] + ) + + self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers + + self.activations = nn.ModuleList( + [TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)] + ) + + def _get_padding(self, kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + def forward(self, hidden_states): + acts1, acts2 = self.activations[::2], self.activations[1::2] + for conv1, conv2, act1, act2 in zip(self.convs1, self.convs2, acts1, acts2): + residual = hidden_states + hidden_states = act1(hidden_states) + hidden_states = conv1(hidden_states) + hidden_states = act2(hidden_states) + hidden_states = conv2(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +@auto_docstring( + custom_intro=""" + The full Qwen2.5Omni Token2WavBigVGAN model. Which take mel spectrogram as input and predict waveform. + """ +) +class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel): + config: Qwen2_5OmniBigVGANConfig + + def __init__(self, config: Qwen2_5OmniBigVGANConfig): + super().__init__(config) + self.num_residual_blocks = len(config.resblock_kernel_sizes) + self.num_upsample_layers = len(config.upsample_rates) + + self.conv_pre = nn.Conv1d(config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3) + + # Removing extra ModuleList breaks official state dict + ups = [ + nn.ModuleList( + [ + nn.ConvTranspose1d( + config.upsample_initial_channel // (2**layer_idx), + config.upsample_initial_channel // (2 ** (layer_idx + 1)), + kernel_size, + stride, + padding=(kernel_size - stride) // 2, + ) + ] + ) + for layer_idx, (stride, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes)) + ] + self.ups = nn.ModuleList(ups) + + self.resblocks = nn.ModuleList( + [ + AMPBlock(config.upsample_initial_channel // (2 ** (layer_idx + 1)), kernel_size, dilation) + for layer_idx in range(self.num_upsample_layers) + for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes) + ] + ) + + self.activation_post = TorchActivation1d( + activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers)) + ) + self.conv_post = nn.Conv1d( + config.upsample_initial_channel // (2**self.num_upsample_layers), 1, 7, 1, padding=3, bias=False + ) + + def normalize_spectrogram(self, spectrogram, max_value, min_db): + return torch.clamp((2 * max_value) * ((spectrogram - min_db) / (-min_db)) - max_value, -max_value, max_value) + + def amplitude_to_db(self, amplitude, min_db_level): + min_level = torch.exp( + torch.tensor(min_db_level / 20.0 * np.log(10), device=amplitude.device, dtype=amplitude.dtype) + ) + return 20 * torch.log10(torch.clamp(amplitude, min=min_level)) + + def process_mel_spectrogram(self, mel_spectrogram): + amplitude_spectrum = torch.exp(mel_spectrogram) + decibel_spectrum = self.amplitude_to_db(amplitude_spectrum, -115) - 20 + return self.normalize_spectrogram(decibel_spectrum, 1, -115) + + def forward(self, mel_spectrogram): + processed_spectrogram = self.process_mel_spectrogram(mel_spectrogram) + hidden_representation = self.conv_pre(processed_spectrogram) + + for layer_index in range(self.num_upsample_layers): + hidden_representation = self.ups[layer_index][0](hidden_representation) + residual_output = sum( + self.resblocks[layer_index * self.num_residual_blocks + block_index](hidden_representation) + for block_index in range(self.num_residual_blocks) + ) + residual_output = residual_output / self.num_residual_blocks + hidden_representation = residual_output + + hidden_representation = self.activation_post(hidden_representation) + output_waveform = self.conv_post(hidden_representation) + return torch.clamp(output_waveform, min=-1.0, max=1.0).squeeze().cpu() + + +class RungeKutta4ODESolver: + def __init__(self, function, initial_value): + self.function = function + self.initial_value = initial_value + + self._one_third = 1 / 3 + self._two_thirds = 2 / 3 + + def _rk4_step(self, function, time_start, time_step, time_end, value_start, function_value_start=None): + k1 = function_value_start if function_value_start is not None else function(time_start, value_start) + k2 = function(time_start + time_step * self._one_third, value_start + time_step * k1 * self._one_third) + k3 = function(time_start + time_step * self._two_thirds, value_start + time_step * (k2 - k1 * self._one_third)) + k4 = function(time_end, value_start + time_step * (k1 - k2 + k3)) + return (k1 + 3 * (k2 + k3) + k4) * time_step / 8 + + def _compute_step(self, function, time_start, time_step, time_end, value_start): + function_value_start = function(time_start, value_start) + return self._rk4_step( + function, time_start, time_step, time_end, value_start, function_value_start=function_value_start + ), function_value_start + + def _linear_interpolation(self, time_start, time_end, value_start, value_end, time_point): + if time_point == time_start: + return value_start + if time_point == time_end: + return value_end + weight = (time_point - time_start) / (time_end - time_start) + return value_start + weight * (value_end - value_start) + + def integrate(self, time_points): + solution = torch.empty( + len(time_points), + *self.initial_value.shape, + dtype=self.initial_value.dtype, + device=self.initial_value.device, + ) + solution[0] = self.initial_value + + current_index = 1 + current_value = self.initial_value + for time_start, time_end in zip(time_points[:-1], time_points[1:]): + time_step = time_end - time_start + delta_value, _ = self._compute_step(self.function, time_start, time_step, time_end, current_value) + next_value = current_value + delta_value + + while current_index < len(time_points) and time_end >= time_points[current_index]: + solution[current_index] = self._linear_interpolation( + time_start, time_end, current_value, next_value, time_points[current_index] + ) + current_index += 1 + + current_value = next_value + + return solution + + +@auto_docstring( + custom_intro=""" + The full Qwen2.5Omni Token2WavDiT model. Which take speech tokens as input and predict mel spectrogram. + """ +) +class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel): + config: Qwen2_5OmniDiTConfig + _no_split_modules = ["DiTDecoderLayer"] + + def __init__(self, config: Qwen2_5OmniDiTConfig): + super().__init__(config) + self.mel_dim = config.mel_dim + self.repeats = config.repeats + self.time_embed = DiTTimestepEmbedding(config.hidden_size) + + self.text_embed = DiTCodecEmbedding(config.num_embeds, config.emb_dim, config.repeats) + self.input_embed = DiTInputEmbedding(config) + + self.rotary_embed = Qwen2_5OmniDiTRotaryEmbedding(config.head_dim) + + self.hidden_size = config.hidden_size + self.layers = config.num_hidden_layers + self.block_size = config.block_size + self.num_attention_heads = config.num_attention_heads + + self.transformer_blocks = nn.ModuleList() + for i in range(config.num_hidden_layers): + self.transformer_blocks.append( + DiTDecoderLayer( + config, + look_ahead_block=1 if i in config.look_ahead_layers else 0, + look_backward_block=1 if i in config.look_backward_layers else 0, + ) + ) + + self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size) # final modulation + self.proj_out = nn.Linear(config.hidden_size, config.mel_dim) + + def _create_block_diff(self, hidden_states): + batch, seq_len = hidden_states.shape[0], hidden_states.shape[1] + block_indices = torch.arange(seq_len, device=hidden_states.device) // self.block_size # [seq_length] + + block_i = block_indices.unsqueeze(1) # [seq_length, 1] + block_j = block_indices.unsqueeze(0) # [1, seq_length] + block_diff = block_j - block_i # (n, n) + + return block_diff.expand(batch, self.num_attention_heads, seq_len, seq_len) + + def forward( + self, + hidden_states, + condition_vector, + speaker_embedding, + quantized_code, + time_step, + drop_audio_conditioning=False, + drop_code=False, + apply_cfg=True, + ): + batch_size = hidden_states.shape[0] + if time_step.ndim == 0: + time_step = time_step.repeat(batch_size) + + # Compute embeddings + time_embedding = self.time_embed(time_step) + text_embedding = self.text_embed(quantized_code, drop_code=False if apply_cfg else drop_code) + text_embedding_unconditioned = self.text_embed(quantized_code, drop_code=True) if apply_cfg else None + + hidden_states = self.input_embed( + hidden_states, + speaker_embedding, + condition_vector, + text_embedding, + drop_audio_cond=drop_audio_conditioning, + code_embed_uncond=text_embedding_unconditioned, + apply_cfg=apply_cfg, + ) + + # Compute positional encodings + position_embeddings = self.rotary_embed(hidden_states) + blockwise_difference = self._create_block_diff(hidden_states) + + # Transformer blocks + for transformer_block in self.transformer_blocks: + hidden_states = transformer_block( + hidden_states, + time_embedding, + position_embeddings=position_embeddings, + block_diff=blockwise_difference, + ) + + hidden_states = self.norm_out(hidden_states, time_embedding) + output = self.proj_out(hidden_states) + + return output + + @torch.no_grad() + def sample( + self, + conditioning_vector, + reference_mel_spectrogram, + quantized_code, + num_steps=10, + guidance_scale=0.5, + sway_coefficient=-1.0, + ): + noise_initialization = torch.randn([1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype) + maximum_duration = quantized_code.shape[1] * self.repeats + initial_state = noise_initialization[:, :maximum_duration].to(quantized_code.device) + batch_size = reference_mel_spectrogram.shape[0] + conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, maximum_duration, 1) + + if batch_size != 1: + raise ValueError("Only batch size = 1 is currently supported") + + def ode_function(time_step, hidden_states): + if guidance_scale < 1e-5: + prediction = self( + hidden_states=hidden_states, + speaker_embedding=conditioning_vector, + condition_vector=reference_mel_spectrogram, + quantized_code=quantized_code, + time_step=time_step, + drop_audio_conditioning=False, + drop_code=False, + ) + return prediction + + model_output = self( + hidden_states=hidden_states, + quantized_code=quantized_code, + speaker_embedding=conditioning_vector, + condition_vector=reference_mel_spectrogram, + time_step=time_step, + apply_cfg=True, + ) + guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0) + return guided_prediction + (guided_prediction - null_prediction) * guidance_scale + + initial_time = 0 + time_embedding = torch.linspace( + initial_time, 1, num_steps, device=quantized_code.device, dtype=conditioning_vector.dtype + ) + + if sway_coefficient is not None: + time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding) + + ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state) + solution_trajectory = ode_solver.integrate(time_embedding) + + generated_waveform = solution_trajectory[-1] + generated_mel_spectrogram = generated_waveform.permute(0, 2, 1) + return generated_mel_spectrogram + + @torch.no_grad() + def fast_block_sample( + self, + conditioning_vector, + reference_mel_spectrogram, + quantized_code, + y0: torch.Tensor, + num_steps=10, + guidance_scale=0.5, + sway_coefficient: Optional[float] = -1.0, + ): + """ + Block-wise ODE sampling starting from provided initial state y0. + + Args: + conditioning_vector: (B, enc_emb_dim) + reference_mel_spectrogram: (B, T_ref, mel_dim) + quantized_code: (B, T_code) + y0: (B, T_target, mel_dim) initial state for ODE + Returns: + mel: (B, mel_dim, T_target) + """ + initial_state = y0.to(quantized_code.device) + batch_size = reference_mel_spectrogram.shape[0] + conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, initial_state.shape[1], 1) + + if batch_size != 1: + raise ValueError("Only batch size = 1 is currently supported") + + def ode_function(time_step, hidden_states): + if guidance_scale < 1e-5: + prediction = self( + hidden_states=hidden_states, + speaker_embedding=conditioning_vector, + condition_vector=reference_mel_spectrogram, + quantized_code=quantized_code, + time_step=time_step, + drop_audio_conditioning=False, + drop_code=False, + ) + return prediction + + model_output = self( + hidden_states=hidden_states, + quantized_code=quantized_code, + speaker_embedding=conditioning_vector, + condition_vector=reference_mel_spectrogram, + time_step=time_step, + apply_cfg=True, + ) + guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0) + return guided_prediction + (guided_prediction - null_prediction) * guidance_scale + + initial_time = 0 + time_embedding = torch.linspace( + initial_time, 1, num_steps, device=quantized_code.device, dtype=conditioning_vector.dtype + ) + + if sway_coefficient is not None: + time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding) + + ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state) + solution_trajectory = ode_solver.integrate(time_embedding) + + generated_waveform = solution_trajectory[-1] + generated_mel_spectrogram = generated_waveform.permute(0, 2, 1) + return generated_mel_spectrogram + + +@auto_docstring( + custom_intro=""" + The full Qwen2.5Omni Token2Wav model. Consists a DiT model take speech tokens as input and predict mel spectrogram and a BigVGAN vocoder take mel spectrogram as input and predict waveform. + """ +) +class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel): + config: Qwen2_5OmniToken2WavConfig + base_model_prefix = "model" + _no_split_modules = ["Qwen2_5OmniToken2WavDiTModel", "Qwen2_5OmniToken2WavBigVGANModel"] + + def __init__(self, config: Qwen2_5OmniToken2WavConfig): + super().__init__(config) + attn_impl = config._attn_implementation + if config._attn_implementation == "flash_attention_2": + logger.warning_once( + "Qwen2_5OmniToken2WavModel must inference with fp32, but flash_attention_2 only supports fp16 and bf16, " + "attention implementation of Qwen2_5OmniToken2WavModel will fallback to sdpa." + ) + attn_impl = "sdpa" + elif config._attn_implementation == "eager": + logger.warning_once( + "Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa" + ) + attn_impl = "sdpa" + self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel._from_config( + config.dit_config, attn_implementation=attn_impl + ) + self.code2wav_bigvgan_model = Qwen2_5OmniToken2WavBigVGANModel._from_config( + config.bigvgan_config, attn_implementation=attn_impl + ) + + # Streaming-related parameters aligned with Qwen2Code2wav + self.factor = self.code2wav_dit_model.repeats # 50Hz=2, 200Hz=4 + # default bs_mel depends on factor + self.bs_mel = 24 if self.factor == 2 else 32 + self.bs_codec = self.bs_mel // self.factor + self.past_cache_size = self.bs_mel * self.factor + self.future_cache_size = self.bs_mel * 1 + self.batched_chunk = 3 + self.chunk_size = self.bs_mel * self.batched_chunk + self.future_size = 20 if self.factor == 2 else 13 + + # codec embedding size for masking EOS out-of-range + try: + self.codec_embed_size = self.code2wav_dit_model.text_embed.codec_embed.weight.size(0) + except Exception: + self.codec_embed_size = -1 + + # vocoder hop length inferred from upsample rates + try: + ups = self.code2wav_bigvgan_model.config.upsample_rates + hop = 1 + for r in ups: + hop *= int(r) + self.vocoder_hop = int(hop) + except Exception: + # fallback to commonly used value + self.vocoder_hop = 240 + + def forward( + self, + code, + conditioning, + reference_mel, + num_steps=10, + guidance_scale=0.5, + sway_coefficient=-1.0, + **kwargs, + ): + """Generates a waveform from input code and conditioning parameters.""" + + mel_spectrogram = self.code2wav_dit_model.sample( + conditioning, + reference_mel, + code, + num_steps=num_steps, + guidance_scale=guidance_scale, + sway_coefficient=sway_coefficient, + ).to(self.code2wav_bigvgan_model.dtype) + + waveform = self.code2wav_bigvgan_model(mel_spectrogram).to(self.dtype) + + return waveform + + # ============== Chunked processing helpers (compat with qwen2_code2wav_dit) ============== + @torch.inference_mode() + def process_chunk_dit_batch( + self, + conditioning: torch.Tensor, + reference_mel: torch.Tensor, + code: torch.Tensor, + y0: torch.Tensor, + steps: int, + ) -> torch.Tensor: + """Block-wise DiT: generate mel from initial state y0 for the given code slice.""" + # prevent codec out-of-range (eos) + if self.codec_embed_size > 0: + code = code.clone() + code[code >= self.codec_embed_size] = 0 + mel = self.code2wav_dit_model.fast_block_sample( + conditioning_vector=conditioning, + reference_mel_spectrogram=reference_mel, + quantized_code=code, + y0=y0, + num_steps=steps, + ) + return mel.to(self.code2wav_bigvgan_model.dtype) + + @torch.inference_mode() + def process_chunk_bigvgan_batch(self, mel_batch: torch.Tensor) -> torch.Tensor: + """Vocoder batch: mel -> waveform.""" + return self.code2wav_bigvgan_model(mel_batch) + + @torch.inference_mode() + def process_little_chunk( + self, + conditioning: torch.Tensor, + reference_mel: torch.Tensor, + codec_all: torch.Tensor, + y_all: torch.Tensor, + i: int, + steps: int, + prev_generated: torch.Tensor, + finished: bool = False, + ) -> Tuple[Optional[torch.Tensor], torch.Tensor]: + """Streaming per small chunk: returns (mel_or_None, audio_slice).""" + start_index = max(i * self.chunk_size - self.past_cache_size, 0) + end_index = min((i + 1) * self.chunk_size + self.future_cache_size, + codec_all.shape[1] * self.factor) + + y0 = y_all[:, start_index:end_index].reshape(1, -1, self.code2wav_dit_model.mel_dim).contiguous() + codec = codec_all[:, start_index // self.factor:end_index // self.factor].reshape(1, -1).contiguous() + + # generate mel for current window (B, mel_dim, T) + generated = self.process_chunk_dit_batch( + conditioning=conditioning, + reference_mel=reference_mel, + code=codec, + y0=y0, + steps=steps, + ) + + # splice and vocode with 50Hz-style rules + return self._process_chunk_for_50hz( + i=i, + start_index=start_index, + end_index=end_index, + finished=finished, + prev_generated=prev_generated, + generated=generated, + ) + + @torch.inference_mode() + def process_chunk( + self, + conditioning: torch.Tensor, + reference_mel: torch.Tensor, + codec_all: torch.Tensor, + y_all: torch.Tensor, + i: int, + steps: int, + prev_generated: Union[torch.Tensor, List[torch.Tensor]], + finished: bool = False, + ) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], torch.Tensor]: + """High-level chunk API aligning to qwen2_code2wav_dit signature.""" + if not isinstance(prev_generated, torch.Tensor): + prev_generated = prev_generated[0] if len(prev_generated) > 0 else None + _mel, audio = self.process_little_chunk( + conditioning=conditioning, + reference_mel=reference_mel, + codec_all=codec_all, + y_all=y_all, + i=i, + steps=steps, + prev_generated=prev_generated, + finished=finished, + ) + return _mel if _mel is not None else prev_generated, audio + + @torch.inference_mode() + def _process_chunk_for_50hz( + self, + i: int, + start_index: int, + end_index: int, + finished: bool, + prev_generated: Optional[torch.Tensor], + generated: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Align mel and audio boundaries for 50Hz-like streaming. + + Shapes: + - generated: (B, mel_dim, T_window) + - prev_generated: (B, mel_dim, T_prev) + Returns: + - mel_chunk: (B, mel_dim, T_chunk) + - audio_slice: (T_audio_chunk,) + """ + # Normalize dtype + generated = generated.to(torch.float32) + if i == 0: + mel = generated[:, :, :self.chunk_size] + elif finished: + mel_trim = generated[:, :, self.past_cache_size:] + mel = torch.cat([prev_generated[:, :, -self.future_size * 2:], mel_trim], dim=2) + else: + if start_index == 0: + mel_trim = generated[:, :, i * self.chunk_size:-self.future_cache_size] + else: + mel_trim = generated[:, :, self.past_cache_size:-self.future_cache_size] + mel = torch.cat([prev_generated[:, :, -self.future_size * 2:], mel_trim], dim=2) + + audio = self.code2wav_bigvgan_model(mel) + if i == 0: + audio_output = audio[:-self.future_size * self.vocoder_hop] + elif finished: + audio_output = audio[self.future_size * self.vocoder_hop:] + else: + audio_output = audio[self.future_size * self.vocoder_hop:-self.future_size * self.vocoder_hop] + return mel, audio_output + + +# ================= vLLM-style wrapper for Token2Wav ================= + +class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP): + logger = init_logger(__name__) + + # Map HF weights -> vLLM module names + hf_to_vllm_mapper = _Vllm_WeightsMapper( + orig_to_new_prefix={ + # HF root is 'model.' + "model.": "token2wav_model.", + }) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + # Expect hf_config to be Token2Wav config + self.config = vllm_config.model_config.hf_config + + # Initialize underlying HF Token2Wav model via registry + self.token2wav = _Vllm_init_vllm_registered_model( + vllm_config=vllm_config, + prefix=_Vllm_maybe_prefix(prefix, "token2wav_model"), + hf_config=self.config, + architectures=["Qwen2_5OmniToken2WavDiTModel"], + ) + + # Provide placeholder to align with vLLM runner expectations + def _empty_intermediate_tensors(): + return None + self.make_empty_intermediate_tensors = _empty_intermediate_tensors + + def get_language_model(self) -> torch.nn.Module: + return self.token2wav + + @property + def sampler(self): + # Token2Wav does not use sampler; return vLLM default for API parity + return get_sampler() + + def forward( + self, + code: torch.Tensor, + conditioning: torch.Tensor, + reference_mel: torch.Tensor, + num_steps: int = 10, + guidance_scale: float = 0.5, + sway_coefficient: float = -1.0, + intermediate_tensors: Optional[IntermediateTensors] = None, + **kwargs, + ) -> torch.Tensor: + # Delegate to HF token2wav model + return self.token2wav( + code=code, + conditioning=conditioning, + reference_mel=reference_mel, + num_steps=num_steps, + guidance_scale=guidance_scale, + sway_coefficient=sway_coefficient, + **kwargs, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[torch.Tensor]: + # Token2Wav outputs waveform; logits are not applicable + return hidden_states + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[SamplerOutput]: + return None + + def load_weights_without_buffers(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]: + loader = _Vllm_AutoWeightsLoader(self) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + # Log load summary + try: + total_bytes = 0 + for _, param in self.named_parameters(): + if param is not None and param.data is not None: + total_bytes += param.data.numel() * param.data.element_size() + device = next(self.parameters()).device + self.logger.info( + "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s", + self.__class__.__name__, True, total_bytes / (1024**2), str(device)) + except Exception: + pass + return loaded + + def find_all_registers(self): + """ + Find all registered buffers in a PyTorch model. + + Args: + Returns: + dict: Dictionary with buffer names as keys and their properties as values + """ + registers = {} + + # Get all named buffers + for name, buf in self.named_buffers(): + if name in self.state_dict(): + registers[name] = { + 'name': name, + 'buffer': buf + } + return registers + + + #remove buffers from the weights and reload them after loading weights + def remove_buffers_from_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], buffers: dict): + weights_to_load = [] + for key, value in weights: + if key in buffers: + buffers[key]['buffer'] = value + continue + weights_to_load.append((key, value)) + return weights_to_load + + def reload_buffers_to_model(self, buffers: dict): + """ + reload stored buffers from weights to model + """ + loaded_buffers = set() + for name, buf_val in self.named_buffers(): + if name in buffers: + buf_val.copy_(buffers[name]['buffer']) + loaded_buffers.add(name) + return loaded_buffers + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], spk_dict_path: str) -> Set[str]: + buffers = self.find_all_registers() + weights_to_load = self.remove_buffers_from_weights(weights, buffers) + loaded = self.load_weights_without_buffers(weights_to_load) + loaded_buffers = self.reload_buffers_to_model(buffers) + #merge loaded and loaded_buffers + loaded.update(loaded_buffers) + self.spk_dict = torch.load(spk_dict_path) + return loaded + + # ============== Optional chunked helpers for API parity ============== + @torch.inference_mode() + def process_chunk_dit_batch( + self, + conditioning: torch.Tensor, + reference_mel: torch.Tensor, + code: torch.Tensor, + y0: torch.Tensor, + steps: int, + ) -> torch.Tensor: + return self.token2wav( + code=code, + conditioning=conditioning, + reference_mel=reference_mel, + num_steps=steps, + ) + + @torch.inference_mode() + def process_chunk_bigvgan_batch(self, mel_batch: torch.Tensor) -> Optional[torch.Tensor]: + # BigVGAN is not part of this wrapper; return None for parity. + return None + + @torch.inference_mode() + def process_little_chunk( + self, + conditioning: torch.Tensor, + reference_mel: torch.Tensor, + codec_all: torch.Tensor, + y_all: torch.Tensor, + i: int, + steps: int, + prev_generated: torch.Tensor, + finished: bool = False, + ) -> Tuple[Optional[torch.Tensor], torch.Tensor]: + mel = self.token2wav( + code=codec_all, + conditioning=conditioning, + reference_mel=reference_mel, + num_steps=steps, + ) + return None, mel + + @torch.inference_mode() + def process_chunk( + self, + conditioning: torch.Tensor, + reference_mel: torch.Tensor, + codec_all: torch.Tensor, + y_all: torch.Tensor, + i: int, + steps: int, + prev_generated: Union[torch.Tensor, List[torch.Tensor]], + finished: bool = False, + ) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], torch.Tensor]: + _mel, out = self.process_little_chunk( + conditioning=conditioning, + reference_mel=reference_mel, + codec_all=codec_all, + y_all=y_all, + i=i, + steps=steps, + prev_generated=prev_generated if isinstance(prev_generated, torch.Tensor) else None, + finished=finished, + ) + return _mel if _mel is not None else prev_generated, out \ No newline at end of file diff --git a/vllm_omni/model_executor/models/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_old.py new file mode 100644 index 000000000..6528bca78 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen2_old.py @@ -0,0 +1,539 @@ +from typing import Iterable, Optional, Set, Tuple, Union + +import torch +from torch import nn +from transformers import Qwen2Config + +from vllm.attention import Attention, AttentionType +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import (MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.pooler import Pooler, PoolingType +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, VocabParallelEmbedding) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, maybe_remap_kv_scale_name) +from vllm.v1.pool.metadata import PoolingMetadata +from vllm.model_executor.sampling_metadata import SamplingMetadata +from vllm.sequence import IntermediateTensors, PoolerOutput + +from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP +from vllm.model_executor.models.utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, make_layers, + maybe_prefix) + +logger = init_logger(__name__) + + +class Qwen2MLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now.") + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Qwen2Attention(nn.Module): + + def __init__(self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position: int = 4096 * 32, + rope_theta: float = 10000, + head_dim: Optional[int] = None, + cache_config: Optional[CacheConfig] = None, + quant_config: Optional[QuantizationConfig] = None, + rope_scaling: Optional[Tuple] = None, + prefix: str = "", + attn_type: str = AttentionType.DECODER) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = head_dim or (hidden_size // self.total_num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=self.rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = Attention(self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.attn", + attn_type=attn_type) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v) + output, _ = self.o_proj(attn_output) + return output + + +class Qwen2DecoderLayer(nn.Module): + + def __init__( + self, + config: Qwen2Config, + cache_config: Optional[CacheConfig] = None, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + # Requires transformers > 4.32.0 + rope_theta = getattr(config, "rope_theta", 1000000) + rope_scaling = getattr(config, "rope_scaling", None) + + # By default, Qwen2 uses causal attention as it is a decoder-only model. + # You can override the HF config with `is_causal=False` to enable + # bidirectional attention, which is used in some embedding models + # (e.g. Alibaba-NLP/gte-Qwen2-7B-instruct) + if getattr(config, "is_causal", True): + attn_type = AttentionType.DECODER + else: + attn_type = AttentionType.ENCODER_ONLY + + self.self_attn = Qwen2Attention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + max_position=config.max_position_embeddings, + num_kv_heads=config.num_key_value_heads, + rope_theta=rope_theta, + cache_config=cache_config, + quant_config=quant_config, + rope_scaling=rope_scaling, + head_dim=getattr(config, "head_dim", None), + prefix=f"{prefix}.self_attn", + attn_type=attn_type, + ) + self.mlp = Qwen2MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.input_layernorm = RMSNorm(config.hidden_size, + eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, + eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm( + hidden_states, residual) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + + # Fully Connected + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + # positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl, + # otherwise (seq_len, ). + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + }) +class Qwen2Model(nn.Module): + + def __init__(self, + *, + vllm_config: VllmConfig, + prefix: str = "", + decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer): + super().__init__() + + config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + + # TODO (@robertgshaw2): see if this can be moved out + if (cache_config.sliding_window is not None + and hasattr(config, "max_window_layers")): + raise ValueError("Sliding window for some but all layers is not " + "supported. This model uses sliding window " + "but `max_window_layers` = {} is less than " + "`num_hidden_layers` = {}. Please open an issue " + "to discuss this feature.".format( + config.max_window_layers, + config.num_hidden_layers, + )) + + self.config = config + self.quant_config = quant_config + self.vocab_size = config.vocab_size + + if get_pp_group().is_first_rank or (config.tie_word_embeddings + and get_pp_group().is_last_rank): + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + getattr(config, "embedding_size", config.hidden_size), + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + # Use the provided decoder layer type or default to Qwen2DecoderLayer + decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: decoder_layer_type(config=config, + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix), + prefix=f"{prefix}.layers", + ) + + self.make_empty_intermediate_tensors = ( + make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size)) + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, IntermediateTensors]: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.get_input_embeddings(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + for layer in self.layers[self.start_layer:self.end_layer]: + hidden_states, residual = layer( + positions, + hidden_states, + residual, + ) + if not get_pp_group().is_last_rank: + return IntermediateTensors({ + "hidden_states": hidden_states, + "residual": residual + }) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + def load_weights(self, weights: Iterable[Tuple[str, + torch.Tensor]]) -> Set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: Set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if (self.quant_config is not None and + (scale_name := self.quant_config.get_cache_scale(name))): + # Loading kv cache quantization scales + param = params_dict[scale_name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else + loaded_weight[0]) + weight_loader(param, loaded_weight) + loaded_params.add(scale_name) + continue + for (param_name, weight_name, shard_id) in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + lora_config = vllm_config.lora_config + + self.config = config + self.lora_config = lora_config + + self.quant_config = quant_config + self.model = Qwen2Model(vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model")) + + if get_pp_group().is_last_rank: + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead(config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix( + prefix, "lm_head")) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config.vocab_size) + self.sampler = get_sampler() + + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.get_input_embeddings(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, IntermediateTensors]: + hidden_states = self.model(input_ids, positions, intermediate_tensors, + inputs_embeds) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[torch.Tensor]: + logits = self.logits_processor(self.lm_head, hidden_states, + sampling_metadata) + return logits + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[SamplerOutput]: + next_tokens = self.sampler(logits, sampling_metadata) + return next_tokens + + def load_weights(self, weights: Iterable[Tuple[str, + torch.Tensor]]) -> Set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=(["lm_head."] + if self.config.tie_word_embeddings else None), + ) + return loader.load_weights(weights) + + +class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP): + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""}) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + lora_config = vllm_config.lora_config + pooler_config = vllm_config.model_config.pooler_config + + self.config = config + self.lora_config = lora_config + + self.quant_config = quant_config + self.model = Qwen2Model(vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model")) + + # TODO: Replace this model class with as_embedding_model( + # Qwen2ForCausalLM) after changing the default pooling method + if pooler_config.pooling_type is None: + logger.warning( + "This embedding model will default to last-token pooling in " + "an upcoming version. To avoid breaking changes, you should " + "pass `--override-pooler-config '{\"pooling_type\": \"MEAN\"}'`" + " explicitly.") + + self._pooler = Pooler.from_config_with_defaults( + pooler_config, + pooling_type=PoolingType.MEAN, + normalize=True, + softmax=False) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + ) -> torch.Tensor: + return self.model(input_ids, positions, intermediate_tensors) + + def pooler( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> Optional[PoolerOutput]: + return self._pooler(hidden_states, pooling_metadata) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + weights = self.hf_to_vllm_mapper.apply(weights) + weights = ((name, data) for name, data in weights + if not name.startswith("lm_head.")) + self.model.load_weights(weights) \ No newline at end of file diff --git a/vllm_omni/model_executor/models/utils.py b/vllm_omni/model_executor/models/utils.py new file mode 100644 index 000000000..7e915d937 --- /dev/null +++ b/vllm_omni/model_executor/models/utils.py @@ -0,0 +1,8 @@ +from vllm.model_executor.models.utils import maybe_prefix + + +def add_prefix_to_loaded_weights(weights: set[str], prefix: str) -> set[str]: + """ + Add a prefix to the names of the loaded weights. + """ + return {maybe_prefix(prefix, name) for name in weights} \ No newline at end of file diff --git a/vllm_omni/model_executor/stage_configs/__init__.py b/vllm_omni/model_executor/stage_configs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml new file mode 100644 index 000000000..49a1222cc --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -0,0 +1,40 @@ +# stage config for running qwen2.5-omni with architecture of OmniLLM. +stage_args: + - stage_id: 0 + engine_args: + model_stage: thinker + model_arch: Qwen2_5OmniForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler + gpu_memory_utilization: 0.32 + enforce_eager: true # need to discuss + trust_remote_code: true + engine_output_type: latent # change the param name,such as pooling_output + final_output: true + final_output_type: text + - stage_id: 1 + engine_args: + model_stage: talker + model_arch: Qwen2_5OmniForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler + gpu_memory_utilization: 0.32 + enforce_eager: true + trust_remote_code: true + engine_output_type: latent + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker + - stage_id: 2 + engine_args: + model_stage: code2wav + model_arch: Qwen2_5OmniForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker + scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler + gpu_memory_utilization: 0.3 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + engine_output_type: audio + engine_input_source: [1] + final_output: true + final_output_type: audio \ No newline at end of file diff --git a/vllm_omni/model_executor/stage_input_processors/__init__.py b/vllm_omni/model_executor/stage_input_processors/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py new file mode 100644 index 000000000..474e2c344 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -0,0 +1,40 @@ +from typing import Union +import torch + +from vllm.inputs import TextPrompt + +from vllm_omni.inputs.data import OmniTokensPrompt + +def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None): + 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") + thinker_outputs = stage_list[source_stage_id].engine_outputs + talker_inputs = [] + multi_modal_data = {thinker_output.request_id: + prompt.get('multi_modal_data', None) for thinker_output, prompt in zip(thinker_outputs, prompt)} + + for i, thinker_output in enumerate(thinker_outputs): + output = thinker_output.outputs[0] + prompt_token_ids = thinker_output.prompt_token_ids + thinker_output_ids = output.token_ids + prompt_token_ids_len = len(prompt_token_ids) + thinker_hidden_states = output.multimodal_output["latent"].clone().detach().cuda() + talker_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0] * (len(prompt_token_ids) + 2), # add 2 for codec pad and start token + additional_information={ + "thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(torch.float32), + "prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(torch.float32), + "prompt_token_ids": prompt_token_ids, + "thinker_output_token_ids": thinker_output_ids, + }, + multi_modal_data=multi_modal_data[thinker_output.request_id] if multi_modal_data is not None else None, + mm_processor_kwargs=None, + ) + ) + return talker_inputs \ No newline at end of file diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py new file mode 100644 index 000000000..73e509bd0 --- /dev/null +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -0,0 +1,578 @@ +"""AR GPU Model Runner for vLLM-omni. + +Exposes per-request hidden representations via ModelRunnerOutput.pooler_output +and also outputs sampled tokens. +""" + +from __future__ import annotations + +from typing import Optional, Union, Any, List +import numpy as np + +import torch + +from vllm import envs +from vllm.v1.worker.gpu_model_runner import ( + EMPTY_MODEL_RUNNER_OUTPUT, + IntermediateTensors, + get_pp_group, + get_tp_group, + has_kv_transfer_group, + set_forward_context, +) +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.multimodal.inputs import MultiModalKwargs +from vllm.v1.attention.backends.utils import CommonAttentionMetadata +from vllm.v1.spec_decode.eagle import EagleProposer + +from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner + + +class GPUARModelRunner(OmniGPUModelRunner): + """Autoregressive GPU model runner that returns hidden states per request. + + This runner follows the same preparation and forward path as GPUModelRunner + (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs), + and additionally performs lightweight sampling so that sampled tokens are + available in outputs. Hidden representations are taken at the same indices + that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`). + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.input_ids = torch.zeros(self.max_num_tokens, + dtype=torch.int64, + device=self.device) + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: "SchedulerOutput", + intermediate_tensors: Optional[IntermediateTensors] = None, + ) -> Union[OmniModelRunnerOutput, IntermediateTensors]: + # Update internal state with the new schedule + self._update_states(scheduler_output) + + # Decode per-request prompt_embeds / additional_hidden_states payloads (if present) into CPU tensors + try: + new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) + if new_reqs: + for nr in new_reqs: + req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) + if req_id is None: + continue + # prompt_embeds + payload_pe = getattr(nr, "prompt_embeds", None) + if payload_pe is not None: + if isinstance(payload_pe, torch.Tensor): + pe_cpu = payload_pe.detach().to("cpu").contiguous() + elif isinstance(payload_pe, PromptEmbedsPayload): + dt = np.dtype(getattr(payload_pe, "dtype", "float32")) + arr = np.frombuffer(payload_pe.data, dtype=dt) + arr = arr.reshape(payload_pe.shape) + pe_cpu = torch.from_numpy(arr) + else: + pe_cpu = None + if pe_cpu is not None and req_id in self.requests: + setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) + # additional_information + payload_info = getattr(nr, "additional_information", None) + if payload_info is not None: + info_dict = {} + if isinstance(payload_info, dict): + # Already decoded + info_dict = payload_info + elif isinstance(payload_info, AdditionalInformationPayload): + for k, entry in payload_info.entries.items(): + if entry.tensor_data is not None: + dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) + arr = np.frombuffer(entry.tensor_data, dtype=dt) + arr = arr.reshape(entry.tensor_shape) + info_dict[k] = torch.from_numpy(arr) + else: + info_dict[k] = entry.list_data + if info_dict and req_id in self.requests: + setattr(self.requests[req_id], "additional_information_cpu", info_dict) + except Exception: + pass + + # If there's no work to do, either return empty output or kv-only path + if not scheduler_output.total_num_scheduled_tokens: + if not has_kv_transfer_group(): + return EMPTY_MODEL_RUNNER_OUTPUT + return self.kv_connector_no_forward(scheduler_output, + self.vllm_config) + + # Prepare decoder inputs and attention metadata + (attn_metadata, attention_cuda_graphs, logits_indices, + spec_decode_metadata, num_scheduled_tokens_np, + spec_decode_common_attn_metadata) = self._prepare_inputs( + scheduler_output) + + # Determine number of input tokens for this iteration + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + if (self.use_cuda_graph + and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]): + num_input_tokens = self.vllm_config.pad_for_cudagraph( + num_scheduled_tokens) + else: + tp_size = self.vllm_config.parallel_config.tensor_parallel_size + if (self.compilation_config.pass_config.enable_sequence_parallelism + and tp_size > 1): + from vllm.utils import round_up # lazy local import + num_input_tokens = round_up(num_scheduled_tokens, tp_size) + else: + num_input_tokens = num_scheduled_tokens + + # DP padding + num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens) + num_input_tokens += num_pad + + # Multimodal handling (encode and gather embeddings if needed) + if self.is_multimodal_model: + self._execute_mm_encoder(scheduler_output) + mm_embeds = self._gather_mm_embeddings(scheduler_output) + else: + mm_embeds = [] + + # Always assemble inputs_embeds on first PP rank; overlay per-request prompt_embeds and collect additional_hidden_states for prefill only + if get_pp_group().is_first_rank: + inputs_embeds_scheduled = self.model.get_input_embeddings( + input_ids=self.input_ids[:num_scheduled_tokens], + multimodal_embeddings=(mm_embeds or None) + if self.is_multimodal_model else None, + ) + + # Copy into persistent buffer to enable CUDA Graph capture + self.inputs_embeds[:num_scheduled_tokens].copy_( + inputs_embeds_scheduled) + + # Reset per-step additional information collector + if hasattr(self, "_forward_additional_information"): + self._forward_additional_information = None + + # Overlay custom prompt_embeds per request for the prompt portion; collect additional_information (tensor/list) for prefill portion only + for req_index, req_id in enumerate(self.input_batch.req_ids): + req_state = self.requests[req_id] + pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) + addi_cpu = getattr(req_state, "additional_information_cpu", None) + num_computed_tokens = int( + self.input_batch.num_computed_tokens_cpu[req_index]) + prompt_len = len(req_state.prompt_token_ids) + prompt_remaining = max(0, prompt_len - num_computed_tokens) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + overlay_len = min(sched_tokens, prompt_remaining) + if overlay_len <= 0: + continue + if pe_cpu is not None: + src = pe_cpu[num_computed_tokens: + num_computed_tokens + overlay_len].to( + dtype=self.dtype, + device=self.device, + non_blocking=True) + start_offset = int(self.query_start_loc_cpu[req_index]) + self.inputs_embeds[start_offset:start_offset + overlay_len] \ + .copy_(src) + # For additional_information: handle arbitrary keys + if addi_cpu is not None and isinstance(addi_cpu, dict): + # Lazy init collector dict + if not hasattr(self, "_forward_additional_information") or \ + self._forward_additional_information is None: + self._forward_additional_information = {} + # Process tensors (slice by scheduled prompt range) and lists (append per-request) + for k, v in addi_cpu.items(): + if isinstance(v, torch.Tensor): + # Slice along token dimension for prefill part + try: + seg = v[num_computed_tokens: + num_computed_tokens + overlay_len].to( + dtype=self.dtype, + device=self.device, + non_blocking=True) + except Exception: + # Fallback: move whole tensor if slicing fails + seg = v.to(dtype=self.dtype, + device=self.device, + non_blocking=True) + prev_val = self._forward_additional_information.get(k) + self._forward_additional_information[k] = ( + torch.cat([prev_val, seg], dim=0) + if isinstance(prev_val, torch.Tensor) else seg.clone()) + elif isinstance(v, list): + prev_val = self._forward_additional_information.get(k) + if prev_val is None: + self._forward_additional_information[k] = [v] + elif isinstance(prev_val, list): + self._forward_additional_information[k].append(v) + else: + # Mixed types: wrap existing into list + self._forward_additional_information[k] = [prev_val, v] + + input_ids = self.input_ids[:num_input_tokens] # preserved for APIs + inputs_embeds = self.inputs_embeds[:num_input_tokens] + model_mm_kwargs = (self._extract_mm_kwargs(scheduler_output) + if self.is_multimodal_model else {}) + else: + input_ids = self.input_ids[:num_input_tokens] + inputs_embeds = None + model_mm_kwargs = {} + + # Positions/mRoPE + if self.uses_mrope: + positions = self.mrope_positions[:, :num_input_tokens] + else: + positions = self.positions[:num_input_tokens] + + # Handle pipeline-parallel intermediate tensors + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + intermediate_tensors = self.sync_and_slice_intermediate_tensors( + num_input_tokens, intermediate_tensors, True) + + # Some attention backends only support CUDA Graphs in pure decode. + skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs + + # Forward pass + with set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_input_tokens, + num_tokens_across_dp=num_tokens_across_dp, + skip_cuda_graphs=skip_cuda_graphs, + ), self.maybe_get_kv_connector_output( + scheduler_output) as kv_connector_output: + + model_kwargs_extra = {} + # Only pass additional_information for the prefill part + if hasattr(self, "_forward_additional_information") and \ + self._forward_additional_information is not None and \ + isinstance(self._forward_additional_information, dict): + model_kwargs_extra["additional_information"] = self._forward_additional_information + model_output = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **MultiModalKwargs.as_kwargs( + model_mm_kwargs, + device=self.device, + ), + sampling_metadata=self.input_batch.sampling_metadata, + logits_index=logits_indices, + sampler=self.sampler, + **model_kwargs_extra, + ) + + if self.use_aux_hidden_state_outputs: + hidden_states, _aux_hidden_states = model_output + else: + hidden_states = model_output + _aux_hidden_states = None + + text_hidden_states, multimodal_outputs = ( + self.extract_multimodal_outputs(hidden_states)) + + # Mid PP stages return intermediate tensors unmodified + if not get_pp_group().is_last_rank: + assert isinstance(text_hidden_states, IntermediateTensors) + text_hidden_states.kv_connector_output = kv_connector_output + return text_hidden_states + + # Broadcast PP output for external_launcher (torchrun) + broadcast_pp_output = \ + self.parallel_config.distributed_executor_backend \ + == "external_launcher" and len(get_pp_group().ranks) > 0 + if not get_pp_group().is_last_rank: + assert isinstance(text_hidden_states, IntermediateTensors) + if not broadcast_pp_output: + text_hidden_states.kv_connector_output = kv_connector_output + return text_hidden_states + get_pp_group().send_tensor_dict(text_hidden_states.tensors, + all_gather_group=get_tp_group()) + logits = None + else: + if self.input_batch.pooling_params: + return self._pool(text_hidden_states, num_scheduled_tokens, + num_scheduled_tokens_np, kv_connector_output) + + sample_hidden_states = text_hidden_states[logits_indices] + logits = self.model.compute_logits(sample_hidden_states, None) + if broadcast_pp_output: + model_output_broadcast_data = { + "logits": logits.contiguous(), + } if logits is not None else {} + model_output_broadcast_data = get_pp_group().broadcast_tensor_dict( + model_output_broadcast_data, src=len(get_pp_group().ranks) - 1) + assert model_output_broadcast_data is not None + logits = model_output_broadcast_data["logits"] + + # Apply structured output bitmasks if present + if scheduler_output.grammar_bitmask is not None: + self.apply_grammar_bitmask(scheduler_output, logits) + + # Sample the next token and get logprobs if needed (with spec decode) + sampling_metadata = self.input_batch.sampling_metadata + if spec_decode_metadata is None: + sampler_output = self.sampler( + logits=logits, + sampling_metadata=sampling_metadata, + ) + else: + assert logits is not None + bonus_logits = logits[spec_decode_metadata.bonus_logits_indices] + sampler_output = self.sampler( + logits=bonus_logits, + sampling_metadata=sampling_metadata, + ) + bonus_token_ids = sampler_output.sampled_token_ids + target_logits = logits[spec_decode_metadata.target_logits_indices] + output_token_ids = self.rejection_sampler( + spec_decode_metadata, + None, + target_logits, + bonus_token_ids, + sampling_metadata, + ) + sampler_output.sampled_token_ids = output_token_ids + + num_nans_in_logits = {} + if envs.VLLM_COMPUTE_NANS_IN_LOGITS: + num_nans_in_logits = self._get_nans_in_logits(logits) + + # Handle partial prefill: discard sampled tokens and rewind RNG + discard_sampled_tokens_req_indices = [] + for i, req_id in enumerate(self.input_batch.req_ids): + req_state = self.requests[req_id] + seq_len = (req_state.num_computed_tokens + + scheduler_output.num_scheduled_tokens[req_id]) + if seq_len < req_state.num_tokens: + generator = self.input_batch.generators.get(i) + if generator is not None: + generator.set_offset(generator.get_offset() - 4) + discard_sampled_tokens_req_indices.append(i) + + # Move CPU sync parts + logprobs_tensors = sampler_output.logprobs_tensors + logprobs_lists = logprobs_tensors.tolists() \ + if logprobs_tensors is not None else None + + # Prompt logprobs if needed + prompt_logprobs_dict = self._get_prompt_logprobs_dict( + text_hidden_states[:num_scheduled_tokens], + scheduler_output, + ) + + # Parse valid sampled tokens + sampled_token_ids = sampler_output.sampled_token_ids + max_gen_len = sampled_token_ids.shape[-1] + if max_gen_len == 1: + valid_sampled_token_ids = sampled_token_ids.tolist() + else: + valid_sampled_token_ids = self.rejection_sampler.parse_output( + sampled_token_ids, + self.input_batch.vocab_size, + ) + for i in discard_sampled_tokens_req_indices: + valid_sampled_token_ids[i].clear() + + # Cache sampled tokens + for req_idx, sampled_ids in enumerate(valid_sampled_token_ids): + if not sampled_ids: + continue + start_idx = self.input_batch.num_tokens_no_spec[req_idx] + end_idx = start_idx + len(sampled_ids) + self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids + self.input_batch.num_tokens_no_spec[req_idx] = end_idx + self.input_batch.num_tokens[req_idx] = end_idx + req_id = self.input_batch.req_ids[req_idx] + req_state = self.requests[req_id] + req_state.output_token_ids.extend(sampled_ids) + + # Speculative decoding draft tokens if configured + if not self.speculative_config: + spec_token_ids = None + else: + assert spec_decode_common_attn_metadata is not None + spec_token_ids = self.propose_draft_token_ids( + scheduler_output, + valid_sampled_token_ids, + sampling_metadata, + text_hidden_states, + sample_hidden_states, + _aux_hidden_states, + spec_decode_metadata, + spec_decode_common_attn_metadata, + ) + + # Convert to per-request tensors on CPU + text_hidden_states_cpu = text_hidden_states.detach().to("cpu").contiguous() + pooler_output: list[Optional[torch.Tensor]] = [] + prev_logits_index = 0 + for logits_index in logits_indices: + pooler_output.append(text_hidden_states_cpu[prev_logits_index:logits_index+1]) + prev_logits_index = logits_index + 1 + + + self.eplb_step() + + return OmniModelRunnerOutput( + req_ids=self.input_batch.req_ids, + req_id_to_index=self.input_batch.req_id_to_index, + sampled_token_ids=valid_sampled_token_ids, + spec_token_ids=spec_token_ids, + logprobs=logprobs_lists, + prompt_logprobs_dict=prompt_logprobs_dict, + pooler_output=pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None, + kv_connector_output=kv_connector_output, + num_nans_in_logits=num_nans_in_logits, + ) + + @torch.inference_mode() + def extract_multimodal_outputs(self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]) -> dict: + if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs: + text_hidden_states = hidden_states.text_hidden_states + multimodal_outputs = hidden_states.multimodal_outputs + + elif isinstance(hidden_states, torch.Tensor): + text_hidden_states = hidden_states + multimodal_outputs = {} + elif isinstance(hidden_states, List): + text_hidden_states = hidden_states[0] + multimodal_outputs = {} + else: + raise ValueError(f"Invalid hidden states type: {type(hidden_states)}") + return text_hidden_states, multimodal_outputs + + @torch.inference_mode() + def _dummy_run( + self, + num_tokens: int, + capture_attn_cudagraph: bool = False, + skip_eplb: bool = False, + is_profile: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + + # Padding for DP + num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) + num_tokens += num_pad + + # Set num_scheduled_tokens based on num_tokens and max_num_seqs + # for dummy run with LoRA so that the num_reqs collectively + # has num_tokens in total. + assert num_tokens <= self.scheduler_config.max_num_batched_tokens + max_num_reqs = self.scheduler_config.max_num_seqs + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs + num_scheduled_tokens_list[-1] += num_tokens % num_reqs + assert sum(num_scheduled_tokens_list) == num_tokens + assert len(num_scheduled_tokens_list) == num_reqs + num_scheduled_tokens = np.array(num_scheduled_tokens_list, + dtype=np.int32) + + attn_metadata: Optional[dict[str, Any]] = None + if capture_attn_cudagraph: + attn_metadata = {} + + # Make sure max_model_len is used at the graph capture time. + self.seq_lens_np[:num_reqs] = self.max_model_len + self.seq_lens_np[num_reqs:] = 0 + self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs], + non_blocking=True) + + for kv_cache_group_id, kv_cache_group_spec in enumerate( + self.kv_cache_config.kv_cache_groups): + common_attn_metadata = CommonAttentionMetadata( + query_start_loc=self.query_start_loc[:num_reqs + 1], + query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + + 1], + seq_lens=self.seq_lens[:num_reqs], + seq_lens_cpu=self.seq_lens_cpu[:num_reqs], + num_computed_tokens_cpu=self.input_batch. + num_computed_tokens_cpu_tensor[:num_reqs], + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + max_query_len=num_tokens, + block_table_tensor=self.input_batch.block_table[ + kv_cache_group_id].get_device_tensor()[:num_reqs], + slot_mapping=self.input_batch. + block_table[kv_cache_group_id].slot_mapping[:num_tokens], + causal=True) + + for attn_group in self.attn_groups[kv_cache_group_id]: + attn_metadata_i = attn_group.metadata_builder\ + .build_for_cudagraph_capture(common_attn_metadata) + for layer_name in kv_cache_group_spec.layer_names: + attn_metadata[layer_name] = attn_metadata_i + + with self.maybe_dummy_run_with_lora(self.lora_config, + num_scheduled_tokens): + if self.is_multimodal_model: + input_ids = None + inputs_embeds = self.inputs_embeds[:num_tokens] + model_mm_kwargs = self._dummy_mm_kwargs(num_reqs) + else: + input_ids = self.input_ids[:num_tokens] + inputs_embeds = None + model_mm_kwargs = {} + + if self.uses_mrope: + positions = self.mrope_positions[:, :num_tokens] + else: + positions = self.positions[:num_tokens] + + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + if self.intermediate_tensors is None: + self.intermediate_tensors = ( + self.model.make_empty_intermediate_tensors( + batch_size=self.max_num_tokens, + dtype=self.model_config.dtype, + device=self.device)) + + intermediate_tensors = self.sync_and_slice_intermediate_tensors( + num_tokens, None, False) + + with self.maybe_randomize_inputs(input_ids), set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp): + outputs = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **MultiModalKwargs.as_kwargs( + model_mm_kwargs, + device=self.device, + ), + ) + + if self.use_aux_hidden_state_outputs: + hidden_states, _ = outputs + else: + hidden_states = outputs + + if self.speculative_config and self.speculative_config.use_eagle(): + assert isinstance(self.drafter, EagleProposer) + self.drafter.dummy_run(num_tokens) + + # This is necessary to avoid blocking DP. + # For dummy runs, we typically skip EPLB since we don't have any real + # requests to process. + # However, in DP settings, there may be cases when some DP ranks do + # not have any requests to process, so they're executing dummy batches. + # In such cases, we still have to trigger EPLB to make sure + # ranks execute the rearrangement in synchronization. + if not skip_eplb: + self.eplb_step(is_dummy=True, is_profile=is_profile) + + logit_indices = np.cumsum(num_scheduled_tokens) - 1 + hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) + return hidden_states, hidden_states[logit_indices] + diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py new file mode 100644 index 000000000..fb22ce15f --- /dev/null +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -0,0 +1,67 @@ +from vllm.v1.worker.gpu_worker import Worker as GPUWorker +import torch +import os +import gc +from vllm.utils import GiB_bytes, MemorySnapshot +from vllm.platforms import current_platform +from vllm.worker.worker import _check_if_gpu_supports_dtype +from vllm.v1.worker.gpu_worker import init_worker_distributed_environment +from vllm.model_executor import set_random_seed +from vllm.v1.utils import report_usage_stats + +from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + + +class GPUARWorker(GPUWorker): + def init_device(self): + if self.device_config.device.type == "cuda": + # torch.distributed.all_reduce does not free the input tensor until + # the synchronization point. This causes the memory usage to grow + # as the number of all_reduce calls increases. This env var disables + # this behavior. + # Related issue: + # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 + os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" + + # This env var set by Ray causes exceptions with graph building. + os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + self.device = torch.device(f"cuda:{self.local_rank}") + current_platform.set_device(self.device) + + _check_if_gpu_supports_dtype(self.model_config.dtype) + gc.collect() + torch.cuda.empty_cache() + + # take current memory snapshot + self.init_snapshot = MemorySnapshot() + self.requested_memory = (self.init_snapshot.total_memory * + self.cache_config.gpu_memory_utilization) + if self.init_snapshot.free_memory < self.requested_memory: + GiB = lambda b: round(b / GiB_bytes, 2) + raise ValueError( + f"Free memory on device " + f"({GiB(self.init_snapshot.free_memory)}/" + f"{GiB(self.init_snapshot.total_memory)} GiB) on startup " + f"is less than desired GPU memory utilization " + f"({self.cache_config.gpu_memory_utilization}, " + f"{GiB(self.requested_memory)} GiB). Decrease GPU memory " + f"utilization or reduce GPU memory used by other processes." + ) + else: + raise RuntimeError( + f"Not support device type: {self.device_config.device}") + # Initialize the distributed environment. + init_worker_distributed_environment(self.vllm_config, self.rank, + self.distributed_init_method, + self.local_rank, + current_platform.dist_backend) + # Set random seed. + set_random_seed(self.model_config.seed) + + # Construct the model runner + self.model_runner: GPUARModelRunner = GPUARModelRunner( + self.vllm_config, self.device) + + if self.rank == 0: + # If usage stat is enabled, collect relevant info. + report_usage_stats(self.vllm_config) \ No newline at end of file diff --git a/vllm_omni/worker/gpu_diffusion_model_runner.py b/vllm_omni/worker/gpu_diffusion_model_runner.py index e46492294..6626844d9 100644 --- a/vllm_omni/worker/gpu_diffusion_model_runner.py +++ b/vllm_omni/worker/gpu_diffusion_model_runner.py @@ -1,133 +1,362 @@ -"""Model runner that wraps a diffusers text-to-image pipeline.""" - from __future__ import annotations -from dataclasses import dataclass -from typing import Any, Optional +from typing import Optional, Union, List +import numpy as np +import gc +import logging -@dataclass -class DiffusionRunnerOutput: - """Structured output returned by the diffusion runner.""" +import torch - prompt: str - images: list - request_id: str = "diffusion_request" - finished: bool = True - output_type: str = "image" +from vllm.v1.worker.gpu_model_runner import ( + EMPTY_MODEL_RUNNER_OUTPUT, + IntermediateTensors, + get_pp_group, + set_forward_context, +) +from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs +from vllm.multimodal.inputs import MultiModalKwargs + +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.attention.backends.utils import CommonAttentionMetadata + +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner -class DiffusionModelRunner: - """Lightweight runner that loads and executes diffusers pipelines.""" +logger = logging.getLogger(__name__) - def __init__( + +class GPUDiffusionModelRunner(OmniGPUModelRunner): + """Diffusion model runner for vLLM-omni (non-autoregressive). + + - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue. + - Does not compute logits or perform token sampling. + - Executes diffusion process and returns tensors via `pooler_output`. + """ + + @torch.inference_mode() + def execute_model( self, - model_path: str, - *, - pipeline_name: Optional[str] = None, - device: Optional[str] = None, - dtype: Optional[str] = None, - ) -> None: - self.model_path = model_path - self.pipeline_name = pipeline_name or "auto" - self.device, self.torch_dtype = self._resolve_device_and_dtype(device, dtype) - self._pipeline = self._load_pipeline() + scheduler_output: "SchedulerOutput", + intermediate_tensors: Optional[IntermediateTensors] = None, + ) -> Union[OmniModelRunnerOutput, IntermediateTensors]: + self._update_states(scheduler_output) - def _resolve_device_and_dtype( - self, - device: Optional[str], - dtype: Optional[str], - ): - import torch + if not scheduler_output.total_num_scheduled_tokens: - resolved_device = device - if resolved_device is None: - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - resolved_device = "mps" - elif torch.cuda.is_available(): - resolved_device = "cuda" - else: - resolved_device = "cpu" + return EMPTY_MODEL_RUNNER_OUTPUT + + # Prepare decoder inputs and attention metadata (for batch/order mapping) + (attn_metadata, attention_cuda_graphs, logits_indices, + spec_decode_metadata, num_scheduled_tokens_np, + spec_decode_common_attn_metadata) = self._prepare_inputs( + scheduler_output) - if dtype is not None: - resolved_dtype = getattr(torch, dtype, None) + # Input token count for this iteration (not used by diffusion, but + # retained to keep DP padding/ordering consistent) + num_input_tokens = scheduler_output.total_num_scheduled_tokens + num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens) + num_input_tokens += num_pad + + # Multimodal conditioning (e.g., text/audio/video encoders) + if self.is_multimodal_model: + self._execute_mm_encoder(scheduler_output) + mm_embeds = self._gather_mm_embeddings(scheduler_output) else: - if resolved_device in {"cuda", "mps"}: - resolved_dtype = torch.float16 - else: - resolved_dtype = torch.float32 + mm_embeds = [] - return resolved_device, resolved_dtype - - def _load_pipeline(self): - import diffusers - - if self.pipeline_name != "auto": - PipelineCls = getattr(diffusers, self.pipeline_name) + # Build inputs to mirror AR runner: input_ids/positions/embeds + if self.is_multimodal_model and get_pp_group().is_first_rank: + inputs_embeds_scheduled = self.model.get_input_embeddings( + input_ids=self.input_ids[:scheduler_output.total_num_scheduled_tokens], + multimodal_embeddings=mm_embeds or None, + ) + self.inputs_embeds[:scheduler_output.total_num_scheduled_tokens].copy_( + inputs_embeds_scheduled) + input_ids = self.input_ids[:num_input_tokens] + inputs_embeds = self.inputs_embeds[:num_input_tokens] + model_mm_kwargs = self._extract_mm_kwargs(scheduler_output) else: - PipelineCls = diffusers.AutoPipelineForText2Image + input_ids = self.input_ids[:num_input_tokens] + inputs_embeds = None + model_mm_kwargs = {} - pipeline = PipelineCls.from_pretrained( - self.model_path, - torch_dtype=self.torch_dtype, + # Positions (mrope or standard) + if self.uses_mrope: + positions = self.mrope_positions[:, :num_input_tokens] + else: + positions = self.positions[:num_input_tokens] + + # Intermediate tensors sync for PP (if any) + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + intermediate_tensors = self.sync_and_slice_intermediate_tensors( + num_input_tokens, intermediate_tensors, True) + + # Set forward context mainly for resource management and kv connector + skip_cuda_graphs = True # diffusion path does not rely on cuda graphs here + with set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_input_tokens, + num_tokens_across_dp=num_tokens_across_dp, + skip_cuda_graphs=skip_cuda_graphs, + ), self.maybe_get_kv_connector_output( + scheduler_output) as kv_connector_output: + + if not get_pp_group().is_last_rank: + # For non-last PP stages, pass through intermediate tensors. + assert intermediate_tensors is not None + intermediate_tensors.kv_connector_output = kv_connector_output + return intermediate_tensors + + outputs = self._run_diffusion( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + multimodal_kwargs=model_mm_kwargs, + logits_indices=logits_indices, + ) + _, multimodal_outputs = ( + self.extract_multimodal_outputs(outputs)) + + # Ensure one tensor per request, map to CPU for output struct + pooler_output: List[Optional[torch.Tensor]] = [] + if isinstance(multimodal_outputs, torch.Tensor): + # If model returned a single stacked tensor, split by requests + assert multimodal_outputs.shape[0] == self.input_batch.num_reqs + for i in range(self.input_batch.num_reqs): + pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous()) + elif isinstance(multimodal_outputs, list): + for out in multimodal_outputs: + pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) + elif isinstance(multimodal_outputs, dict): + for out in multimodal_outputs.values(): + pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) + else: + raise RuntimeError("Unsupported diffusion output type") + + self.eplb_step() + + return OmniModelRunnerOutput( + req_ids=self.input_batch.req_ids, + req_id_to_index=self.input_batch.req_id_to_index, + sampled_token_ids=[], + spec_token_ids=None, + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=pooler_output, + kv_connector_output=kv_connector_output, + num_nans_in_logits={}, ) - try: - pipeline = pipeline.to(self.device) - except Exception: - pipeline = pipeline.to("cpu") - - self._apply_optimisations(pipeline) - return pipeline - - @staticmethod - def _apply_optimisations(pipeline) -> None: - # Enable attention slicing / VAE tiling where available to reduce memory. - try: - if hasattr(pipeline, "enable_attention_slicing"): - pipeline.enable_attention_slicing() - except Exception: - pass - try: - vae = getattr(pipeline, "vae", None) - if vae and hasattr(vae, "enable_tiling"): - vae.enable_tiling() - except Exception: - pass - - def generate( + def _run_diffusion( self, - prompt: str, *, - height: int = 512, - width: int = 512, - num_inference_steps: int = 30, - guidance_scale: float = 5.0, - negative_prompt: Optional[str] = None, - seed: Optional[int] = None, - image: Optional[Any] = None, - ) -> DiffusionRunnerOutput: - import torch + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors], + inputs_embeds: Optional[torch.Tensor], + multimodal_kwargs: dict, + logits_indices: torch.Tensor, + ) -> Union[torch.Tensor, list[torch.Tensor]]: + """Runs the diffusion process and returns per-request tensors. - generator = None - if seed is not None: - try: - generator = torch.Generator(device=self.device).manual_seed(int(seed)) - except Exception: - # Some backends (e.g., CPU without full support) may not accept the device. - generator = torch.Generator().manual_seed(int(seed)) - - output = self._pipeline( - prompt=prompt, - negative_prompt=negative_prompt, - num_inference_steps=int(num_inference_steps), - guidance_scale=float(guidance_scale), - height=int(height), - width=int(width), - generator=generator, - image=image, + Tries model interfaces in the following order for maximal compatibility: + 1) model.sample(condition=..., **kwargs) + 2) model.forward(condition=..., **kwargs) + 3) model.diffuse(condition=..., **kwargs) + """ + # Keep inputs identical to AR runner + kwargs = dict( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device), + sampling_metadata=self.input_batch.sampling_metadata, + logits_index=logits_indices, + sampler=self.sampler, ) - return DiffusionRunnerOutput( - prompt=prompt, - images=getattr(output, "images", []) or [], - ) + # For Qwen 2.5 Omni's current implementation, we only support the forward method + if hasattr(self.model, "forward"): + return self.model.forward(**kwargs) + #TODO: add the diffuse method for other models + + raise RuntimeError( + "The loaded model does not expose diffusion interfaces 'sample', " + "'forward', or 'diffuse'. Please implement one of them or adapt the runner.") + + + @torch.inference_mode() + def _dummy_run( + self, + num_tokens: int, + capture_attn_cudagraph: bool = False, + skip_eplb: bool = False, + is_profile: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Padding for DP + num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) + num_tokens += num_pad + + # Set num_scheduled_tokens based on num_tokens and max_num_seqs + # for dummy run with LoRA so that the num_reqs collectively + # has num_tokens in total. + assert num_tokens <= self.scheduler_config.max_num_batched_tokens + max_num_reqs = self.scheduler_config.max_num_seqs + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs + num_scheduled_tokens_list[-1] += num_tokens % num_reqs + assert sum(num_scheduled_tokens_list) == num_tokens + assert len(num_scheduled_tokens_list) == num_reqs + num_scheduled_tokens = np.array(num_scheduled_tokens_list, + dtype=np.int32) + + attn_metadata: Optional[dict[str, dict]] = None + if capture_attn_cudagraph: + attn_metadata = {} + + # Make sure max_model_len is used at the graph capture time. + self.seq_lens_np[:num_reqs] = self.max_model_len + self.seq_lens_np[num_reqs:] = 0 + self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs], + non_blocking=True) + + for kv_cache_group_id, kv_cache_group_spec in enumerate( + self.kv_cache_config.kv_cache_groups): + common_attn_metadata = CommonAttentionMetadata( + query_start_loc=self.query_start_loc[:num_reqs + 1], + query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1], + seq_lens=self.seq_lens[:num_reqs], + seq_lens_cpu=self.seq_lens_cpu[:num_reqs], + num_computed_tokens_cpu=self.input_batch. + num_computed_tokens_cpu_tensor[:num_reqs], + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + max_query_len=num_tokens, + block_table_tensor=self.input_batch.block_table[ + kv_cache_group_id].get_device_tensor()[:num_reqs], + slot_mapping=self.input_batch.block_table[ + kv_cache_group_id].slot_mapping[:num_tokens], + causal=True) + + for attn_group in self.attn_groups[kv_cache_group_id]: + attn_metadata_i = attn_group.metadata_builder \ + .build_for_cudagraph_capture(common_attn_metadata) + for layer_name in kv_cache_group_spec.layer_names: + attn_metadata[layer_name] = attn_metadata_i + + with self.maybe_dummy_run_with_lora(self.lora_config, + num_scheduled_tokens): + if self.is_multimodal_model: + input_ids = None + inputs_embeds = self.inputs_embeds[:num_tokens] + model_mm_kwargs = self._dummy_mm_kwargs(num_reqs) + else: + input_ids = self.input_ids[:num_tokens] + inputs_embeds = None + model_mm_kwargs = {} + + if self.uses_mrope: + positions = self.mrope_positions[:, :num_tokens] + else: + positions = self.positions[:num_tokens] + + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + if self.intermediate_tensors is None: + self.intermediate_tensors = ( + self.model.make_empty_intermediate_tensors( + batch_size=self.max_num_tokens, + dtype=self.model_config.dtype, + device=self.device)) + + intermediate_tensors = self.sync_and_slice_intermediate_tensors( + num_tokens, None, False) + + # Diffusion path: avoid CUDA graphs; we only use context for resource wiring + with self.maybe_randomize_inputs(input_ids), set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp): + outputs = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **MultiModalKwargs.as_kwargs( + model_mm_kwargs, + device=self.device, + ), + sampler=None, + ) + + if self.use_aux_hidden_state_outputs: + hidden_states, _ = outputs + else: + hidden_states = outputs + + # Extract multimodal outputs if present; we ignore them here because + # dummy run returns tensors only. The actual diffusion runner returns + # multimodal outputs via pooler_output in execute_model. + text_hidden_states, _ = self.extract_multimodal_outputs(hidden_states) + + if not skip_eplb: + self.eplb_step(is_dummy=True, is_profile=is_profile) + + logit_indices = np.cumsum(num_scheduled_tokens) - 1 + return text_hidden_states, None + + @torch.inference_mode() + def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: + logger.warning("Dummy sampler run is not implemented for diffusion model") + return None + + def profile_run(self) -> None: + # Profile with multimodal encoder & encoder cache, similar to base but + # without any logits/sampler warming. + if self.is_multimodal_model: + mm_budget = self.mm_budget + assert mm_budget is not None + + # TODO: handle encoder-decoder models once supported. + if (encoder_budget := mm_budget.get_encoder_budget()) > 0: + ( + dummy_modality, + max_tokens, + ) = mm_budget.get_modality_with_max_tokens() + ( + max_mm_items_per_prompt, + max_mm_items_per_batch, + ) = mm_budget.get_max_items(dummy_modality, max_tokens) + + batched_dummy_mm_inputs = self._get_mm_dummy_batch( + dummy_modality, + max_mm_items_per_batch, + ) + + dummy_encoder_outputs = self.model.get_multimodal_embeddings( + **batched_dummy_mm_inputs) + + sanity_check_mm_encoder_outputs( + dummy_encoder_outputs, + expected_num_items=max_mm_items_per_batch, + ) + + self.encoder_cache["tmp"] = dict( + enumerate(dummy_encoder_outputs)) + + hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True) + if get_pp_group().is_last_rank: + pass # No sampler/pooler warmup for diffusion + self._sync_device() + del hidden_states + self.encoder_cache.clear() + gc.collect() diff --git a/vllm_omni/worker/gpu_diffusion_worker.py b/vllm_omni/worker/gpu_diffusion_worker.py index 74f2474c4..be4c2047d 100644 --- a/vllm_omni/worker/gpu_diffusion_worker.py +++ b/vllm_omni/worker/gpu_diffusion_worker.py @@ -1,56 +1,67 @@ -"""Worker abstraction for executing diffusion model runners.""" +from vllm.v1.worker.gpu_worker import Worker as GPUWorker +import torch +import os +import gc +from vllm.utils import GiB_bytes, MemorySnapshot +from vllm.platforms import current_platform +from vllm.worker.worker import _check_if_gpu_supports_dtype +from vllm.v1.worker.gpu_worker import init_worker_distributed_environment +from vllm.model_executor import set_random_seed +from vllm.v1.utils import report_usage_stats -from __future__ import annotations - -from typing import Any, Optional - -from vllm.v1.worker.worker_base import WorkerBase - -from .gpu_diffusion_model_runner import DiffusionModelRunner, DiffusionRunnerOutput +from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner -class DiffusionGPUWorker(WorkerBase): - """Minimal worker wrapping DiffusionModelRunner. +class GPUDiffusionWorker(GPUWorker): + def init_device(self): + if self.device_config.device.type == "cuda": + # torch.distributed.all_reduce does not free the input tensor until + # the synchronization point. This causes the memory usage to grow + # as the number of all_reduce calls increases. This env var disables + # this behavior. + # Related issue: + # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 + os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" - The worker interface mirrors the shape expected by vLLM executors: it owns - a model runner instance and exposes a `generate` method that can be called - by an executor implementation. - """ + # This env var set by Ray causes exceptions with graph building. + os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + self.device = torch.device(f"cuda:{self.local_rank}") + current_platform.set_device(self.device) - def __init__( - self, - model_path: str, - *, - pipeline_name: Optional[str] = None, - device: Optional[str] = None, - dtype: Optional[str] = None, - ) -> None: - self.model_runner = DiffusionModelRunner( - model_path, - pipeline_name=pipeline_name, - device=device, - dtype=dtype, - ) + _check_if_gpu_supports_dtype(self.model_config.dtype) + gc.collect() + torch.cuda.empty_cache() - def generate( - self, - prompt: str, - *, - height: int = 512, - width: int = 512, - num_inference_steps: int = 30, - guidance_scale: float = 5.0, - negative_prompt: Optional[str] = None, - seed: Optional[int] = None, - image: Optional[Any] = None, - ) -> DiffusionRunnerOutput: - return self.model_runner.generate( - prompt=prompt, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - negative_prompt=negative_prompt, - seed=seed, - image=image, - ) + # take current memory snapshot + self.init_snapshot = MemorySnapshot() + self.requested_memory = (self.init_snapshot.total_memory * + self.cache_config.gpu_memory_utilization) + if self.init_snapshot.free_memory < self.requested_memory: + GiB = lambda b: round(b / GiB_bytes, 2) + raise ValueError( + f"Free memory on device " + f"({GiB(self.init_snapshot.free_memory)}/" + f"{GiB(self.init_snapshot.total_memory)} GiB) on startup " + f"is less than desired GPU memory utilization " + f"({self.cache_config.gpu_memory_utilization}, " + f"{GiB(self.requested_memory)} GiB). Decrease GPU memory " + f"utilization or reduce GPU memory used by other processes." + ) + else: + raise RuntimeError( + f"Not support device type: {self.device_config.device}") + # Initialize the distributed environment. + init_worker_distributed_environment(self.vllm_config, self.rank, + self.distributed_init_method, + self.local_rank, + current_platform.dist_backend) + # Set random seed. + set_random_seed(self.model_config.seed) + + # Construct the model runner + self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner( + self.vllm_config, self.device) + + if self.rank == 0: + # If usage stat is enabled, collect relevant info. + report_usage_stats(self.vllm_config) \ No newline at end of file diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py new file mode 100644 index 000000000..d1d73e6f3 --- /dev/null +++ b/vllm_omni/worker/gpu_model_runner.py @@ -0,0 +1,755 @@ + +from typing import TYPE_CHECKING, Any, Optional, Union, cast + +import numpy as np +import torch + +import vllm.envs as envs +from vllm.distributed.parallel_state import get_pp_group, get_tp_group +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces_base import VllmModelForPooling +from vllm.multimodal.inputs import MultiModalKwargs +from vllm.sequence import IntermediateTensors +from vllm.utils import LazyLoader +from vllm.v1.attention.backends.utils import CommonAttentionMetadata +from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT +from vllm.v1.worker.gpu_input_batch import CachedRequestState +from vllm.sampling_params import SamplingType +from vllm.distributed.kv_transfer import has_kv_transfer_group +from vllm.utils import round_up +from vllm.v1.spec_decode.eagle import EagleProposer + +from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding + +if TYPE_CHECKING: + from vllm.v1.core.sched.output import SchedulerOutput +else: + xgr = LazyLoader("xgr", globals(), "xgrammar") + xgr_torch_compile = LazyLoader( + "xgr_torch_compile", globals(), + "xgrammar.kernels.apply_token_bitmask_inplace_torch_compile") + +logger = init_logger(__name__) + + +class OmniGPUModelRunner(GPUModelRunner): + def _update_states(self, scheduler_output: "SchedulerOutput") -> None: + """Update the cached states and the persistent batch with the scheduler + output. + + The updated states are used by the `_prepare_inputs` function to create + the input GPU tensors for the model. + + The SamplingMetadata is updated and copied to the GPU if there is a + new/resumed/paused/finished request in the batch. + """ + # Remove finished requests from the cached states. + for req_id in scheduler_output.finished_req_ids: + self.requests.pop(req_id, None) + self.encoder_cache.pop(req_id, None) + # Remove the finished requests from the persistent batch. + # NOTE(woosuk): There could be an edge case where finished_req_ids and + # scheduled_req_ids overlap. This happens when a request is aborted and + # then resubmitted with the same ID. In this case, we treat them as two + # distinct requests - clearing the cached states for the first request + # and handling the second as a new request. + for req_id in scheduler_output.finished_req_ids: + self.input_batch.remove_request(req_id) + + # Free the cached encoder outputs. + for req_id, input_id in scheduler_output.free_encoder_input_ids: + encoder_outputs = self.encoder_cache.get(req_id) + if encoder_outputs is not None: + encoder_outputs.pop(input_id, None) + if not encoder_outputs: + self.encoder_cache.pop(req_id, None) + + # Remove the unscheduled requests from the persistent batch. + # NOTE(woosuk): The unscheduled requests are either preempted requests + # or running requests that are not scheduled in this step. We remove + # them from the persistent batch but keep their cached states since + # they will be scheduled again sometime in the future. + scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys() + cached_req_ids = self.input_batch.req_id_to_index.keys() + unscheduled_req_ids = cached_req_ids - scheduled_req_ids + # NOTE(woosuk): The persistent batch optimization assumes that + # consecutive batches contain mostly the same requests. If batches + # have low request overlap (e.g., alternating between two distinct + # sets of requests), this optimization becomes very inefficient. + for req_id in unscheduled_req_ids: + self.input_batch.remove_request(req_id) + + req_ids_to_add: list[str] = [] + # Add new requests to the cached states. + for new_req_data in scheduler_output.scheduled_new_reqs: + req_id = new_req_data.req_id + sampling_params = new_req_data.sampling_params + pooling_params = new_req_data.pooling_params + + if sampling_params and \ + sampling_params.sampling_type == SamplingType.RANDOM_SEED: + generator = torch.Generator(device=self.device) + generator.manual_seed(sampling_params.seed) + else: + generator = None + + if pooling_params: + assert (task := pooling_params.task) is not None, ( + "You did not set `task` in the API") + + model = cast(VllmModelForPooling, self.model) + to_update = model.pooler.get_pooling_updates(task) + to_update.apply(pooling_params) + + self.requests[req_id] = CachedRequestState( + req_id=req_id, + prompt_token_ids=new_req_data.prompt_token_ids, + mm_inputs=new_req_data.mm_inputs, + mm_positions=new_req_data.mm_positions, + sampling_params=sampling_params, + pooling_params=pooling_params, + generator=generator, + block_ids=new_req_data.block_ids, + num_computed_tokens=new_req_data.num_computed_tokens, + output_token_ids=[], + lora_request=new_req_data.lora_request, + ) + + # If prompt embeddings are provided, decode and attach to inter_data + try: + if getattr(new_req_data, "prompt_embeds", None) is not None: + payload = new_req_data.prompt_embeds + dtype = getattr(np, payload.dtype) + arr = np.frombuffer(payload.data, dtype=dtype) + arr = arr.reshape(payload.shape) + pe_cpu = torch.from_numpy(arr) + # Store temporarily on CPU; later moved to device in builder + setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) + # Also replace payload with Tensor for user visibility in scheduler_output + try: + new_req_data.prompt_embeds = pe_cpu # type: ignore[assignment] + except Exception: + pass + except Exception: + pass + # Decode additional_information payloads (dictionary) + try: + if getattr(new_req_data, "additional_information", None) is not None: + payload_info = new_req_data.additional_information + info_dict = {} + if isinstance(payload_info, dict): + info_dict = payload_info + else: + from vllm.v1.engine import AdditionalInformationPayload + if isinstance(payload_info, AdditionalInformationPayload): + for k, entry in payload_info.entries.items(): + if entry.tensor_data is not None: + dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) + arr = np.frombuffer(entry.tensor_data, dtype=dt) + arr = arr.reshape(entry.tensor_shape) + info_dict[k] = torch.from_numpy(arr) + else: + info_dict[k] = entry.list_data + if info_dict: + setattr(self.requests[req_id], "additional_information_cpu", info_dict) + except Exception: + pass + + # Only relevant for models using M-RoPE (e.g, Qwen2-VL) + if self.uses_mrope: + image_grid_thw = [] + video_grid_thw = [] + second_per_grid_ts = [] + audio_feature_lengths = [] + use_audio_in_video = False + for mm_input in self.requests[req_id].mm_inputs: + if mm_input.get("image_grid_thw") is not None: + image_grid_thw.extend( + mm_input["image_grid_thw"].tolist()) + if mm_input.get("video_grid_thw") is not None: + video_grid_thw.extend( + mm_input["video_grid_thw"].tolist()) + if mm_input.get("second_per_grid_ts") is not None: + second_per_grid_ts.extend( + mm_input["second_per_grid_ts"]) + if mm_input.get("audio_feature_lengths") is not None: + audio_feature_lengths.extend( + mm_input["audio_feature_lengths"]) + if mm_input.get("use_audio_in_video") is True: + use_audio_in_video = True + + hf_config = self.model_config.hf_config + + self.requests[req_id].mrope_positions, \ + self.requests[req_id].mrope_position_delta = \ + MRotaryEmbedding.get_input_positions_tensor( + self.requests[req_id].prompt_token_ids, + hf_config=hf_config, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + second_per_grid_ts=second_per_grid_ts, + audio_feature_lengths=audio_feature_lengths, + use_audio_in_video=use_audio_in_video, + ) + + req_ids_to_add.append(req_id) + + # Update the states of the running/resumed requests. + is_last_rank = get_pp_group().is_last_rank + req_data = scheduler_output.scheduled_cached_reqs + for i, req_id in enumerate(req_data.req_ids): + req_state = self.requests[req_id] + num_computed_tokens = req_data.num_computed_tokens[i] + new_block_ids = req_data.new_block_ids[i] + resumed_from_preemption = req_data.resumed_from_preemption[i] + + # Update the cached states. + req_state.num_computed_tokens = num_computed_tokens + + if not is_last_rank: + # When using PP, the scheduler sends the sampled tokens back, + # because there's no direct communication between the first- + # stage worker and the last-stage worker. + new_token_ids = req_data.new_token_ids[i] + # Add the sampled token(s) from the previous step (if any). + # This doesn't include "unverified" tokens like spec tokens. + num_new_tokens = (num_computed_tokens + len(new_token_ids) - + req_state.num_tokens) + if num_new_tokens == 1: + # Avoid slicing list in most common case. + req_state.output_token_ids.append(new_token_ids[-1]) + elif num_new_tokens > 0: + req_state.output_token_ids.extend( + new_token_ids[-num_new_tokens:]) + + # Update the block IDs. + if not resumed_from_preemption: + # Append the new blocks to the existing block IDs. + for block_ids, new_ids in zip(req_state.block_ids, + new_block_ids): + block_ids.extend(new_ids) + else: + # The request is resumed from preemption. + # Replace the existing block IDs with the new ones. + req_state.block_ids = new_block_ids + + req_index = self.input_batch.req_id_to_index.get(req_id) + if req_index is None: + # The request is not in the persistent batch. + # The request was either preempted and resumed later, or was not + # scheduled in the previous step and needs to be added again. + req_ids_to_add.append(req_id) + continue + + # Update the persistent batch. + self.input_batch.num_computed_tokens_cpu[req_index] = ( + num_computed_tokens) + self.input_batch.block_table.append_row(new_block_ids, req_index) + + # For the last rank, we don't need to update the token_ids_cpu + # because the sampled tokens are already cached. + if not is_last_rank: + # Add new_token_ids to token_ids_cpu. + start_token_index = num_computed_tokens + end_token_index = num_computed_tokens + len(new_token_ids) + self.input_batch.token_ids_cpu[ + req_index, + start_token_index:end_token_index] = new_token_ids + self.input_batch.num_tokens_no_spec[ + req_index] = end_token_index + self.input_batch.num_tokens[req_index] = end_token_index + + # Add spec_token_ids to token_ids_cpu. + spec_token_ids = ( + scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())) + if spec_token_ids: + num_spec_tokens = len(spec_token_ids) + start_index = self.input_batch.num_tokens_no_spec[req_index] + end_token_index = start_index + num_spec_tokens + self.input_batch.token_ids_cpu[ + req_index, start_index:end_token_index] = spec_token_ids + # NOTE(woosuk): `num_tokens` here may include spec tokens. + self.input_batch.num_tokens[req_index] += num_spec_tokens + + # Add the new or resumed requests to the persistent batch. + # The smaller empty indices are filled first. + for req_id in req_ids_to_add: + req_state = self.requests[req_id] + self.input_batch.add_request(req_state) + + # Condense the batched states if there are gaps left by removed requests + self.input_batch.condense() + # Allow attention backend to reorder the batch, potentially + self._may_reorder_batch(scheduler_output) + # Refresh batch metadata with any pending updates. + self.input_batch.refresh_metadata() + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: "SchedulerOutput", + intermediate_tensors: Optional[IntermediateTensors] = None, + ) -> Union[OmniModelRunnerOutput, IntermediateTensors]: + self._update_states(scheduler_output) + if not scheduler_output.total_num_scheduled_tokens: + if not has_kv_transfer_group(): + # Return empty ModelRunnerOutput if there's no work to do. + return EMPTY_MODEL_RUNNER_OUTPUT + + return self.kv_connector_no_forward(scheduler_output, + self.vllm_config) + + # Prepare the decoder inputs. + (attn_metadata, attention_cuda_graphs, logits_indices, + spec_decode_metadata, num_scheduled_tokens_np, + spec_decode_common_attn_metadata) = ( + self._prepare_inputs(scheduler_output)) + + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + if (self.use_cuda_graph + and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]): + # Use piecewise CUDA graphs. + # Add padding to the batch size. + num_input_tokens = self.vllm_config.pad_for_cudagraph( + num_scheduled_tokens) + else: + # Eager mode. + # Pad tokens to multiple of tensor_parallel_size when + # enabled collective fusion for SP + tp_size = self.vllm_config.parallel_config.tensor_parallel_size + if self.compilation_config.pass_config. \ + enable_sequence_parallelism and tp_size > 1: + num_input_tokens = round_up(num_scheduled_tokens, tp_size) + else: + num_input_tokens = num_scheduled_tokens + + # Padding for DP + num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens) + num_input_tokens += num_pad + + # _prepare_inputs may reorder the batch, so we must gather multi + # modal outputs after that to ensure the correct order + if self.is_multimodal_model: + # Run the multimodal encoder if any. + self._execute_mm_encoder(scheduler_output) + mm_embeds = self._gather_mm_embeddings(scheduler_output) + else: + mm_embeds = [] + + if get_pp_group().is_first_rank: + # Always prepare inputs_embeds on the first PP rank so both + # multimodal and text-only models can consume embeddings. + inputs_embeds_scheduled = self.model.get_input_embeddings( + input_ids=self.input_ids[:num_scheduled_tokens], + multimodal_embeddings=(mm_embeds or None) + if self.is_multimodal_model else None, + ) + + # TODO(woosuk): Avoid the copy. Optimize. + self.inputs_embeds[:num_scheduled_tokens].copy_( + inputs_embeds_scheduled) + + # Overlay per-request custom prompt_embeds for the prompt portion + for req_index, req_id in enumerate(self.input_batch.req_ids): + req_state = self.requests[req_id] + pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) + if pe_cpu is None: + continue + num_computed_tokens = int( + self.input_batch.num_computed_tokens_cpu[req_index]) + prompt_len = len(req_state.prompt_token_ids) + prompt_remaining = max(0, prompt_len - num_computed_tokens) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + overlay_len = min(sched_tokens, prompt_remaining) + if overlay_len <= 0: + continue + src = pe_cpu[num_computed_tokens: + num_computed_tokens + overlay_len].to( + dtype=self.dtype, + device=self.device, + non_blocking=True) + start_offset = int(self.query_start_loc_cpu[req_index]) + self.inputs_embeds[start_offset:start_offset + overlay_len]\ + .copy_(src) + + input_ids = self.input_ids[:num_input_tokens] # preserved for APIs + inputs_embeds = self.inputs_embeds[:num_input_tokens] + model_mm_kwargs = (self._extract_mm_kwargs(scheduler_output) + if self.is_multimodal_model else {}) + else: + # For non-first PP ranks, use token ids as usual; embeddings are + # only consumed on the first rank. + input_ids = self.input_ids[:num_input_tokens] + inputs_embeds = None + model_mm_kwargs = {} + if self.uses_mrope: + positions = self.mrope_positions[:, :num_input_tokens] + else: + positions = self.positions[:num_input_tokens] + + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + intermediate_tensors = self.sync_and_slice_intermediate_tensors( + num_input_tokens, intermediate_tensors, True) + + # Some attention backends only support CUDA Graphs in pure decode. + # If attention doesn't support CUDA Graphs for this batch, but we + # compiled with full CUDA graphs, we have to skip them entirely. + skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs + + # Run the model. + # Use persistent buffers for CUDA graphs. + with set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_input_tokens, + num_tokens_across_dp=num_tokens_across_dp, + skip_cuda_graphs=skip_cuda_graphs, + ), self.maybe_get_kv_connector_output( + scheduler_output) as kv_connector_output: + + model_output = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **MultiModalKwargs.as_kwargs( + model_mm_kwargs, + device=self.device, + ), + sampling_metadata=self.input_batch.sampling_metadata, + logits_index=logits_indices, + sampler=self.sampler, + ) + + if self.use_aux_hidden_state_outputs: + hidden_states, aux_hidden_states = model_output + else: + hidden_states = model_output + aux_hidden_states = None + + text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) + + # Broadcast PP output for external_launcher (torchrun) + # to make sure we are synced across pp ranks + # TODO: Support overlapping mirco-batches + # https://github.com/vllm-project/vllm/issues/18019 + broadcast_pp_output = \ + self.parallel_config.distributed_executor_backend \ + == "external_launcher" and len(get_pp_group().ranks) > 0 + if not get_pp_group().is_last_rank: + # For mid-pipeline stages, return the hidden states. + assert isinstance(text_hidden_states, IntermediateTensors) + if not broadcast_pp_output: + text_hidden_states.kv_connector_output = kv_connector_output + return text_hidden_states + get_pp_group().send_tensor_dict(text_hidden_states.tensors, + all_gather_group=get_tp_group()) + logits = None + else: + if self.input_batch.pooling_params: + return self._pool(text_hidden_states, num_scheduled_tokens, + num_scheduled_tokens_np, kv_connector_output) + + sample_hidden_states = text_hidden_states[logits_indices] + logits = self.model.compute_logits(sample_hidden_states, None) + if broadcast_pp_output: + model_output_broadcast_data = { + "logits": logits.contiguous(), + } if logits is not None else {} + model_output_broadcast_data = get_pp_group().broadcast_tensor_dict( + model_output_broadcast_data, src=len(get_pp_group().ranks) - 1) + assert model_output_broadcast_data is not None + logits = model_output_broadcast_data["logits"] + + # Apply structured output bitmasks if present + if scheduler_output.grammar_bitmask is not None: + self.apply_grammar_bitmask(scheduler_output, logits) + + # Sample the next token and get logprobs if needed. + sampling_metadata = self.input_batch.sampling_metadata + if spec_decode_metadata is None: + sampler_output = self.sampler( + logits=logits, + sampling_metadata=sampling_metadata, + ) + else: + # When indexing with a tensor (bonus_logits_indices), PyTorch + # creates a new tensor with separate storage from the original + # logits tensor. This means any in-place operations on bonus_logits + # won't affect the original logits tensor. + assert logits is not None + bonus_logits = logits[spec_decode_metadata.bonus_logits_indices] + sampler_output = self.sampler( + logits=bonus_logits, + sampling_metadata=sampling_metadata, + ) + bonus_token_ids = sampler_output.sampled_token_ids + + # Just like `bonus_logits`, `target_logits` is a new tensor with + # separate storage from the original `logits` tensor. Therefore, + # it is safe to update `target_logits` in place. + target_logits = logits[spec_decode_metadata.target_logits_indices] + output_token_ids = self.rejection_sampler( + spec_decode_metadata, + None, # draft_probs + target_logits, + bonus_token_ids, + sampling_metadata, + ) + sampler_output.sampled_token_ids = output_token_ids + + num_nans_in_logits = {} + if envs.VLLM_COMPUTE_NANS_IN_LOGITS: + num_nans_in_logits = self._get_nans_in_logits(logits) + + # TODO(woosuk): The following loop can be slow since it iterates over + # the requests one by one. Optimize. + discard_sampled_tokens_req_indices = [] + for i, req_id in enumerate(self.input_batch.req_ids): + req_state = self.requests[req_id] + seq_len = (req_state.num_computed_tokens + + scheduler_output.num_scheduled_tokens[req_id]) + if seq_len < req_state.num_tokens: + # Ignore the sampled token for partial prefills. + # Rewind the generator state as if the token was not sampled. + # This relies on cuda-specific torch-internal impl details + generator = self.input_batch.generators.get(i) + if generator is not None: + generator.set_offset(generator.get_offset() - 4) + # Record the index of the request that should not be sampled, + # so that we could clear the sampled tokens before returning. + discard_sampled_tokens_req_indices.append(i) + + # NOTE: GPU -> CPU Sync happens here. + # Move as many CPU operations as possible before this sync point. + logprobs_tensors = sampler_output.logprobs_tensors + logprobs_lists = logprobs_tensors.tolists() \ + if logprobs_tensors is not None else None + + # Compute prompt logprobs if needed. + prompt_logprobs_dict = self._get_prompt_logprobs_dict( + text_hidden_states[:num_scheduled_tokens], + scheduler_output, + ) + + # Get the valid generated tokens. + import os + sampled_token_ids = sampler_output.sampled_token_ids + max_gen_len = sampled_token_ids.shape[-1] + if max_gen_len == 1: + # No spec decode tokens. + valid_sampled_token_ids = sampled_token_ids.tolist() + else: + # Includes spec decode tokens. + valid_sampled_token_ids = self.rejection_sampler.parse_output( + sampled_token_ids, + self.input_batch.vocab_size, + ) + # Mask out the sampled tokens that should not be sampled. + for i in discard_sampled_tokens_req_indices: + valid_sampled_token_ids[i].clear() + + # Cache the sampled tokens in the model runner, so that the scheduler + # doesn't need to send them back. + # NOTE(woosuk): As an exception, when using PP, the scheduler sends + # the sampled tokens back, because there's no direct communication + # between the first-stage worker and the last-stage worker. + for req_idx, sampled_ids in enumerate(valid_sampled_token_ids): + if not sampled_ids: + continue + + start_idx = self.input_batch.num_tokens_no_spec[req_idx] + end_idx = start_idx + len(sampled_ids) + assert end_idx <= self.max_model_len, ( + "Sampled token IDs exceed the max model length. " + f"Total number of tokens: {end_idx} > max_model_len: " + f"{self.max_model_len}") + + self.input_batch.token_ids_cpu[req_idx, + start_idx:end_idx] = sampled_ids + self.input_batch.num_tokens_no_spec[req_idx] = end_idx + self.input_batch.num_tokens[req_idx] = end_idx + req_id = self.input_batch.req_ids[req_idx] + req_state = self.requests[req_id] + req_state.output_token_ids.extend(sampled_ids) + + if not self.speculative_config: + # Speculative decoding is not enabled. + spec_token_ids = None + else: + assert spec_decode_common_attn_metadata is not None + spec_token_ids = self.propose_draft_token_ids( + scheduler_output, + valid_sampled_token_ids, + sampling_metadata, + text_hidden_states, + sample_hidden_states, + aux_hidden_states, + spec_decode_metadata, + spec_decode_common_attn_metadata, + ) + + self.eplb_step() + + return OmniModelRunnerOutput( + req_ids=self.input_batch.req_ids, + req_id_to_index=self.input_batch.req_id_to_index, + sampled_token_ids=valid_sampled_token_ids, + spec_token_ids=spec_token_ids, + logprobs=logprobs_lists, + prompt_logprobs_dict=prompt_logprobs_dict, + pooler_output=[], + kv_connector_output=kv_connector_output, + num_nans_in_logits=num_nans_in_logits, + multimodal_outputs=multimodal_outputs, + ) + + @torch.inference_mode() + def extract_multimodal_outputs(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, dict]: + if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs: + text_hidden_states = hidden_states.text_hidden_states + multimodal_outputs = hidden_states.multimodal_outputs + + else: + text_hidden_states = hidden_states + multimodal_outputs = {} + return text_hidden_states, multimodal_outputs + + @torch.inference_mode() + def _dummy_run( + self, + num_tokens: int, + capture_attn_cudagraph: bool = False, + skip_eplb: bool = False, + is_profile: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + + # Padding for DP + num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) + num_tokens += num_pad + + # Set num_scheduled_tokens based on num_tokens and max_num_seqs + # for dummy run with LoRA so that the num_reqs collectively + # has num_tokens in total. + assert num_tokens <= self.scheduler_config.max_num_batched_tokens + max_num_reqs = self.scheduler_config.max_num_seqs + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs + num_scheduled_tokens_list[-1] += num_tokens % num_reqs + assert sum(num_scheduled_tokens_list) == num_tokens + assert len(num_scheduled_tokens_list) == num_reqs + num_scheduled_tokens = np.array(num_scheduled_tokens_list, + dtype=np.int32) + + attn_metadata: Optional[dict[str, Any]] = None + if capture_attn_cudagraph: + attn_metadata = {} + + # Make sure max_model_len is used at the graph capture time. + self.seq_lens_np[:num_reqs] = self.max_model_len + self.seq_lens_np[num_reqs:] = 0 + self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs], + non_blocking=True) + + for kv_cache_group_id, kv_cache_group_spec in enumerate( + self.kv_cache_config.kv_cache_groups): + common_attn_metadata = CommonAttentionMetadata( + query_start_loc=self.query_start_loc[:num_reqs + 1], + query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + + 1], + seq_lens=self.seq_lens[:num_reqs], + seq_lens_cpu=self.seq_lens_cpu[:num_reqs], + num_computed_tokens_cpu=self.input_batch. + num_computed_tokens_cpu_tensor[:num_reqs], + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + max_query_len=num_tokens, + block_table_tensor=self.input_batch.block_table[ + kv_cache_group_id].get_device_tensor()[:num_reqs], + slot_mapping=self.input_batch. + block_table[kv_cache_group_id].slot_mapping[:num_tokens], + causal=True) + + for attn_group in self.attn_groups[kv_cache_group_id]: + attn_metadata_i = attn_group.metadata_builder\ + .build_for_cudagraph_capture(common_attn_metadata) + for layer_name in kv_cache_group_spec.layer_names: + attn_metadata[layer_name] = attn_metadata_i + + with self.maybe_dummy_run_with_lora(self.lora_config, + num_scheduled_tokens): + if self.is_multimodal_model: + input_ids = None + inputs_embeds = self.inputs_embeds[:num_tokens] + model_mm_kwargs = self._dummy_mm_kwargs(num_reqs) + else: + input_ids = self.input_ids[:num_tokens] + inputs_embeds = None + model_mm_kwargs = {} + + if self.uses_mrope: + positions = self.mrope_positions[:, :num_tokens] + else: + positions = self.positions[:num_tokens] + + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + if self.intermediate_tensors is None: + self.intermediate_tensors = ( + self.model.make_empty_intermediate_tensors( + batch_size=self.max_num_tokens, + dtype=self.model_config.dtype, + device=self.device)) + + intermediate_tensors = self.sync_and_slice_intermediate_tensors( + num_tokens, None, False) + + with self.maybe_randomize_inputs(input_ids), set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens, + num_tokens_across_dp=num_tokens_across_dp): + outputs = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **MultiModalKwargs.as_kwargs( + model_mm_kwargs, + device=self.device, + ), + sampler = None + ) + + if self.use_aux_hidden_state_outputs: + hidden_states, _ = outputs + else: + hidden_states = outputs + + text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) + + + if self.speculative_config and self.speculative_config.use_eagle(): + assert isinstance(self.drafter, EagleProposer) + self.drafter.dummy_run(num_tokens) + + # This is necessary to avoid blocking DP. + # For dummy runs, we typically skip EPLB since we don't have any real + # requests to process. + # However, in DP settings, there may be cases when some DP ranks do + # not have any requests to process, so they're executing dummy batches. + # In such cases, we still have to trigger EPLB to make sure + # ranks execute the rearrangement in synchronization. + if not skip_eplb: + self.eplb_step(is_dummy=True, is_profile=is_profile) + + logit_indices = np.cumsum(num_scheduled_tokens) - 1 + return text_hidden_states, text_hidden_states[logit_indices] \ No newline at end of file diff --git a/vllm_omni/worker/model_runner.py b/vllm_omni/worker/model_runner.py new file mode 100644 index 000000000..ca7b78d98 --- /dev/null +++ b/vllm_omni/worker/model_runner.py @@ -0,0 +1,199 @@ +from dataclasses import dataclass +from typing import Optional, Dict, Any, List +import itertools + +import torch + +from vllm.worker.model_runner import ModelInputForGPU, ModelInputForGPUBuilder + +from vllm.utils import flatten_2d_lists, async_tensor_h2d +from vllm.lora.layers import LoRAMapping +from vllm.sequence import MultiModalKwargs + + +@dataclass(frozen=True) +class OmniModelInputForGPU(ModelInputForGPU): + # New: optional additional information dict for current scheduled tokens + additional_information: Optional[Dict[str, Any]] = None + + +class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder): + def build(self) -> OmniModelInputForGPU: + """Finalize the builder intermediate data and + create on-device tensors. + """ + # Combine and flatten intermediate data. + input_tokens = list[int]() + inputs_embeds_list = list[torch.Tensor]() + token_types = list[int]() + for inter_data in self.inter_data_list: + for cur_input_tokens in inter_data.input_tokens: + input_tokens.extend(cur_input_tokens) + for cur_token_types in inter_data.token_types: + token_types.extend(cur_token_types) + if inter_data.inputs_embeds is not None: + inputs_embeds_list.append( + inter_data.inputs_embeds.to( + dtype=self.runner.model_config.dtype, + device=self.runner.device)) + # Support v1 direct-transfer prompt embeds attached on request state + if inter_data.inputs_embeds is None: + try: + # Locate req_id -> request state to fetch decoded CPU embeds + req_id = inter_data.request_id # type: ignore[attr-defined] + req_state = getattr(self.runner, "requests", {}).get(req_id) + pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) + if pe_cpu is not None: + inter_data.inputs_embeds = pe_cpu.to( + dtype=self.runner.model_config.dtype, + device=self.runner.device) + inputs_embeds_list.append(inter_data.inputs_embeds) + except Exception: + pass + inputs_embeds: Optional[torch.Tensor] + if len(inputs_embeds_list) == 0: + inputs_embeds = None + else: + inputs_embeds = torch.cat(inputs_embeds_list, dim=0).to( + dtype=self.runner.model_config.dtype, + device=self.runner.device) + assert len(inputs_embeds) == len(input_tokens) + + if not input_tokens and inputs_embeds is None: + # This may happen when all prefill requests hit + # prefix caching and there is no decode request. + return self.model_input_cls() + + mrope_input_positions: Optional[List[List[int]]] = None + if any(inter_data.mrope_input_positions is not None + for inter_data in self.inter_data_list): + mrope_input_positions = [[] for _ in range(3)] + for idx in range(3): + for inter_data in self.inter_data_list: + msections = inter_data.mrope_input_positions + if msections is None: + for _seq_input_positions in inter_data.input_positions: + mrope_input_positions[idx].extend( + _seq_input_positions) + else: + for _seq_mrope_input_positions in msections: + mrope_input_positions[idx].extend( + _seq_mrope_input_positions[idx]) + input_positions = None + else: + input_positions = [] + for inter_data in self.inter_data_list: + for cur_input_positions in inter_data.input_positions: + input_positions.extend(cur_input_positions) + + seq_lens = [] + query_lens = [] + max_decode_seq_len = 0 + max_encoder_seq_len = 0 + for inter_data in self.inter_data_list: + seq_lens.extend(inter_data.seq_lens) + query_lens.extend(inter_data.query_lens) + if not inter_data.is_prompt: + max_decode_seq_len = max(max_decode_seq_len, + max(inter_data.seq_lens)) + if self.runner.model_config.is_encoder_decoder: + max_encoder_seq_len = max(max_encoder_seq_len, + inter_data.encoder_seq_len) + + # Mapping from request IDs to sequence IDs. Used for Jamba models + # that manages the cache by itself. + request_ids_to_seq_ids = { + data.request_id: data.seq_ids + for data in self.inter_data_list + } + + cuda_graph_pad_size = self._get_cuda_graph_pad_size( + num_seqs=len(seq_lens), + max_decode_seq_len=max_decode_seq_len, + max_encoder_seq_len=max_encoder_seq_len) + + batch_size = len(input_tokens) + if cuda_graph_pad_size != -1: + # If cuda graph can be used, pad tensors accordingly. + # See `capture_model` API for more details. + # vLLM uses cuda graph only for decoding requests. + batch_size += cuda_graph_pad_size + + # Tokens and positions. + if cuda_graph_pad_size: + input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size)) + assert self.runner.device is not None + input_tokens_tensor = async_tensor_h2d(input_tokens, torch.long, + self.runner.device, + self.runner.pin_memory) + + token_types_tensor = async_tensor_h2d(token_types, torch.long, + self.runner.device, + self.runner.pin_memory) \ + if token_types else None + + if mrope_input_positions is not None: + for idx in range(3): + mrope_input_positions[idx].extend( + itertools.repeat(0, cuda_graph_pad_size)) + input_positions_tensor = async_tensor_h2d(mrope_input_positions, + torch.long, + self.runner.device, + self.runner.pin_memory) + else: + input_positions.extend(itertools.repeat(0, cuda_graph_pad_size)) + input_positions_tensor = async_tensor_h2d(input_positions, + torch.long, + self.runner.device, + self.runner.pin_memory) + # Sequence and query lengths. + if cuda_graph_pad_size: + seq_lens.extend(itertools.repeat(1, cuda_graph_pad_size)) + + # Attention metadata. + attn_metadata = self.attn_metadata_builder.build( + seq_lens, query_lens, cuda_graph_pad_size, batch_size) + + # LoRA data. + lora_requests = set() + lora_mapping = None + if self.enable_lora: + lora_requests = set(r for data in self.inter_data_list + for r in data.lora_requests) + lora_index_mapping = flatten_2d_lists([ + flatten_2d_lists(inter_data.lora_index_mapping) + for inter_data in self.inter_data_list + ]) + if cuda_graph_pad_size: + lora_index_mapping.extend( + itertools.repeat(0, cuda_graph_pad_size)) + lora_prompt_mapping = flatten_2d_lists([ + flatten_2d_lists(inter_data.lora_prompt_mapping) + for inter_data in self.inter_data_list + ]) + + lora_mapping = LoRAMapping( + **dict(index_mapping=lora_index_mapping, + prompt_mapping=lora_prompt_mapping, + is_prefill=not self.decode_only)) + + # Multi-modal data. + multi_modal_kwargs_list = [ + data.multi_modal_kwargs for data in self.inter_data_list + if data.multi_modal_kwargs is not None + ] + multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list) + + return self.model_input_cls( + input_tokens=input_tokens_tensor, + inputs_embeds=inputs_embeds, + input_positions=input_positions_tensor, + token_types=token_types_tensor, + attn_metadata=attn_metadata, + seq_lens=seq_lens, + query_lens=query_lens, + lora_mapping=lora_mapping, + lora_requests=lora_requests, + multi_modal_kwargs=multi_modal_kwargs, + request_ids_to_seq_ids=request_ids_to_seq_ids, + finished_requests_ids=self.finished_requests_ids) \ No newline at end of file