diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 001c46645..7c8c4a89e 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -54,7 +54,6 @@ steps: - docker#v5.2.0: image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT always-pull: true - shm-size: "8gb" propagate-environment: true shm-size: "8gb" environment: @@ -62,38 +61,13 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test" - timeout_in_minutes: 30 + - label: ":full_moon: Diffusion Model Test with H100" + timeout_in_minutes: 60 depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" - commands: - - | - huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk - agents: - queue: "gpu_4_queue" - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - - label: ":full_moon: Omni Model Perf Test with H100" - key: nightly-performance - timeout_in_minutes: 180 - depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" + # if: build.env("NIGHTLY") == "1" commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export BENCHMARK_DIR=tests - - pytest -s -v tests/perf/scripts/run_benchmark.py - - buildkite-agent artifact upload "tests/*.json" + - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model" agents: queue: "mithril-h100-pool" plugins: @@ -123,17 +97,82 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":email: Nightly Perf Collection & Email" + - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test" + timeout_in_minutes: 30 + depends_on: upload-nightly-pipeline + if: build.env("NIGHTLY") == "1" + commands: + - | + huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk + agents: + queue: "gpu_4_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + + - label: ":full_moon: Omni Model Perf Test & Test Case Statistics" + key: nightly-performance + timeout_in_minutes: 180 + depends_on: upload-nightly-pipeline + if: build.env("NIGHTLY") == "1" + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - export BENCHMARK_DIR=tests + - pytest -s -v tests/perf/scripts/run_benchmark.py + - buildkite-agent artifact upload "tests/*.json" + - python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html + - buildkite-agent artifact upload "tests/*.html" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":email: Nightly Collection & Email" key: nightly-perf-distribution - depends_on: nightly-performance + depends_on: + - nightly-performance if: build.env("NIGHTLY") == "1" commands: - pip install openpyxl - export DEFAULT_INPUT_DIR=tests - export DEFAULT_OUTPUT_DIR=tests - buildkite-agent artifact download "tests/*.json" . --step nightly-performance + - buildkite-agent artifact download "tests/*.html" . --step nightly-performance - python tools/nightly/generate_nightly_perf_excel.py - - python tools/nightly/send_nightly_perf_email.py + - python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html" - buildkite-agent artifact upload "tests/*.xlsx" agents: queue: "cpu_queue_premerge" diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 092d6f11a..66598f3ea 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -34,15 +34,7 @@ steps: - label: "Diffusion Model CPU offloading Test" depends_on: upload-ready-pipeline commands: - - | - timeout 20m bash -c ' - set +e - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py - EXIT1=$$? - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py - EXIT2=$$? - exit $$((EXIT1 | EXIT2)) - ' + - timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py agents: queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU plugins: diff --git a/README.md b/README.md index df5d093b1..81d24a1a3 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone

-| Documentation | User Forum | Developer Slack | WeChat | Paper | Slides | +| Documentation | User Forum | Developer Slack | WeChat | Paper | Slides |

@@ -17,6 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone *Latest News* πŸ”₯ +- [2026/03] Joined vLLM HK First Meetup! Video on [youtube](https://youtu.be/sgwNfsNnR9I) - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more. - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**β€”while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation. - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results. diff --git a/apps/ComfyUI-vLLM-Omni/README.md b/apps/ComfyUI-vLLM-Omni/README.md index 23bbe5b72..54f2fdf2e 100644 --- a/apps/ComfyUI-vLLM-Omni/README.md +++ b/apps/ComfyUI-vLLM-Omni/README.md @@ -47,6 +47,7 @@ If no, check your shell running the ComfyUI process. There may be some error mes This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**): - **Generate Image** for text-to-image and image-to-image tasks +- **Generate Video** for text-to-video and image-to-video tasks - **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks - **TTS** and **TTS Voice Clone** for TTS tasks @@ -59,7 +60,7 @@ To build a simple workflow yourself, - Drag a generation node onto the canvas. - Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio** -- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node. +- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**. - If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**. - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node. - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters. @@ -104,6 +105,20 @@ You can configure per-stage sampling parameters for multi-stage models. > [!TIP] > The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not. +### Text-to-video and image-to-video generation (e.g., Wan) + +(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**) + +

+ + + vLLM-Omni video generation + +

+ +> [!TIP] +> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not. + ### TTS (e.g., Qwen TTS series) (Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**) diff --git a/apps/ComfyUI-vLLM-Omni/__init__.py b/apps/ComfyUI-vLLM-Omni/__init__.py index 641824add..e89b2ca1e 100644 --- a/apps/ComfyUI-vLLM-Omni/__init__.py +++ b/apps/ComfyUI-vLLM-Omni/__init__.py @@ -14,17 +14,21 @@ from .comfyui_vllm_omni.nodes import ( VLLMOmniARSampling, VLLMOmniDiffusionSampling, VLLMOmniGenerateImage, + VLLMOmniGenerateVideo, VLLMOmniQwenTTSParams, + VLLMOmniRemoteLoRA, VLLMOmniSamplingParamsList, VLLMOmniTTS, VLLMOmniUnderstanding, VLLMOmniVoiceClone, + VLLMOmniWanParams, ) # A dictionary that contains all nodes you want to export with their names NODE_CLASS_MAPPINGS = { # === Generation === "VLLMOmniGenerateImage": VLLMOmniGenerateImage, + "VLLMOmniGenerateVideo": VLLMOmniGenerateVideo, "VLLMOmniUnderstanding": VLLMOmniUnderstanding, "VLLMOmniTTS": VLLMOmniTTS, "VLLMOmniVoiceClone": VLLMOmniVoiceClone, @@ -32,13 +36,16 @@ NODE_CLASS_MAPPINGS = { "VLLMOmniARSampling": VLLMOmniARSampling, "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling, "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList, + "VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA, "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams, + "VLLMOmniWanParams": VLLMOmniWanParams, } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { # === Generation === "VLLMOmniGenerateImage": "Generate Image", + "VLLMOmniGenerateVideo": "Generate Video", "VLLMOmniUnderstanding": "Multimodality Understanding", "VLLMOmniTTS": "TTS (Text to Speech)", "VLLMOmniVoiceClone": "TTS Voice Cloning", @@ -46,7 +53,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "VLLMOmniARSampling": "AR Sampling Params", "VLLMOmniDiffusionSampling": "Diffusion Sampling Params", "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List", + "VLLMOmniRemoteLoRA": "LoRA", "VLLMOmniQwenTTSParams": "Qwen TTS Params", + "VLLMOmniWanParams": "Wan Video Params", } WEB_DIRECTORY = "./web" diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py index 5caa3869e..bfea93998 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py @@ -11,6 +11,7 @@ from .utils.types import ( AutoregressionSamplingParams, DiffusionSamplingParams, QwenTTSModelSpecificParams, + WanModelSpecificParams, ) from .utils.validators import ( add_sampling_parameters_to_stage, @@ -56,6 +57,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase): # "video": ("VIDEO",), # "audio": ("AUDIO",), "sampling_params": ("SAMPLING_PARAMS",), + "lora": ("REMOTE_LORA",), }, } @@ -76,9 +78,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase): audio: AudioInput | None = None, # Hidden & unused video: VideoInput | None = None, # Hidden & unused sampling_params: dict | list[dict] | None = None, + lora: dict | None = None, **kwargs, ): - logger.info("Uncaught kwargs: %s", kwargs) + if kwargs: + logger.info("Uncaught kwargs: %s", kwargs) logger.debug("Got sampling params: %s", sampling_params) validate_model_and_sampling_params_types(model, sampling_params) if image is None and mask is not None: @@ -105,6 +109,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase): height=height, negative_prompt=negative_prompt, sampling_params=sampling_params, + lora=lora, ) return (output,) elif image is not None and audio is None and video is None: @@ -119,6 +124,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase): negative_prompt=negative_prompt, mask=mask, sampling_params=sampling_params, + lora=lora, ) return (output,) @@ -136,11 +142,91 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase): audio=audio, video=video, sampling_params=sampling_params, + lora=lora, ) return (output,) +class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ("STRING", {"default": "http://localhost:8000/v1"}), + "model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}), + "prompt": ("STRING", {"multiline": True}), + "negative_prompt": ("STRING", {"multiline": True, "default": ""}), + "width": ("INT", {"default": 832, "min": 1}), + "height": ("INT", {"default": 480, "min": 1}), + "fps": ("INT", {"default": 16, "min": 1}), + "num_frames": ("INT", {"default": 41, "min": 1}), + }, + "optional": { + "image": ("IMAGE",), + "sampling_params": ("SAMPLING_PARAMS",), + "lora": ("REMOTE_LORA",), + "model_params": ("VIDEO_PARAMS",), + }, + } + + RETURN_TYPES = ("VIDEO",) + RETURN_NAMES = ("video",) + FUNCTION = "generate" + + async def generate( + self, + url: str, + model: str, + prompt: str, + width: int, + height: int, + fps: int, + num_frames: int, + negative_prompt: str | None = None, + image: torch.Tensor | None = None, + sampling_params: dict | list[dict] | None = None, + model_params: dict | None = None, + lora: dict | None = None, + **kwargs, + ): + if kwargs: + logger.info("Uncaught kwargs: %s", kwargs) + logger.debug("Got sampling params: %s", sampling_params) + logger.debug("Got model params: %s", model_params) + validate_model_and_sampling_params_types(model, sampling_params) + + # Currently, all video generation models are single-stage diffusion models + if isinstance(sampling_params, list): + if len(sampling_params) != 1: + raise ValueError( + "Video generation expects a single sampling params group. " + "Please provide one Diffusion sampling node." + ) + sampling_params = sampling_params[0] + + if sampling_params is not None: + sampling_params.pop("type", None) # internal fields + if model_params is not None: + model_params.pop("type", None) # internal fields + + client = VLLMOmniClient(url) + output = await client.generate_video( + model=model, + prompt=prompt, + image=image, # image present => i2v, absent => t2v + width=width, + height=height, + num_frames=num_frames, + fps=fps, + negative_prompt=negative_prompt, + sampling_params=sampling_params, + lora=lora, + model_params=model_params, + ) + return (output,) + + class VLLMOmniUnderstanding(_VLLMOmniGenerateBase): @classmethod def INPUT_TYPES(cls): @@ -188,7 +274,8 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase): use_audio_in_video: bool = True, **kwargs, ) -> tuple[str, AudioInput]: - logger.info("Uncaught kwargs: %s", kwargs) + if kwargs: + logger.info("Uncaught kwargs: %s", kwargs) logger.debug("Got sampling params: %s", sampling_params) validate_model_and_sampling_params_types(model, sampling_params) @@ -547,6 +634,52 @@ class VLLMOmniSamplingParamsList: return (params,) +class VLLMOmniRemoteLoRA: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "local_path": ("STRING", {"default": ""}), + "name": ("STRING", {"default": ""}), + "scale": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1}, + ), + "int_id": ( + "INT", + { + "default": 0, + "min": 0, + "step": 1, + "tooltip": "0 means it is not set and the server can derive it.", + }, + ), + } + } + + RETURN_TYPES = ("REMOTE_LORA",) + RETURN_NAMES = ("lora",) + FUNCTION = "get_lora" + CATEGORY = "vLLM-Omni" + + @classmethod + def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]: + if not local_path.strip() or not name.strip(): + return "Both local_path and name must be provided." + return True + + def get_lora(self, local_path: str, name: str, scale: float, int_id: int): + local_path = local_path.strip() + name = name.strip() + lora = { + "local_path": local_path or None, + "name": name or None, + "scale": float(scale), + "int_id": int(int_id) if int_id > 0 else None, + } + return (lora,) + + class VLLMOmniQwenTTSParams: @classmethod def INPUT_TYPES(cls): @@ -572,3 +705,32 @@ class VLLMOmniQwenTTSParams: def get_params(self, **kwargs): return (QwenTTSModelSpecificParams(kwargs),) + + +class VLLMOmniWanParams: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "guidance_scale_2": ( + "FLOAT", + {"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1}, + ), + "boundary_ratio": ( + "FLOAT", + {"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001}, + ), + "flow_shift": ( + "FLOAT", + {"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1}, + ), + } + } + + RETURN_TYPES = ("VIDEO_PARAMS",) + RETURN_NAMES = ("Wan video params",) + FUNCTION = "get_params" + CATEGORY = "vLLM-Omni/Video Params" + + def get_params(self, **kwargs): + return (WanModelSpecificParams(kwargs),) diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py index d340073bc..2c936d33f 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py @@ -6,6 +6,7 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License. """ +import json from typing import Any import aiohttp @@ -17,6 +18,7 @@ from .format import ( audio_to_base64, base64_to_audio, base64_to_image_tensor, + base64_to_video, bytes_to_audio, image_tensor_to_base64, image_tensor_to_png_bytes, @@ -30,7 +32,7 @@ logger = get_logger(__name__) class VLLMOmniClient: - def __init__(self, base_url: str, timeout: float = 300.0): + def __init__(self, base_url: str, timeout: float | None = None): self.base_url = base_url self.timeout = aiohttp.ClientTimeout(total=timeout) @@ -43,12 +45,13 @@ class VLLMOmniClient: height: int, negative_prompt: str | None = None, sampling_params: dict | None = None, + lora: dict | None = None, ) -> torch.Tensor: """Run text-to-image generation via DALLE API""" await self._check_model_exist(model) size = f"{width}x{height}" - payload = { + payload: dict[str, Any] = { "model": model, "prompt": prompt, "size": size, @@ -58,6 +61,8 @@ class VLLMOmniClient: payload["negative_prompt"] = negative_prompt if sampling_params is not None: payload.update(sampling_params) + if lora is not None: + payload["lora"] = lora logger.debug("img gen payload: %s", payload) url = self.base_url + "/images/generations" @@ -110,6 +115,7 @@ class VLLMOmniClient: negative_prompt: str | None = None, mask: torch.Tensor | None = None, sampling_params: dict | None = None, + lora: dict | None = None, ) -> torch.Tensor: """Run image editing via DALLE API""" await self._check_model_exist(model) @@ -131,6 +137,8 @@ class VLLMOmniClient: if sampling_params is not None: for k, v in sampling_params.items(): form.add_field(k, str(v)) + if lora is not None: + form.add_field("lora", json.dumps(lora, ensure_ascii=False)) if mask is not None: mask_filename = "mask.png" form.add_field( @@ -183,6 +191,7 @@ class VLLMOmniClient: audio: AudioInput | None = None, video: VideoInput | None = None, sampling_params: dict | list[dict] | None = None, + lora: dict | None = None, ) -> torch.Tensor: payload = VLLMOmniClient._prepare_chat_completion_messages( model=model, @@ -193,6 +202,8 @@ class VLLMOmniClient: video=video, sampling_params=sampling_params, modalities=["image"], + # === below are additional `extra_body` fields, handled by **kwargs === + lora=lora, ) choices = await self._generate_base_chat_completion(model, payload) @@ -206,6 +217,79 @@ class VLLMOmniClient: return torch.stack(image_tensors, dim=0) + async def generate_video( + self, + *, + model: str, + prompt: str, + width: int, + height: int, + num_frames: int, + fps: int, + negative_prompt: str | None = None, + image: torch.Tensor | None = None, + sampling_params: dict | None = None, + model_params: dict | None = None, + lora: dict | None = None, + **extra_body, + ) -> VideoInput: + form = aiohttp.FormData() + form.add_field("model", model) + form.add_field("prompt", prompt) + form.add_field("width", str(width)) + form.add_field("height", str(height)) + form.add_field("num_frames", str(num_frames)) + form.add_field("fps", str(fps)) + if negative_prompt: + form.add_field("negative_prompt", negative_prompt) + if sampling_params is not None: + for k, v in sampling_params.items(): + form.add_field(k, str(v)) + if model_params is not None: + for k, v in model_params.items(): + form.add_field(k, str(v)) + if lora is not None: + form.add_field("lora", json.dumps(lora, ensure_ascii=False)) + if extra_body: + form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False)) + + if image is not None: + image_filename = "image.png" # Required for multipart form + form.add_field( + "input_reference", + image_tensor_to_png_bytes(image, image_filename), + filename=image_filename, + content_type="image/png", + ) + + url = self.base_url + "/videos" + async with aiohttp.ClientSession(timeout=self.timeout) as session: + try: + async with session.post(url, data=form) as response: + if not response.ok: + error_text = await response.text() + raise (ValueError if response.status < 500 else RuntimeError)( + f"vLLM-Omni API returned status {response.status}: {error_text}" + ) + + try: + data = await response.json() + except aiohttp.ContentTypeError as e: + raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}") + except aiohttp.ClientError as e: + raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}") + + if "data" not in data: + raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format") + if not data["data"]: + raise RuntimeError("API returned empty data array") + try: + base64_str = data["data"][0]["b64_json"] + except (KeyError, IndexError): + raise RuntimeError("API response missing 'b64_json' field in first data item") + + return base64_to_video(base64_str) + async def generate_understanding_chat_completion( self, *, diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py index 0ddf86f9c..642c83c8d 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py @@ -6,12 +6,17 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und import base64 import mimetypes +from fractions import Fraction from io import BytesIO import av import numpy as np import torch +from av.audio.frame import AudioFrame +from av.audio.resampler import AudioResampler +from av.video.frame import VideoFrame from comfy_api.input import AudioInput, VideoInput +from comfy_api.latest import InputImpl, Types from comfy_extras import nodes_audio from PIL import Image @@ -138,6 +143,71 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str: return f"data:{mime_type};base64,{base64_str}" +def base64_to_video(base64_str: str) -> VideoInput: + if base64_str.startswith("data:video"): + _, base64_str = base64_str.split(",", 1) + + try: + video_bytes = base64.b64decode(base64_str) + except Exception as e: + raise ValueError(f"Invalid base64 string: {e}") + + video_buffer = BytesIO(video_bytes) + + try: + with av.open(video_buffer, mode="r") as container: + video_stream = next((s for s in container.streams if s.type == "video"), None) + if video_stream is None: + raise ValueError("No video stream found in decoded payload.") + + frames = [] + for frame in container.decode(video_stream): + if not isinstance(frame, VideoFrame): + continue + image = frame.to_ndarray(format="rgb24") + frames.append(torch.from_numpy(image).float() / 255.0) + + if len(frames) == 0: + raise ValueError("No video frames found in decoded payload.") + + images = torch.stack(frames, dim=0) + frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1) + + audio: AudioInput | None = None + if len(container.streams.audio): + audio_stream = container.streams.audio[-1] + audio_frames = [] + resampler = AudioResampler(format="fltp") + for frame in container.decode(audio_stream): + if not isinstance(frame, AudioFrame): + continue + resampled = resampler.resample(frame) + if not isinstance(resampled, list): + resampled = [resampled] + for audio_frame in resampled: + if audio_frame is not None: + audio_frames.append(audio_frame.to_ndarray()) + + if len(audio_frames) > 0: + audio_data = np.concatenate(audio_frames, axis=1) + sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1 + audio = { + "waveform": torch.from_numpy(audio_data).unsqueeze(0), + "sample_rate": sample_rate, + } + + components = Types.VideoComponents( + images=images, + frame_rate=frame_rate, + audio=audio, + metadata=container.metadata if container.metadata else None, + ) + except Exception as e: + raise RuntimeError(f"Failed to decode video: {e}") + + return InputImpl.VideoFromComponents(components) + + def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO: waveform = audio["waveform"][0] # Shape: (C, T) sample_rate = audio["sample_rate"] diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py index 6cc79f3ed..c7d254eb9 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py @@ -23,6 +23,10 @@ class QwenTTSModelSpecificParams(dict): pass +class WanModelSpecificParams(dict): + pass + + class ModelMode(Enum): IMAGE_GENERATION = auto() VIDEO_GENERATION = auto() diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py index 2abdd2100..c5ac4bdf4 100644 --- a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py @@ -16,7 +16,7 @@ def validate_model_and_sampling_params_types( # Skip if no spec or no sampling params pipeline_spec, _ = lookup_model_spec(model_name) if pipeline_spec is None: - logger.info("skipping sampling params check because spec is not found") + logger.info(f"skipping sampling params check because spec for {model_name} is not found") return if sampling_param_list is None: return diff --git a/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg new file mode 100644 index 000000000..a449a4eee Binary files /dev/null and b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg differ diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json new file mode 100644 index 000000000..636360ab8 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Video Generation.json @@ -0,0 +1,513 @@ +{ + "id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21", + "revision": 0, + "last_node_id": 14, + "last_link_id": 28, + "nodes": [ + { + "id": 10, + "type": "MarkdownNote", + "pos": [ + 227.99306819575278, + -231.24143306069843 + ], + "size": [ + 240.3326478922474, + 136.3505820791881 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "title": "Note: Task and Input", + "properties": {}, + "widgets_values": [ + "vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not." + ], + "color": "#432", + "bgcolor": "#000" + }, + { + "id": 4, + "type": "LoadImage", + "pos": [ + 496.31859627609606, + -229.71277089860084 + ], + "size": [ + 270, + 314 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "localized_name": "image", + "name": "image", + "type": "COMBO", + "widget": { + "name": "image" + }, + "link": null + }, + { + "localized_name": "choose file to upload", + "name": "upload", + "type": "IMAGEUPLOAD", + "widget": { + "name": "upload" + }, + "link": null + } + ], + "outputs": [ + { + "localized_name": "IMAGE", + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 25 + ] + }, + { + "localized_name": "MASK", + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "cute-cat.jpg", + "image" + ] + }, + { + "id": 8, + "type": "VLLMOmniDiffusionSampling", + "pos": [ + 465.5140218220927, + 148.80072646828987 + ], + "size": [ + 284.205078125, + 226 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "localized_name": "n", + "name": "n", + "type": "INT", + "widget": { + "name": "n" + }, + "link": null + }, + { + "localized_name": "num_inference_steps", + "name": "num_inference_steps", + "type": "INT", + "widget": { + "name": "num_inference_steps" + }, + "link": null + }, + { + "localized_name": "guidance_scale", + "name": "guidance_scale", + "type": "FLOAT", + "widget": { + "name": "guidance_scale" + }, + "link": null + }, + { + "localized_name": "true_cfg_scale", + "name": "true_cfg_scale", + "type": "FLOAT", + "widget": { + "name": "true_cfg_scale" + }, + "link": null + }, + { + "localized_name": "vae_use_slicing", + "name": "vae_use_slicing", + "type": "BOOLEAN", + "widget": { + "name": "vae_use_slicing" + }, + "link": null + }, + { + "localized_name": "vae_use_tiling", + "name": "vae_use_tiling", + "type": "BOOLEAN", + "widget": { + "name": "vae_use_tiling" + }, + "link": null + }, + { + "localized_name": "seed", + "name": "seed", + "type": "INT", + "widget": { + "name": "seed" + }, + "link": null + } + ], + "outputs": [ + { + "localized_name": "diffusion sampling params", + "name": "diffusion sampling params", + "type": "SAMPLING_PARAMS", + "links": [ + 27 + ] + } + ], + "properties": { + "Node name for S&R": "VLLMOmniDiffusionSampling" + }, + "widgets_values": [ + 1, + 50, + 1, + 1, + false, + false, + 1992, + "randomize" + ] + }, + { + "id": 14, + "type": "VLLMOmniWanParams", + "pos": [ + 518.6005521845747, + 428.4455663719664 + ], + "size": [ + 270, + 106 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "localized_name": "guidance_scale_2", + "name": "guidance_scale_2", + "type": "FLOAT", + "widget": { + "name": "guidance_scale_2" + }, + "link": null + }, + { + "localized_name": "boundary_ratio", + "name": "boundary_ratio", + "type": "FLOAT", + "widget": { + "name": "boundary_ratio" + }, + "link": null + }, + { + "localized_name": "flow_shift", + "name": "flow_shift", + "type": "FLOAT", + "widget": { + "name": "flow_shift" + }, + "link": null + } + ], + "outputs": [ + { + "localized_name": "Wan video params", + "name": "Wan video params", + "type": "VIDEO_PARAMS", + "links": [ + 28 + ] + } + ], + "properties": { + "Node name for S&R": "VLLMOmniWanParams" + }, + "widgets_values": [ + 4, + 0.875, + 5 + ] + }, + { + "id": 12, + "type": "VLLMOmniGenerateVideo", + "pos": [ + 827.2566336087867, + -73.77449831827545 + ], + "size": [ + 400, + 346 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "localized_name": "image", + "name": "image", + "shape": 7, + "type": "IMAGE", + "link": 25 + }, + { + "localized_name": "sampling_params", + "name": "sampling_params", + "shape": 7, + "type": "SAMPLING_PARAMS", + "link": 27 + }, + { + "localized_name": "lora", + "name": "lora", + "shape": 7, + "type": "REMOTE_LORA", + "link": null + }, + { + "localized_name": "model_params", + "name": "model_params", + "shape": 7, + "type": "VIDEO_PARAMS", + "link": 28 + }, + { + "localized_name": "url", + "name": "url", + "type": "STRING", + "widget": { + "name": "url" + }, + "link": null + }, + { + "localized_name": "model", + "name": "model", + "type": "STRING", + "widget": { + "name": "model" + }, + "link": null + }, + { + "localized_name": "prompt", + "name": "prompt", + "type": "STRING", + "widget": { + "name": "prompt" + }, + "link": null + }, + { + "localized_name": "negative_prompt", + "name": "negative_prompt", + "type": "STRING", + "widget": { + "name": "negative_prompt" + }, + "link": null + }, + { + "localized_name": "width", + "name": "width", + "type": "INT", + "widget": { + "name": "width" + }, + "link": null + }, + { + "localized_name": "height", + "name": "height", + "type": "INT", + "widget": { + "name": "height" + }, + "link": null + }, + { + "localized_name": "fps", + "name": "fps", + "type": "INT", + "widget": { + "name": "fps" + }, + "link": null + }, + { + "localized_name": "num_frames", + "name": "num_frames", + "type": "INT", + "widget": { + "name": "num_frames" + }, + "link": null + } + ], + "outputs": [ + { + "localized_name": "video", + "name": "video", + "type": "VIDEO", + "links": [ + 26 + ] + } + ], + "properties": { + "Node name for S&R": "VLLMOmniGenerateVideo" + }, + "widgets_values": [ + "http://localhost:8000/v1", + "Wan-AI/Wan2.2-TI2V-5B-Diffusers", + "Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.", + "Cartoonish", + 900, + 1200, + 12, + 41 + ] + }, + { + "id": 13, + "type": "SaveVideo", + "pos": [ + 1257.5441369501657, + -44.12956394419863 + ], + "size": [ + 374, + 372.34896826553796 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "localized_name": "video", + "name": "video", + "type": "VIDEO", + "link": 26 + }, + { + "localized_name": "filename_prefix", + "name": "filename_prefix", + "type": "STRING", + "widget": { + "name": "filename_prefix" + }, + "link": null + }, + { + "localized_name": "format", + "name": "format", + "type": "COMBO", + "widget": { + "name": "format" + }, + "link": null + }, + { + "localized_name": "codec", + "name": "codec", + "type": "COMBO", + "widget": { + "name": "codec" + }, + "link": null + } + ], + "outputs": [], + "properties": {}, + "widgets_values": [ + "video/ComfyUI", + "auto", + "auto" + ] + } + ], + "links": [ + [ + 25, + 4, + 0, + 12, + 0, + "IMAGE" + ], + [ + 26, + 12, + 0, + 13, + 0, + "VIDEO" + ], + [ + 27, + 8, + 0, + 12, + 1, + "SAMPLING_PARAMS" + ], + [ + 28, + 14, + 0, + 12, + 3, + "VIDEO_PARAMS" + ] + ], + "groups": [ + { + "id": 1, + "title": "Input", + "bounding": [ + 213.1706010087147, + -313.20095750667554, + 560.8246471437027, + 407.4246532472182 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "workflowRendererVersion": "LG", + "ds": { + "scale": 1.1469075819486894, + "offset": [ + 101.7632950847089, + 492.29121889347806 + ] + } + }, + "version": 0.4 +} diff --git a/benchmarks/qwen3-tts/README.md b/benchmarks/qwen3-tts/README.md new file mode 100644 index 000000000..b038669e8 --- /dev/null +++ b/benchmarks/qwen3-tts/README.md @@ -0,0 +1,100 @@ +# Qwen3-TTS Benchmark + +Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference. + +## Prerequisites + +```bash +pip install matplotlib aiohttp soundfile numpy tqdm +pip install qwen_tts # for HF baseline +``` + +## Quick Start + +Run the full benchmark (vllm-omni + HF baseline) with a single command: + +```bash +cd benchmarks/qwen3-tts +bash run_benchmark.sh +``` + +Results (JSON + PNG plots) are saved to `results/`. + +### Common options + +```bash +# Only vllm-omni (skip HF baseline) +bash run_benchmark.sh --async-only + +# Only HF baseline +bash run_benchmark.sh --hf-only + +# Use a different model (e.g. 1.7B) +MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only + +# Use batch_size=4 config for higher throughput +STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only + +# Custom GPU, prompt count, concurrency levels +GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh +``` + +## Manual Steps + +### 1) Start the vLLM-Omni server + +```bash +CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \ + "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \ + --omni --host 127.0.0.1 --port 8000 \ + --stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \ + --trust-remote-code +``` + +### 2) Run online serving benchmark + +```bash +python vllm_omni/bench_tts_serve.py \ + --port 8000 \ + --num-prompts 50 \ + --max-concurrency 1 4 10 \ + --config-name "async_chunk" \ + --result-dir results/ +``` + +### 3) Run HuggingFace baseline + +```bash +python transformers/bench_tts_hf.py \ + --model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \ + --num-prompts 50 \ + --gpu-device 0 \ + --result-dir results/ +``` + +### 4) Generate comparison plots + +```bash +python plot_results.py \ + --results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \ + --labels "vllm-omni" "hf_transformers" \ + --output results/comparison.png +``` + +## Stage Configs + +| Config | Batch Size | Description | +|--------|:----------:|-------------| +| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) | +| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) | + +Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages. + +The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants. + +## Metrics + +- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency) +- **E2E (End-to-End Latency)**: Total time from request to complete audio response +- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis +- **Throughput**: Total audio seconds generated per wall-clock second diff --git a/benchmarks/qwen3-tts/plot_results.py b/benchmarks/qwen3-tts/plot_results.py new file mode 100644 index 000000000..e750101e3 --- /dev/null +++ b/benchmarks/qwen3-tts/plot_results.py @@ -0,0 +1,254 @@ +"""Plot Qwen3-TTS benchmark results. + +Generates comparison bar charts similar to the async_chunk design doc: +- TTFP (Time-to-First-Packet) across concurrency levels +- E2E latency across concurrency levels +- RTF (Real-Time Factor) across concurrency levels + +Usage: + # Compare two configs (async_chunk vs no_async_chunk): + python plot_results.py \ + --results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \ + --labels "async_chunk" "no_async_chunk" \ + --output results/qwen3_tts_benchmark.png + + # Single config: + python plot_results.py \ + --results results/bench_async_chunk_*.json \ + --labels "async_chunk" \ + --output results/qwen3_tts_benchmark.png +""" + +import argparse +import json +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + + +def load_results(result_files: list[str]) -> list[list[dict]]: + """Load benchmark results from JSON files.""" + all_results = [] + for f in result_files: + with open(f) as fh: + data = json.load(fh) + all_results.append(data) + return all_results + + +def plot_comparison( + all_results: list[list[dict]], + labels: list[str], + output_path: str, + title_prefix: str = "Qwen3-TTS", +): + """Generate comparison bar charts.""" + n_configs = len(all_results) + + # Collect concurrency levels present in ALL configs (skip missing data) + all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results] + concurrencies = sorted(set.union(*all_concurrencies)) + + # Build data arrays, using None for missing concurrency levels + ttfp_data = {label: [] for label in labels} + e2e_data = {label: [] for label in labels} + rtf_data = {label: [] for label in labels} + throughput_data = {label: [] for label in labels} + + for results, label in zip(all_results, labels): + conc_map = {r["concurrency"]: r for r in results} + for c in concurrencies: + r = conc_map.get(c) + ttfp_data[label].append(r["mean_ttfp_ms"] if r else None) + e2e_data[label].append(r["mean_e2e_ms"] if r else None) + rtf_data[label].append(r["mean_rtf"] if r else None) + throughput_data[label].append(r["audio_throughput"] if r else None) + + fig, axes = plt.subplots(2, 2, figsize=(14, 10)) + fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold") + + x = np.arange(len(concurrencies)) + width = 0.35 if n_configs == 2 else 0.5 + if n_configs > 1: + offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs) + else: + offsets = [0] + + colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"] + + def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"): + bars = [] + for i, (label, values) in enumerate(data_dict.items()): + # Replace None with 0 for plotting, but track which are missing + plot_values = [v if v is not None else 0 for v in values] + color = colors[i % len(colors)] + bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85) + bars.append(bar) + # Add value labels on bars (skip None/missing data) + max_val = max((v for v in values if v is not None), default=1) + for rect, val in zip(bar, values): + if val is not None and val > 0: + ax.text( + rect.get_x() + rect.get_width() / 2, + rect.get_height() + max_val * 0.02, + f"{val:{fmt}}", + ha="center", + va="bottom", + fontsize=9, + fontweight="bold", + ) + ax.set_xlabel("Concurrency", fontsize=12) + ax.set_ylabel(ylabel, fontsize=12) + ax.set_title(title, fontsize=13, fontweight="bold") + ax.set_xticks(x) + ax.set_xticklabels([str(c) for c in concurrencies]) + ax.legend(fontsize=10) + ax.grid(axis="y", alpha=0.3) + ax.set_axisbelow(True) + + plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)") + plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)") + plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f") + plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f") + + plt.tight_layout() + plt.savefig(output_path, dpi=150, bbox_inches="tight") + print(f"Plot saved to {output_path}") + plt.close() + + +def plot_single_summary(results: list[dict], label: str, output_path: str): + """Generate a single-config summary with percentile breakdown.""" + concurrencies = [r["concurrency"] for r in results] + + fig, axes = plt.subplots(1, 3, figsize=(16, 5)) + fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold") + + # TTFP breakdown + ax = axes[0] + means = [r["mean_ttfp_ms"] for r in results] + medians = [r["median_ttfp_ms"] for r in results] + p90s = [r["p90_ttfp_ms"] for r in results] + p99s = [r["p99_ttfp_ms"] for r in results] + x = np.arange(len(concurrencies)) + w = 0.2 + ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3") + ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50") + ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800") + ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336") + ax.set_xticks(x) + ax.set_xticklabels([str(c) for c in concurrencies]) + ax.set_xlabel("Concurrency") + ax.set_ylabel("TTFP (ms)") + ax.set_title("Time to First Audio Packet") + ax.legend(fontsize=9) + ax.grid(axis="y", alpha=0.3) + + # E2E breakdown + ax = axes[1] + means = [r["mean_e2e_ms"] for r in results] + medians = [r["median_e2e_ms"] for r in results] + p90s = [r["p90_e2e_ms"] for r in results] + p99s = [r["p99_e2e_ms"] for r in results] + ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3") + ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50") + ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800") + ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336") + ax.set_xticks(x) + ax.set_xticklabels([str(c) for c in concurrencies]) + ax.set_xlabel("Concurrency") + ax.set_ylabel("E2E Latency (ms)") + ax.set_title("End-to-End Latency") + ax.legend(fontsize=9) + ax.grid(axis="y", alpha=0.3) + + # RTF + ax = axes[2] + means = [r["mean_rtf"] for r in results] + medians = [r["median_rtf"] for r in results] + ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3") + ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50") + ax.set_xticks(x) + ax.set_xticklabels([str(c) for c in concurrencies]) + ax.set_xlabel("Concurrency") + ax.set_ylabel("RTF") + ax.set_title("Real-Time Factor") + ax.legend(fontsize=9) + ax.grid(axis="y", alpha=0.3) + + plt.tight_layout() + plt.savefig(output_path, dpi=150, bbox_inches="tight") + print(f"Plot saved to {output_path}") + plt.close() + + +def print_comparison_table(all_results: list[list[dict]], labels: list[str]): + """Print a markdown-formatted comparison table.""" + concurrencies = sorted(set(r["concurrency"] for r in all_results[0])) + + print("\n## Benchmark Results\n") + header = "| Metric | Concurrency |" + sep = "| --- | --- |" + for label in labels: + header += f" {label} |" + sep += " --- |" + print(header) + print(sep) + + for metric, key, fmt in [ + ("TTFP (ms)", "mean_ttfp_ms", ".1f"), + ("E2E (ms)", "mean_e2e_ms", ".1f"), + ("RTF", "mean_rtf", ".3f"), + ("Throughput (audio-s/s)", "audio_throughput", ".2f"), + ]: + for c in concurrencies: + row = f"| {metric} | {c} |" + for results in all_results: + conc_map = {r["concurrency"]: r for r in results} + val = conc_map.get(c, {}).get(key, 0) + row += f" {val:{fmt}} |" + print(row) + + # Improvement calculation (only if 2 configs) + if len(all_results) == 2: + print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n") + print("| Metric | Concurrency | Improvement |") + print("| --- | --- | --- |") + for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]: + for c in concurrencies: + m0 = {r["concurrency"]: r for r in all_results[0]} + m1 = {r["concurrency"]: r for r in all_results[1]} + v0 = m0.get(c, {}).get(key, 0) + v1 = m1.get(c, {}).get(key, 0) + if v1 > 0: + pct = (v1 - v0) / v1 * 100 + print(f"| {metric} | {c} | {pct:+.1f}% |") + + +def parse_args(): + parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results") + parser.add_argument( + "--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)" + ) + parser.add_argument( + "--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)" + ) + parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path") + parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot") + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + assert len(args.results) == len(args.labels), "--results and --labels must have the same count" + + all_results = load_results(args.results) + print_comparison_table(all_results, args.labels) + + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + + if len(all_results) == 1: + plot_single_summary(all_results[0], args.labels[0], args.output) + else: + plot_comparison(all_results, args.labels, args.output, title_prefix=args.title) diff --git a/benchmarks/qwen3-tts/results/.gitignore b/benchmarks/qwen3-tts/results/.gitignore new file mode 100644 index 000000000..5b6759ef7 --- /dev/null +++ b/benchmarks/qwen3-tts/results/.gitignore @@ -0,0 +1,3 @@ +# Benchmark results are machine-specific - do not commit +* +!.gitignore diff --git a/benchmarks/qwen3-tts/run_benchmark.sh b/benchmarks/qwen3-tts/run_benchmark.sh new file mode 100755 index 000000000..ef85d64d6 --- /dev/null +++ b/benchmarks/qwen3-tts/run_benchmark.sh @@ -0,0 +1,272 @@ +#!/bin/bash +# Qwen3-TTS Benchmark Runner +# +# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference. +# Produces JSON results and comparison plots. +# +# Usage: +# # Full comparison (vllm-omni + HF): +# bash run_benchmark.sh +# +# # Only vllm-omni async_chunk config: +# bash run_benchmark.sh --async-only +# +# # Only HuggingFace baseline: +# bash run_benchmark.sh --hf-only +# +# # vllm-omni only (skip HF): +# bash run_benchmark.sh --skip-hf +# +# # Custom settings: +# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh +# +# # Use 1.7B model: +# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only +# +# # Use batch_size=4 config: +# STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only +# +# Environment variables: +# GPU_DEVICE - GPU index to use (default: 0) +# NUM_PROMPTS - Number of prompts per concurrency level (default: 50) +# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10") +# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice) +# PORT - Server port (default: 8000) +# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3) +# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2) +# STAGE_CONFIG - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml) + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" + +# Defaults +GPU_DEVICE="${GPU_DEVICE:-0}" +NUM_PROMPTS="${NUM_PROMPTS:-50}" +CONCURRENCY="${CONCURRENCY:-1 4 10}" +MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}" +PORT="${PORT:-8000}" +GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}" +GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}" +NUM_WARMUPS="${NUM_WARMUPS:-3}" +STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}" +RESULT_DIR="${SCRIPT_DIR}/results" +TIMESTAMP="$(date +%Y%m%d_%H%M%S)" + +# Parse args +RUN_ASYNC=true +RUN_HF=true +for arg in "$@"; do + case "$arg" in + --async-only) RUN_HF=false ;; + --hf-only) RUN_ASYNC=false ;; + --skip-hf) RUN_HF=false ;; + esac +done + +mkdir -p "${RESULT_DIR}" + +echo "============================================================" +echo " Qwen3-TTS Benchmark" +echo "============================================================" +echo " GPU: ${GPU_DEVICE}" +echo " Model: ${MODEL}" +echo " Prompts: ${NUM_PROMPTS}" +echo " Concurrency: ${CONCURRENCY}" +echo " Port: ${PORT}" +echo " Stage config: ${STAGE_CONFIG}" +echo " Results: ${RESULT_DIR}" +echo "============================================================" + +# Prepare stage config with correct GPU device and memory settings +prepare_config() { + local config_template="$1" + local config_name="$2" + local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml" + + # Use sed to patch GPU device and memory utilization + sed \ + -e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \ + -e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \ + -e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \ + "${config_template}" > "${output_path}" + + echo "${output_path}" +} + +# Start server and wait for it to be ready +start_server() { + local stage_config="$1" + local config_name="$2" + local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log" + + echo "" + echo "Starting server with config: ${config_name}" + echo " Stage config: ${stage_config}" + echo " Log file: ${log_file}" + + VLLM_WORKER_MULTIPROC_METHOD=spawn \ + CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \ + python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \ + --omni \ + --host 127.0.0.1 \ + --port "${PORT}" \ + --stage-configs-path "${stage_config}" \ + --stage-init-timeout 120 \ + --trust-remote-code \ + --disable-log-stats \ + > "${log_file}" 2>&1 & + + SERVER_PID=$! + echo " Server PID: ${SERVER_PID}" + + # Wait for server to be ready + echo " Waiting for server to be ready..." + local max_wait=300 + local waited=0 + while [ ${waited} -lt ${max_wait} ]; do + if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then + echo " Server is ready! (waited ${waited}s)" + return 0 + fi + # Check if process is still alive + if ! kill -0 ${SERVER_PID} 2>/dev/null; then + echo " ERROR: Server process died. Check log: ${log_file}" + tail -20 "${log_file}" + return 1 + fi + sleep 2 + waited=$((waited + 2)) + done + + echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}" + kill ${SERVER_PID} 2>/dev/null || true + return 1 +} + +# Stop the server +stop_server() { + if [ -n "${SERVER_PID:-}" ]; then + echo " Stopping server (PID: ${SERVER_PID})..." + kill ${SERVER_PID} 2>/dev/null || true + wait ${SERVER_PID} 2>/dev/null || true + # Kill any remaining child processes on the port + local pids + pids=$(lsof -ti:${PORT} 2>/dev/null || true) + if [ -n "${pids}" ]; then + echo " Cleaning up remaining processes on port ${PORT}..." + echo "${pids}" | xargs kill -9 2>/dev/null || true + fi + echo " Server stopped." + SERVER_PID="" + fi +} + +# Cleanup on exit +trap 'stop_server' EXIT + +# Run benchmark for a given config +run_bench() { + local config_name="$1" + local config_template="$2" + + echo "" + echo "============================================================" + echo " Benchmarking: ${config_name}" + echo "============================================================" + + local stage_config + stage_config=$(prepare_config "${config_template}" "${config_name}") + + start_server "${stage_config}" "${config_name}" + + # Convert concurrency string to args + local conc_args="" + for c in ${CONCURRENCY}; do + conc_args="${conc_args} ${c}" + done + + cd "${PROJECT_ROOT}" + python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \ + --host 127.0.0.1 \ + --port "${PORT}" \ + --num-prompts "${NUM_PROMPTS}" \ + --max-concurrency ${conc_args} \ + --num-warmups "${NUM_WARMUPS}" \ + --config-name "${config_name}" \ + --result-dir "${RESULT_DIR}" + + stop_server + + # Allow GPU memory to settle + sleep 5 +} + +# Run vllm-omni benchmark +if [ "${RUN_ASYNC}" = true ]; then + run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}" +fi + +# Run HuggingFace baseline benchmark +if [ "${RUN_HF}" = true ]; then + echo "" + echo "============================================================" + echo " Benchmarking: HuggingFace transformers (offline)" + echo "============================================================" + + cd "${PROJECT_ROOT}" + python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \ + --model "${MODEL}" \ + --num-prompts "${NUM_PROMPTS}" \ + --num-warmups "${NUM_WARMUPS}" \ + --gpu-device "${GPU_DEVICE}" \ + --config-name "hf_transformers" \ + --result-dir "${RESULT_DIR}" + + # Allow GPU memory to settle + sleep 5 +fi + +# Plot results +echo "" +echo "============================================================" +echo " Generating plots..." +echo "============================================================" + +RESULT_FILES="" +LABELS="" + +if [ "${RUN_ASYNC}" = true ]; then + ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1) + if [ -n "${ASYNC_FILE}" ]; then + RESULT_FILES="${ASYNC_FILE}" + LABELS="async_chunk" + fi +fi + +if [ "${RUN_HF}" = true ]; then + HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1) + if [ -n "${HF_FILE}" ]; then + if [ -n "${RESULT_FILES}" ]; then + RESULT_FILES="${RESULT_FILES} ${HF_FILE}" + LABELS="${LABELS} hf_transformers" + else + RESULT_FILES="${HF_FILE}" + LABELS="hf_transformers" + fi + fi +fi + +if [ -n "${RESULT_FILES}" ]; then + python "${SCRIPT_DIR}/plot_results.py" \ + --results ${RESULT_FILES} \ + --labels ${LABELS} \ + --output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png" +fi + +echo "" +echo "============================================================" +echo " Benchmark complete!" +echo " Results: ${RESULT_DIR}" +echo "============================================================" diff --git a/benchmarks/qwen3-tts/transformers/bench_tts_hf.py b/benchmarks/qwen3-tts/transformers/bench_tts_hf.py new file mode 100644 index 000000000..0e0ef8e9e --- /dev/null +++ b/benchmarks/qwen3-tts/transformers/bench_tts_hf.py @@ -0,0 +1,281 @@ +"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library). + +Measures E2E latency, RTF, and audio duration for offline (non-serving) inference. +Results are saved in the same JSON format as bench_tts_serve.py for unified plotting. + +Usage: + python bench_tts_hf.py \ + --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \ + --num-prompts 50 \ + --num-warmups 3 \ + --gpu-device 0 \ + --result-dir results/ +""" + +import argparse +import json +import time +from dataclasses import asdict, dataclass, field +from datetime import datetime +from pathlib import Path + +import numpy as np +import soundfile as sf +import torch + +PROMPTS = [ + "Hello, welcome to the voice synthesis benchmark test.", + "She said she would be here by noon, but nobody showed up.", + "The quick brown fox jumps over the lazy dog near the riverbank.", + "I can't believe how beautiful the sunset looks from up here on the mountain.", + "Please remember to bring your identification documents to the appointment tomorrow morning.", + "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?", + "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.", + "After the meeting, we should discuss the quarterly results and plan for the next phase.", + "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.", + "The train leaves at half past seven, so we need to arrive at the station before then.", + "Could you please turn down the music a little bit, I'm trying to concentrate on my work.", + "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.", +] + + +@dataclass +class BenchmarkResult: + config_name: str = "" + concurrency: int = 1 # always 1 for offline + num_prompts: int = 0 + completed: int = 0 + failed: int = 0 + duration_s: float = 0.0 + # TTFP stats - not applicable for HF offline, set to E2E for compatibility + mean_ttfp_ms: float = 0.0 + median_ttfp_ms: float = 0.0 + std_ttfp_ms: float = 0.0 + p90_ttfp_ms: float = 0.0 + p95_ttfp_ms: float = 0.0 + p99_ttfp_ms: float = 0.0 + # E2E stats (ms) + mean_e2e_ms: float = 0.0 + median_e2e_ms: float = 0.0 + std_e2e_ms: float = 0.0 + p90_e2e_ms: float = 0.0 + p95_e2e_ms: float = 0.0 + p99_e2e_ms: float = 0.0 + # RTF stats + mean_rtf: float = 0.0 + median_rtf: float = 0.0 + std_rtf: float = 0.0 + p99_rtf: float = 0.0 + # Audio stats + mean_audio_duration_s: float = 0.0 + total_audio_duration_s: float = 0.0 + audio_throughput: float = 0.0 + request_throughput: float = 0.0 + # Per-request details + per_request: list = field(default_factory=list) + + +def run_benchmark(args): + from qwen_tts import Qwen3TTSModel + + device = f"cuda:{args.gpu_device}" + print(f"Loading model: {args.model} on {device}") + model = Qwen3TTSModel.from_pretrained( + args.model, + device_map=device, + dtype=torch.bfloat16, + ) + print("Model loaded.") + + # Build prompt list + prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)] + + # Warmup + if args.num_warmups > 0: + print(f"Warming up with {args.num_warmups} requests...") + for i in range(args.num_warmups): + p = PROMPTS[i % len(PROMPTS)] + wavs, sr = model.generate_custom_voice( + text=p, + language=args.language, + speaker=args.voice, + ) + # Sync GPU + torch.cuda.synchronize(device) + print("Warmup done.") + + # Benchmark + print(f"Running {args.num_prompts} requests sequentially...") + e2e_times = [] + rtfs = [] + audio_durations = [] + per_request = [] + failed = 0 + + audio_dir = None + if args.save_audio: + audio_dir = Path(args.result_dir) / "audio_hf" + audio_dir.mkdir(parents=True, exist_ok=True) + + total_start = time.perf_counter() + + for i, prompt in enumerate(prompts): + try: + torch.cuda.synchronize(device) + st = time.perf_counter() + + wavs, sr = model.generate_custom_voice( + text=prompt, + language=args.language, + speaker=args.voice, + ) + + torch.cuda.synchronize(device) + elapsed = time.perf_counter() - st + + # Compute audio duration + audio_samples = wavs[0] + if isinstance(audio_samples, torch.Tensor): + audio_samples = audio_samples.cpu().numpy() + audio_dur = len(audio_samples) / sr + + rtf = elapsed / audio_dur if audio_dur > 0 else 0.0 + + e2e_times.append(elapsed) + rtfs.append(rtf) + audio_durations.append(audio_dur) + per_request.append( + { + "e2e_ms": elapsed * 1000, + "ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E + "rtf": rtf, + "audio_duration_s": audio_dur, + "prompt": prompt, + } + ) + + if audio_dir: + sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr) + + if (i + 1) % 10 == 0 or i == 0: + print( + f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s" + ) + + except Exception as e: + print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}") + failed += 1 + + total_duration = time.perf_counter() - total_start + completed = len(e2e_times) + + # Compute stats + result = BenchmarkResult( + config_name=args.config_name, + concurrency=1, + num_prompts=args.num_prompts, + completed=completed, + failed=failed, + duration_s=total_duration, + ) + + if e2e_times: + e2e_ms = [t * 1000 for t in e2e_times] + + result.mean_e2e_ms = float(np.mean(e2e_ms)) + result.median_e2e_ms = float(np.median(e2e_ms)) + result.std_e2e_ms = float(np.std(e2e_ms)) + result.p90_e2e_ms = float(np.percentile(e2e_ms, 90)) + result.p95_e2e_ms = float(np.percentile(e2e_ms, 95)) + result.p99_e2e_ms = float(np.percentile(e2e_ms, 99)) + + # For HF offline, TTFP = E2E (no streaming) + result.mean_ttfp_ms = result.mean_e2e_ms + result.median_ttfp_ms = result.median_e2e_ms + result.std_ttfp_ms = result.std_e2e_ms + result.p90_ttfp_ms = result.p90_e2e_ms + result.p95_ttfp_ms = result.p95_e2e_ms + result.p99_ttfp_ms = result.p99_e2e_ms + + result.mean_rtf = float(np.mean(rtfs)) + result.median_rtf = float(np.median(rtfs)) + result.std_rtf = float(np.std(rtfs)) + result.p99_rtf = float(np.percentile(rtfs, 99)) + + result.mean_audio_duration_s = float(np.mean(audio_durations)) + result.total_audio_duration_s = float(np.sum(audio_durations)) + result.audio_throughput = result.total_audio_duration_s / total_duration + result.request_throughput = completed / total_duration + result.per_request = per_request + + # Print summary in standardized performance template + W = 50 + print("") + print(f"{'=' * W}") + print(f"{'Serving Benchmark Result':^{W}}") + print(f"{'=' * W}") + print(f"{'Successful requests:':<40}{completed:<10}") + print(f"{'Failed requests:':<40}{failed:<10}") + print(f"{'Maximum request concurrency:':<40}{1:<10}") + print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}") + print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}") + print(f"{'-' * W}") + print(f"{'End-to-end Latency':^{W}}") + print(f"{'-' * W}") + print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}") + print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}") + print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}") + print(f"{'=' * W}") + print(f"{'Audio Result':^{W}}") + print(f"{'=' * W}") + print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}") + print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}") + print(f"{'-' * W}") + print(f"{'Time to First Packet':^{W}}") + print(f"{'-' * W}") + print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}") + print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}") + print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}") + print(f"{'-' * W}") + print(f"{'Real Time Factor':^{W}}") + print(f"{'-' * W}") + print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}") + print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}") + print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}") + print(f"{'=' * W}") + print("") + + # Save results (as a list with single concurrency=1 entry, matching serve format) + result_dir = Path(args.result_dir) + result_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json" + + with open(result_file, "w") as f: + json.dump([asdict(result)], f, indent=2) + print(f"Results saved to {result_file}") + + return result + + +def parse_args(): + parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark") + parser.add_argument( + "--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path" + ) + parser.add_argument("--num-prompts", type=int, default=50) + parser.add_argument("--num-warmups", type=int, default=3) + parser.add_argument("--gpu-device", type=int, default=0) + parser.add_argument("--voice", type=str, default="Vivian") + parser.add_argument("--language", type=str, default="English") + parser.add_argument( + "--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)" + ) + parser.add_argument("--result-dir", type=str, default="results") + parser.add_argument("--save-audio", action="store_true", help="Save generated audio files") + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + run_benchmark(args) diff --git a/benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py b/benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py new file mode 100644 index 000000000..91e4ecbbb --- /dev/null +++ b/benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py @@ -0,0 +1,348 @@ +"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint. + +Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor) +across configurable concurrency levels. Saves results as JSON for plotting. + +Usage: + python bench_tts_serve.py \ + --host 127.0.0.1 --port 8000 \ + --num-prompts 50 \ + --max-concurrency 1 4 10 \ + --result-dir results/ +""" + +import argparse +import asyncio +import json +import time +from dataclasses import asdict, dataclass, field +from datetime import datetime +from pathlib import Path + +import aiohttp +import numpy as np +from tqdm.asyncio import tqdm + +PROMPTS = [ + "Hello, welcome to the voice synthesis benchmark test.", + "She said she would be here by noon, but nobody showed up.", + "The quick brown fox jumps over the lazy dog near the riverbank.", + "I can't believe how beautiful the sunset looks from up here on the mountain.", + "Please remember to bring your identification documents to the appointment tomorrow morning.", + "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?", + "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.", + "After the meeting, we should discuss the quarterly results and plan for the next phase.", + "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.", + "The train leaves at half past seven, so we need to arrive at the station before then.", + "Could you please turn down the music a little bit, I'm trying to concentrate on my work.", + "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.", +] + + +@dataclass +class RequestResult: + success: bool = False + ttfp: float = 0.0 # Time to first audio packet (seconds) + e2e: float = 0.0 # End-to-end latency (seconds) + audio_bytes: int = 0 # Total audio bytes received + audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM) + rtf: float = 0.0 # Real-time factor = e2e / audio_duration + prompt: str = "" + error: str = "" + + +@dataclass +class BenchmarkResult: + config_name: str = "" + concurrency: int = 0 + num_prompts: int = 0 + completed: int = 0 + failed: int = 0 + duration_s: float = 0.0 + # TTFP stats (ms) + mean_ttfp_ms: float = 0.0 + median_ttfp_ms: float = 0.0 + std_ttfp_ms: float = 0.0 + p90_ttfp_ms: float = 0.0 + p95_ttfp_ms: float = 0.0 + p99_ttfp_ms: float = 0.0 + # E2E stats (ms) + mean_e2e_ms: float = 0.0 + median_e2e_ms: float = 0.0 + std_e2e_ms: float = 0.0 + p90_e2e_ms: float = 0.0 + p95_e2e_ms: float = 0.0 + p99_e2e_ms: float = 0.0 + # RTF stats + mean_rtf: float = 0.0 + median_rtf: float = 0.0 + std_rtf: float = 0.0 + p99_rtf: float = 0.0 + # Audio stats + mean_audio_duration_s: float = 0.0 + total_audio_duration_s: float = 0.0 + audio_throughput: float = 0.0 # audio_duration / wall_time + request_throughput: float = 0.0 # requests / second + # Per-request details + per_request: list = field(default_factory=list) + + +def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float: + """Convert raw PCM byte count to duration in seconds.""" + num_samples = num_bytes / sample_width + return num_samples / sample_rate + + +async def send_tts_request( + session: aiohttp.ClientSession, + api_url: str, + prompt: str, + voice: str = "vivian", + language: str = "English", + pbar: tqdm | None = None, +) -> RequestResult: + """Send a streaming TTS request and measure latency metrics.""" + payload = { + "input": prompt, + "voice": voice, + "language": language, + "stream": True, + "response_format": "pcm", + } + + result = RequestResult(prompt=prompt) + st = time.perf_counter() + + try: + async with session.post(api_url, json=payload) as response: + if response.status != 200: + result.error = f"HTTP {response.status}: {await response.text()}" + result.success = False + return result + + first_chunk = True + total_bytes = 0 + + async for chunk in response.content.iter_any(): + if first_chunk and len(chunk) > 0: + result.ttfp = time.perf_counter() - st + first_chunk = False + total_bytes += len(chunk) + + result.e2e = time.perf_counter() - st + result.audio_bytes = total_bytes + result.audio_duration = pcm_bytes_to_duration(total_bytes) + + if result.audio_duration > 0: + result.rtf = result.e2e / result.audio_duration + result.success = True + + except Exception as e: + result.error = str(e) + result.success = False + result.e2e = time.perf_counter() - st + + if pbar: + pbar.update(1) + return result + + +async def run_benchmark( + host: str, + port: int, + num_prompts: int, + max_concurrency: int, + num_warmups: int = 3, + voice: str = "vivian", + language: str = "English", +) -> BenchmarkResult: + """Run benchmark at a given concurrency level.""" + api_url = f"http://{host}:{port}/v1/audio/speech" + + connector = aiohttp.TCPConnector( + limit=max_concurrency, + limit_per_host=max_concurrency, + keepalive_timeout=60, + ) + session = aiohttp.ClientSession( + connector=connector, + timeout=aiohttp.ClientTimeout(total=600), + ) + + # Warmup + if num_warmups > 0: + print(f" Warming up with {num_warmups} requests...") + warmup_tasks = [] + for i in range(num_warmups): + prompt = PROMPTS[i % len(PROMPTS)] + warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language)) + await asyncio.gather(*warmup_tasks) + print(" Warmup done.") + + # Build request list + request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)] + + # Run benchmark + print(f" Running {num_prompts} requests with concurrency={max_concurrency}...") + semaphore = asyncio.Semaphore(max_concurrency) + pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}") + + async def limited_request(prompt): + async with semaphore: + return await send_tts_request(session, api_url, prompt, voice, language, pbar) + + start_time = time.perf_counter() + tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts] + results: list[RequestResult] = await asyncio.gather(*tasks) + duration = time.perf_counter() - start_time + pbar.close() + + await session.close() + + # Compute stats + successful = [r for r in results if r.success] + failed = [r for r in results if not r.success] + + bench = BenchmarkResult( + concurrency=max_concurrency, + num_prompts=num_prompts, + completed=len(successful), + failed=len(failed), + duration_s=duration, + ) + + if successful: + ttfps = [r.ttfp * 1000 for r in successful] # convert to ms + e2es = [r.e2e * 1000 for r in successful] + rtfs = [r.rtf for r in successful] + audio_durs = [r.audio_duration for r in successful] + + bench.mean_ttfp_ms = float(np.mean(ttfps)) + bench.median_ttfp_ms = float(np.median(ttfps)) + bench.std_ttfp_ms = float(np.std(ttfps)) + bench.p90_ttfp_ms = float(np.percentile(ttfps, 90)) + bench.p95_ttfp_ms = float(np.percentile(ttfps, 95)) + bench.p99_ttfp_ms = float(np.percentile(ttfps, 99)) + + bench.mean_e2e_ms = float(np.mean(e2es)) + bench.median_e2e_ms = float(np.median(e2es)) + bench.std_e2e_ms = float(np.std(e2es)) + bench.p90_e2e_ms = float(np.percentile(e2es, 90)) + bench.p95_e2e_ms = float(np.percentile(e2es, 95)) + bench.p99_e2e_ms = float(np.percentile(e2es, 99)) + + bench.mean_rtf = float(np.mean(rtfs)) + bench.median_rtf = float(np.median(rtfs)) + bench.std_rtf = float(np.std(rtfs)) + bench.p99_rtf = float(np.percentile(rtfs, 99)) + + bench.mean_audio_duration_s = float(np.mean(audio_durs)) + bench.total_audio_duration_s = float(np.sum(audio_durs)) + bench.audio_throughput = bench.total_audio_duration_s / duration + bench.request_throughput = len(successful) / duration + + bench.per_request = [ + { + "ttfp_ms": r.ttfp * 1000, + "e2e_ms": r.e2e * 1000, + "rtf": r.rtf, + "audio_duration_s": r.audio_duration, + "prompt": r.prompt, + } + for r in successful + ] + + # Print summary in standardized performance template + W = 50 + print("") + print(f"{'=' * W}") + print(f"{'Serving Benchmark Result':^{W}}") + print(f"{'=' * W}") + print(f"{'Successful requests:':<40}{bench.completed:<10}") + print(f"{'Failed requests:':<40}{bench.failed:<10}") + print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}") + print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}") + print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}") + print(f"{'-' * W}") + print(f"{'End-to-end Latency':^{W}}") + print(f"{'-' * W}") + print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}") + print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}") + print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}") + print(f"{'=' * W}") + print(f"{'Audio Result':^{W}}") + print(f"{'=' * W}") + print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}") + print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}") + print(f"{'-' * W}") + print(f"{'Time to First Packet':^{W}}") + print(f"{'-' * W}") + print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}") + print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}") + print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}") + print(f"{'-' * W}") + print(f"{'Real Time Factor':^{W}}") + print(f"{'-' * W}") + print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}") + print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}") + print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}") + print(f"{'=' * W}") + print("") + + if failed: + for r in failed[:3]: + print(f" [ERROR] {r.error[:200]}") + + return bench + + +async def main(args): + all_results = [] + + for concurrency in args.max_concurrency: + result = await run_benchmark( + host=args.host, + port=args.port, + num_prompts=args.num_prompts, + max_concurrency=concurrency, + num_warmups=args.num_warmups, + voice=args.voice, + language=args.language, + ) + result.config_name = args.config_name + all_results.append(asdict(result)) + + # Save results + result_dir = Path(args.result_dir) + result_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json" + + with open(result_file, "w") as f: + json.dump(all_results, f, indent=2) + print(f"Results saved to {result_file}") + + return all_results + + +def parse_args(): + parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client") + parser.add_argument("--host", type=str, default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level") + parser.add_argument( # noqa: E501 + "--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test" + ) + parser.add_argument("--num-warmups", type=int, default=3) + parser.add_argument("--voice", type=str, default="vivian") + parser.add_argument("--language", type=str, default="English") + parser.add_argument( + "--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)" + ) + parser.add_argument("--result-dir", type=str, default="results") + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + asyncio.run(main(args)) diff --git a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml new file mode 100644 index 000000000..1597f8aa2 --- /dev/null +++ b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml @@ -0,0 +1,96 @@ +# Qwen3-TTS batch_size=1 config (streaming with async_chunk) +# 2-stage pipeline: Talker -> Code2Wav +async_chunk: true +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: qwen3_tts + model_arch: Qwen3TTSTalkerForConditionalGeneration + hf_overrides: + architectures: [Qwen3TTSTalkerForConditionalGeneration] + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: false + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.3 + distributed_executor_backend: "mp" + max_num_batched_tokens: 512 + max_model_len: 4096 + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + detokenize: false + repetition_penalty: 1.05 + stop_token_ids: [2150] + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: code2wav + model_arch: Qwen3TTSCode2Wav + hf_overrides: + architectures: [Qwen3TTSCode2Wav] + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.2 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 32768 + engine_input_source: [0] + final_output: true + final_output_type: audio + input_connectors: + from_stage_0: connector_of_shared_memory + tts_args: + max_instructions_length: 500 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + detokenize: true + repetition_penalty: 1.0 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 25 + + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml new file mode 100644 index 000000000..18493f3ae --- /dev/null +++ b/benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml @@ -0,0 +1,97 @@ +# Qwen3-TTS batch_size=4 config (streaming with async_chunk) +# Enables concurrent request processing with max_inflight=4 +# 2-stage pipeline: Talker -> Code2Wav +async_chunk: true +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 4 + engine_args: + model_stage: qwen3_tts + model_arch: Qwen3TTSTalkerForConditionalGeneration + hf_overrides: + architectures: [Qwen3TTSTalkerForConditionalGeneration] + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: false + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.3 + distributed_executor_backend: "mp" + max_num_batched_tokens: 512 + max_model_len: 4096 + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + detokenize: false + repetition_penalty: 1.05 + stop_token_ids: [2150] + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 4 + engine_args: + model_stage: code2wav + model_arch: Qwen3TTSCode2Wav + hf_overrides: + architectures: [Qwen3TTSCode2Wav] + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.2 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 32768 + engine_input_source: [0] + final_output: true + final_output_type: audio + input_connectors: + from_stage_0: connector_of_shared_memory + tts_args: + max_instructions_length: 500 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + detokenize: true + repetition_penalty: 1.0 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 4 + + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 25 + + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/docs/.nav.yml b/docs/.nav.yml index 4aa8fee43..eb135c291 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -13,9 +13,14 @@ nav: - examples/README.md - Offline Inference: - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md + - CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md + - Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md + - Helios Video Generation: user_guide/examples/offline_inference/helios.md + - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md + - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md @@ -25,6 +30,8 @@ nav: - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md - Online Serving: - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md + - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md + - Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md - Image-To-Image: user_guide/examples/online_serving/image_to_image.md - Image-To-Video: user_guide/examples/online_serving/image_to_video.md diff --git a/docs/README.md b/docs/README.md index 363f579c9..66fc8ef46 100644 --- a/docs/README.md +++ b/docs/README.md @@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in For more information, checkout the following: -- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true) +- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1) - [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html) diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg index c5e88d616..ce90f1753 100644 Binary files a/docs/assets/WeChat.jpg and b/docs/assets/WeChat.jpg differ diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 412f3d524..ed69cdfe1 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in For offline (Assume necessary dependencies have ben imported): ```python model_name = "Qwen/Qwen2.5-Omni-7B" -omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml") +omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml") ``` For online serving: diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md index c8493af0a..54af850bb 100644 --- a/docs/contributing/ci/CI_5levels.md +++ b/docs/contributing/ci/CI_5levels.md @@ -508,7 +508,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu **2.4.5 Request Execution** ```python - openai_client.send_request(request_config, request_num=1) + openai_client.send_omni_request(request_config, request_num=1) # for omni-understanding models + # or + openai_client.send_diffusion_request(request_config, request_num=1) # for diffusion models ``` **Explanation**: @@ -517,7 +519,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu **Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n. - **Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation. + **Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation. - ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model` diff --git a/docs/contributing/ci/tests_markers.md b/docs/contributing/ci/tests_markers.md index c9e5d18e6..6200ca5ae 100644 --- a/docs/contributing/ci/tests_markers.md +++ b/docs/contributing/ci/tests_markers.md @@ -32,7 +32,7 @@ Defined in `pyproject.toml`: | `slow` | Slow tests (may skip in quick CI) | | `benchmark` | Benchmark tests | -\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator. +\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility). ### Example usage for markers @@ -49,6 +49,7 @@ from tests.utils import hardware_test def test_video_to_audio() ... ``` + ### Decorator: `@hardware_test` This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions: @@ -98,6 +99,24 @@ This decorator is intended to make hardware-aware, cross-platform test authoring - ROCm only: `pytest -m "distributed_rocm and MI325"` - NPU only: `pytest -m "distributed_npu"` +### Function: `hardware_marks` + +`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function. + +```python +from tests.utils import hardware_marks + +MULTI_CARD_MARKS = hardware_marks( + res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2 +) + +@pytest.mark.parametrize("omni_server", [ + pytest.param(OmniServerParams(...), id="case_001", marks=MULTI_CARD_MARKS), +], indirect=True) +def test_feature(omni_server): + ... +``` + ## Add Support for a New Platform If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps: @@ -132,12 +151,13 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator # Optionally: add skipif_tpu when implemented return [test_platform, test_resource, test_distributed] ``` -3. **Update `hardware_test` to recognize your new platform**: - In the relevant place (see the `hardware_test` implementation), add: +3. **Update `hardware_marks` to recognize your new platform**: + In the relevant place (see the `hardware_marks` implementation), add: ```python if platform == "tpu": marks = tpu_marks(res=resource, num_cards=cards) ``` + (`hardware_test` calls `hardware_marks` internally, so both will pick up the change.) 4. **(Recommended) Add a test using your new markers**: ```python @hardware_test( @@ -151,7 +171,7 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator **Summary**: - Add pytest markers for your new platform/resources - Implement a marker function (`xxx_marks`) -- Plug into `hardware_test` -- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation! +- Plug into `hardware_marks` +- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation! See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`). diff --git a/docs/contributing/ci/tests_style.md b/docs/contributing/ci/tests_style.md index ab94f2f56..a24f71864 100644 --- a/docs/contributing/ci/tests_style.md +++ b/docs/contributing/ci/tests_style.md @@ -158,6 +158,7 @@ vllm_omni/ tests/ β”œβ”€β”€ test_diffusion_layerwise_offload.py β”œβ”€β”€ test_diffusion_lora.py β”œβ”€β”€ test_sequence_parallel.py + β”œβ”€β”€ test_qwen_image_edit_expansion.py └── stage_configs/ β”œβ”€β”€ qwen2_5_omni_ci.yaml β”œβ”€β”€ qwen3_omni_ci.yaml diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index c1094802c..818bd15b7 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -651,8 +651,27 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num #### 5.2 Functionality Check in CI -To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes. -For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md). +To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes. + +For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md). +The following tests are required to add: + +- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/). + - Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features: + - 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise) + - 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2 + - 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2 + - 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2 + - 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2 + - 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2 + - Validation: test that the multimodal output files of your model have the correct shapes. + - Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/) + - To maximize code reuse, you may refer to + - `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper. + - `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`. + - [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases. + - Doc: add a concise dostring for each test function. + - Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py). --- diff --git a/docs/design/feature/omni_connectors/mooncake_store_connector.md b/docs/design/feature/omni_connectors/mooncake_store_connector.md index 5812a6abe..2ff116859 100644 --- a/docs/design/feature/omni_connectors/mooncake_store_connector.md +++ b/docs/design/feature/omni_connectors/mooncake_store_connector.md @@ -72,3 +72,241 @@ Parameters: For more details, refer to the [Mooncake repository](https://github.com/kvcache-ai/Mooncake). + +--- + +## Design + +### 1. Overview + +`MooncakeStoreConnector` is the storage-oriented remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake distributed store APIs and provides a uniform `put()` / `get()` abstraction for transferring stage payloads across nodes. + +Compared with `SharedMemoryConnector`, this connector is intended for multi-node deployments. Compared with `MooncakeTransferEngineConnector`, it is simpler and more object-oriented: data is written into a distributed store and fetched back by key, rather than transferred through an explicit remote-memory write protocol. + +Its primary role is to provide a general-purpose remote connector for stage payloads when: + +- the pipeline spans multiple hosts +- a shared-memory connector is not possible +- a simpler remote backend is preferred over the RDMA-oriented transfer engine + +### 2. Relationship with the OmniConnector System + +`MooncakeStoreConnector` implements `OmniConnectorBase`, so it integrates with the same connector lifecycle as the other backends: + +- `OmniConnectorFactory` instantiates it from `ConnectorSpec` +- `load_omni_transfer_config()` maps edge config to the connector spec +- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract + +This means the rest of the pipeline does not need store-specific logic. It only relies on the generic connector contract. + +### 3. Design Goals + +The connector is designed around the following goals: + +- **Cross-node object transfer** using a single key-based abstraction +- **Minimal connector-specific control plane**, since the store itself is the shared medium +- **Compatibility with arbitrary Python payloads** through the Omni serializer +- **Operational simplicity** compared with the transfer-engine-based remote connector + +The connector is intentionally not optimized for zero-copy tensor movement or direct remote-memory access. + +### 4. Core Design + +#### 4.1 Object-Oriented Transfer Model + +`MooncakeStoreConnector` treats the transport backend as a distributed object store: + +1. Serialize the Python object into bytes. +2. Generate a unique connector key. +3. Store the bytes in Mooncake. +4. Fetch the bytes from Mooncake on the receiver side. +5. Deserialize the bytes back into the original object. + +This design makes the connector conceptually close to a distributed key-value transport. + +#### 4.2 Key Construction + +The connector uses `OmniConnectorBase._make_key()` to derive the internal store key: + +```text +{request_key}@{from_stage}_{to_stage} +``` + +This adds stage routing information to the logical request key and avoids collisions between different pipeline edges. + +#### 4.3 No Extra Metadata Path + +Unlike `SharedMemoryConnector` and `MooncakeTransferEngineConnector`, this connector returns: + +```python +metadata = None +``` + +from `put()`. + +This is a meaningful design difference: + +- the store itself is the shared rendezvous point +- consumers only need the same key +- no transport handle, address, or side-channel metadata needs to be propagated + +As a result, the control plane is simpler for this connector than for the SHM and RDMA variants. + +### 5. Initialization + +#### 5.1 Required Mooncake Components + +The connector requires the Mooncake Python bindings to expose: + +- `MooncakeDistributedStore` +- `ReplicateConfig` + +If these symbols are unavailable, construction fails immediately with `ImportError`. This keeps startup failures explicit and avoids silent fallback behavior. + +#### 5.2 Store Setup + +During `_init_store()`, the connector: + +1. creates `MooncakeDistributedStore()` +2. calls `store.setup(...)` +3. validates the return code +4. creates a `ReplicateConfig` +5. enables `with_soft_pin = True` + +The setup step is completed during connector construction, not lazily at first use. + +### 6. Put / Get Flow + +#### 6.1 Producer Flow: `put()` + +The producer-side flow is: + +1. Validate that the store has been initialized. +2. Serialize the payload into bytes. +3. Build the internal key with stage routing information. +4. Call `store.put(key, serialized_data, self.pin)`. +5. Update metrics and return success. + +The returned tuple is: + +```python +(True, len(serialized_data), None) +``` + +#### 6.2 Consumer Flow: `get()` + +The consumer-side flow is: + +1. Validate that the store has been initialized. +2. Build the same internal key. +3. Poll `store.get(key)` for up to 20 retries. +4. Sleep 50 ms between retries. +5. If data is found, deserialize it and return `(data, payload_size)`. +6. If all retries are exhausted, record a timeout and return `None`. + +This design gives the connector a bounded waiting model rather than an indefinite blocking get. + +### 7. Integration with Stage Communication + +All callers use the connector through the same `put()` / `get()` contract: + +- the sender calls `put()` to serialize and store the payload +- the receiver calls `get()` to retrieve and deserialize it +- no connector-specific metadata is required, since the store key is the rendezvous point + +Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, which makes the connector functional but not the highest-performance option for large payloads such as KV cache blocks. + +### 8. Data Flow in the Pipeline + +The end-to-end transfer model is: + +```mermaid +sequenceDiagram + participant SenderStage + participant MooncakeStoreConnector + participant MooncakeStore + participant ReceiverStage + + SenderStage->>MooncakeStoreConnector: put(from_stage, to_stage, put_key, data) + MooncakeStoreConnector->>MooncakeStoreConnector: serialize object + MooncakeStoreConnector->>MooncakeStore: put(internal_key, bytes) + MooncakeStoreConnector-->>SenderStage: (success, size, None) + + ReceiverStage->>MooncakeStoreConnector: get(from_stage, to_stage, get_key) + MooncakeStoreConnector->>MooncakeStore: get(internal_key) + MooncakeStoreConnector->>MooncakeStoreConnector: deserialize bytes + MooncakeStoreConnector-->>ReceiverStage: (data, size) +``` + +This is a store-mediated remote transfer model rather than a direct peer-to-peer transport model. + +### 9. Strengths and Trade-offs + +#### Strengths + +- Simple conceptual model: store by key, fetch by key. +- No connector-specific metadata handoff is required. +- Works naturally for remote multi-node stage transfer. +- Easy to integrate into the existing connector abstraction. + +#### Trade-offs + +- Full serialization and deserialization are always required. +- Large payloads are more expensive than in direct-memory transports. +- Runtime behavior depends on external Mooncake services being available. +- Cleanup semantics are weaker than request-scoped local buffer management. + +### 10. Important Implementation Characteristics + +#### 10.1 Cleanup Is a No-op + +`cleanup()` only logs a debug message and does not actively delete remote data. The current design assumes Mooncake-side lifecycle management rather than explicit per-request removal. + +This keeps the connector implementation simple but means request-level reclamation is not modeled inside the connector itself. + +#### 10.2 Close Releases the Store Handle + +Unlike `SharedMemoryConnector`, where `close()` is a no-op, `MooncakeStoreConnector.close()` performs a meaningful shutdown: + +1. Calls `self.store.close()` to release the Mooncake store handle. +2. Sets `self.store = None` to prevent further operations. + +This ensures the connector releases its connection to the external Mooncake service on shutdown. Errors during close are logged but do not propagate. + +#### 10.3 Health Output Uses a Shared Schema + +`health()` reports: + +- `host` +- `metadata_server` +- `master` +- metrics + +and also includes placeholder fields such as `protocol`, `pool_device`, and `pool_size` to keep a more uniform shape across connector health outputs. This reflects a system-level consistency choice rather than a store-specific design need. + +#### 10.4 Get Is Retry-Based, Not Event-Driven + +The receiver side polls the store with a short retry loop. This is simple and robust, but it means the connector is latency-sensitive to: + +- store visibility delay +- network jitter +- payload size + +If the payload is not visible within the retry window, the connector reports a timeout. + +### 11. Summary + +`MooncakeStoreConnector` is the remote, store-based connector in the OmniConnector stack. Its design is straightforward: + +- serialize the payload +- store it in Mooncake under a stage-qualified key +- fetch it by the same key on the receiver side +- deserialize it back into the original object + +This connector fills an important gap in the system: + +1. It enables cross-node transfer without requiring shared memory. +2. It keeps the stage communication model uniform. +3. It provides a simpler operational alternative to the transfer-engine-based remote connector. + +Its simplicity is also its main trade-off: it favors a clean object-transfer model over the fast-path and peer-to-peer optimizations implemented by `MooncakeTransferEngineConnector`. diff --git a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md index 96f93ef9a..798644b96 100644 --- a/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md +++ b/docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md @@ -205,3 +205,589 @@ for test-specific setup instructions. For more details on the underlying engine, refer to the [Mooncake repository](https://github.com/kvcache-ai/Mooncake). + +--- + +## Design + +### 1. Overview + +`MooncakeTransferEngineConnector` is the high-performance remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake `TransferEngine` and combines: + +- a **direct data plane** for remote memory writes +- a **ZMQ side channel** for metadata lookup, handshake, and completion signaling +- a **managed local memory pool** for both send and receive buffers + +Unlike `MooncakeStoreConnector`, which treats the backend as a distributed store, `MooncakeTransferEngineConnector` is designed as a peer-to-peer transport. Its goal is to move large stage payloads efficiently while still fitting the common `put()` / `get()` API defined by `OmniConnectorBase`. + +It is the most performance-oriented connector in the current OmniConnector family and is intended for large remote payloads such as: + +- KV cache transfer +- stage hidden-state payloads +- streaming chunk payloads +- other binary-heavy inter-stage artifacts + +### 2. Relationship with the OmniConnector System + +`MooncakeTransferEngineConnector` implements the same connector contract as the other backends: + +- `put(from_stage, to_stage, put_key, data)` +- `get(from_stage, to_stage, get_key, metadata=None)` +- `cleanup(request_id, ...)` +- `health()` +- `close()` + +It is integrated into the system through the standard connector plumbing: + +- `OmniConnectorFactory` constructs the connector from `ConnectorSpec` +- `load_omni_transfer_config()` resolves the edge-level connector configuration +- `get_connectors_config_for_stage()` and `resolve_omni_kv_config_for_stage()` inject the connector role +- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract + +The key system-level distinction is that this connector is **role-aware**: + +- sender instances expose data and listen for pull requests +- receiver instances allocate buffers and actively pull data from the sender + +### 3. Design Goals + +The connector is designed around four primary goals: + +1. **High-throughput remote transfer** + Avoid store-mediated round trips and write directly into the receiver memory region. + +2. **Fast path for raw payloads** + Support `torch.Tensor`, `bytes`, and `ManagedBuffer` without forcing all traffic through full object serialization. + +3. **Unified connector abstraction** + Preserve the same `put()` / `get()` interface used by the rest of the OmniConnector stack. + +4. **Safe lifecycle management** + Manage allocation, reuse, cleanup, and failure recovery for a registered memory pool. + +### 4. Architecture Overview + +At a high level, the connector is composed of four main subsystems: + +```mermaid +classDiagram + class OmniConnectorBase { + <> + +put(from_stage, to_stage, put_key, data) + +get(from_stage, to_stage, get_key, metadata) + +cleanup(request_id) + +health() + +close() + } + + class MooncakeTransferEngineConnector { + +supports_raw_data: bool + -engine: TransferEngine + -allocator: BufferAllocator + -pool: torch.Tensor + -zmq_ctx: zmq.Context + -_local_buffers: dict + -_sender_executor: ThreadPoolExecutor + -_listener_thread: threading.Thread + +put(...) + +get(...) + +update_sender_info(sender_host, sender_zmq_port) + +get_connection_info() + +cleanup(request_id, from_stage, to_stage) + +close() + } + + class BufferAllocator { + -total_size: int + -alignment: int + -free_blocks: list + +alloc(size) int + +free(offset, size) + } + + class ManagedBuffer { + -allocator: BufferAllocator + -offset: int + -size: int + -pool_tensor: torch.Tensor + +tensor + +as_tensor(dtype, shape) torch.Tensor + +to_bytes() bytes + +release() + } + + class TransferEngine { + +initialize(host, handshake, protocol, device_name) + +register_memory(base_ptr, size) + +batch_transfer_sync_write(remote_session, src_addrs, dst_addrs, lengths) + +unregister_memory(base_ptr) + +get_rpc_port() int + } + + class QueryRequest { + +request_id: str + } + + class QueryResponse { + +request_id: str + +data_size: int + +is_fast_path: bool + } + + class MooncakeAgentMetadata { + +remote_hostname: str + +remote_port: int + +request_id: str + +dst_addrs: list[int] + +lengths: list[int] + } + + OmniConnectorBase <|-- MooncakeTransferEngineConnector + MooncakeTransferEngineConnector *-- BufferAllocator + MooncakeTransferEngineConnector *-- TransferEngine + ManagedBuffer --> BufferAllocator : releases to + ManagedBuffer --> "1" torch.Tensor : views + MooncakeTransferEngineConnector ..> ManagedBuffer : returns / retains + MooncakeTransferEngineConnector ..> QueryRequest : decodes + MooncakeTransferEngineConnector ..> QueryResponse : encodes + MooncakeTransferEngineConnector ..> MooncakeAgentMetadata : exchanges +``` + +#### 4.1 Transfer Engine + +Mooncake `TransferEngine` is responsible for the actual data-plane transfer. It registers local memory and performs synchronous remote writes through: + +```python +batch_transfer_sync_write(...) +``` + +#### 4.2 Managed Memory Pool + +Each connector instance owns a large pre-registered memory pool: + +- CPU pinned memory when `memory_pool_device == "cpu"` +- GPU memory when `memory_pool_device == "cuda"` + +This avoids repeated memory registration and allows each transfer to allocate subranges from one long-lived pool. + +#### 4.3 Buffer Manager + +Two helper classes control local memory ownership: + +- `BufferAllocator` + Manages aligned subrange allocation and free-list merging. + +- `ManagedBuffer` + Represents one live slice of the pool and exposes: + - `.tensor` + - `.as_tensor(dtype, shape)` + - `.to_bytes()` + - `.release()` + +#### 4.4 ZMQ Side Channel + +ZMQ is used for transport coordination, not for the data payload itself. It handles: + +- metadata query from receiver to sender +- pull request submission +- completion or error signaling +- internal notification from worker threads back to the listener thread + +This split makes the control plane lightweight while keeping the bulk payload on the transfer engine data plane. + +### 5. Role Model + +#### 5.1 Sender Role + +A sender connector: + +- accepts `put()` calls +- stores live transfer-ready buffers in `_local_buffers` +- starts a ZMQ listener thread +- responds to metadata queries and pull requests from receivers + +#### 5.2 Receiver Role + +A receiver connector: + +- does not bind the sender-side ZMQ listener +- accepts `get()` calls +- allocates receive buffers from its own pool +- requests metadata or transfer service from the sender + +The role is not inferred dynamically. It is injected by the stage configuration layer: + +- incoming edge for a stage -> `role="receiver"` +- outgoing edge for a stage -> `role="sender"` + +This is important because incorrect role assignment would break initialization semantics. + +#### 5.3 Host Auto-Detection + +The `host` configuration field supports the special value `"auto"`. When set, the connector auto-detects the local IP address that would be used for external communication (via a UDP socket probe to `8.8.8.8`). If that fails, it falls back to hostname resolution, and ultimately to `127.0.0.1`. + +This is useful in environments where the operator does not want to hard-code IP addresses in the connector config. + +#### 5.4 RDMA Device Filtering + +The `device_name` configuration field allows the operator to specify which RDMA NICs to use (comma-separated, e.g. `"mlx5_0,mlx5_1"`). If not set in config, the connector also checks the `RDMA_DEVICE_NAME` environment variable. + +This is important in environments with mixed InfiniBand/RoCE NICs, where not all devices are suitable for the transfer engine. + +### 6. Local Memory Management + +#### 6.1 Memory Pool Registration + +During initialization, the connector: + +1. allocates a large pool tensor +2. records its base pointer +3. registers that memory with Mooncake +4. creates a `BufferAllocator` for subrange management + +This means every later transfer only allocates offsets inside the pre-registered pool rather than registering memory per request. + +#### 6.2 BufferAllocator + +`BufferAllocator` maintains a sorted free list of `(offset, size)` blocks and enforces alignment. Its responsibilities include: + +- aligned allocation +- freeing previously allocated blocks +- adjacent block merging +- double-free detection +- overlap detection to catch corruption + +This is a critical piece of the connector because both sender and receiver depend on long-lived pool reuse. + +#### 6.3 ManagedBuffer + +`ManagedBuffer` is the main fast-path data wrapper. It can: + +- expose the pool slice as a zero-copy 1D `uint8` tensor +- reinterpret that slice as a typed tensor +- copy out the contents as Python `bytes` +- release the slice back to the allocator + +The connector uses `ManagedBuffer` in two different ways: + +- as a send-side holder to keep the pool slice alive +- as a receive-side return type when `is_fast_path=True` + +### 7. Put Flow + +#### 7.1 High-Level Behavior + +`put()` is only valid in sender mode. Its job is to expose a payload for later remote pull by the receiver. + +The high-level flow is: + +1. validate connector state and role +2. convert the input into a pool-backed transferable representation +3. store the transfer metadata in `_local_buffers` +4. return lightweight metadata describing how the receiver can fetch the data + +#### 7.2 Payload Type Handling + +`put()` supports three payload classes: + +**A. `ManagedBuffer`** + +If the buffer belongs to the same pool, the connector can use it directly without copying. This is the most efficient path. + +If the buffer comes from a different pool, the connector falls back to a copy path. + +**B. `torch.Tensor` or `bytes`** + +These payloads are copied into the local pool and marked as fast-path data: + +- no Omni object serialization is required +- receiver can get a `ManagedBuffer` back + +**C. Generic Python object** + +Any other payload is serialized via `OmniSerializer.serialize(...)` and then copied into the pool. + +In this case: + +- `is_fast_path=False` +- the receiver will deserialize back into a Python object + +#### 7.3 Sender Metadata + +The sender returns: + +```python +{ + "source_host": self.host, + "source_port": self.zmq_port, + "data_size": size, + "is_fast_path": is_fast_path, +} +``` + +This metadata is intentionally lightweight. It tells the receiver: + +- where the sender-side control plane lives +- how large the remote transfer will be +- whether the payload should be returned as a `ManagedBuffer` or a deserialized object + +#### 7.4 Sender Buffer Table + +The sender stores each live payload in `_local_buffers` under the stage-qualified key. Each entry contains: + +- source addresses +- lengths +- the holder object +- ownership information (`should_release`) +- `is_fast_path` +- creation time + +This table is the sender-side truth source for both metadata queries and pull requests. + +### 8. Get Flow + +#### 8.1 High-Level Behavior + +`get()` runs on the receiver side and performs four steps: + +1. resolve metadata +2. allocate a destination buffer in the local pool +3. request the sender to write into that destination buffer +4. return either a `ManagedBuffer` or a deserialized object + +#### 8.2 Metadata Resolution Paths + +The `metadata` parameter in `get()` is optional. The connector supports two resolution modes depending on whether the caller supplies it. + +**With metadata** + +When the caller passes metadata, the connector uses it directly. The metadata carries: + +- `source_host` / `source_port` β€” sender ZMQ endpoint +- `data_size` β€” payload byte count +- `is_fast_path` β€” whether the receiver gets a `ManagedBuffer` or a deserialized object + +This mode is suitable when the control plane already forwards the sender's `put()` output to the receiver. + +**Without metadata** + +When `get(metadata=None)` is called, the connector queries the sender over ZMQ to discover the same fields (`data_size`, `is_fast_path`). The caller must first call: + +```python +update_sender_info(sender_host, sender_zmq_port) +``` + +so that the connector knows where to send the query. + +This mode is suitable for polling-based flows (e.g. KV transfer, async chunk transfer) where the receiver does not have metadata from the control plane. + +#### 8.3 Destination Allocation + +Once metadata is resolved, the receiver: + +1. allocates a subrange from its own local pool +2. wraps it in a `ManagedBuffer` +3. builds a `MooncakeAgentMetadata` request containing: + - receiver hostname + - receiver RPC port + - request ID + - destination addresses + - transfer lengths + +This tells the sender exactly where to write the incoming data. + +#### 8.4 Transfer Completion + +The receiver then sends the pull request over ZMQ and waits for: + +- `TRANS_DONE` +- or `TRANS_ERROR` + +If the transfer succeeds: + +- for `is_fast_path=True`, the receiver returns `(ManagedBuffer, size)` +- for `is_fast_path=False`, the receiver copies to bytes, deserializes, releases the buffer, and returns `(object, size)` + +### 9. Sender-Side Listener Design + +#### 9.1 Listener Thread + +In sender mode, the connector starts `_zmq_listener_loop()` after initialization. The listener: + +- binds `tcp://{host}:{zmq_port}` +- receives incoming requests +- uses a poller for socket events and internal notifications +- periodically reclaims stale buffers + +If the bind fails, initialization fails immediately. The code does not silently downgrade the connector role. + +#### 9.2 Worker Thread Pool + +The listener hands work to `_sender_executor` so that the listener thread does not block on transfer work. + +There are two request types: + +- metadata query -> `_handle_query_request(...)` +- transfer request -> `_handle_pull_request(...)` + +#### 9.3 Query Handling + +For metadata queries, the sender looks up the request ID in `_local_buffers` and returns: + +- data size +- fast-path flag + +This supports consumers that only know the sender endpoint but not the original sender metadata. + +#### 9.4 Pull Handling + +For a transfer request, the sender: + +1. locates the source addresses in `_local_buffers` +2. constructs the remote session identifier +3. calls `batch_transfer_sync_write(...)` +4. replies `TRANS_DONE` or `TRANS_ERROR` + +On success, the sender immediately calls `cleanup(meta.request_id)` and frees the producer-side buffer if it owns it. + +This choice is important: it makes the connector effectively a **single-consumer transfer model** for each successful put/get pair. + +### 10. Fast Path Semantics + +This connector explicitly advertises: + +```python +supports_raw_data = True +``` + +That means it can move raw payloads without forcing everything through the Omni object serializer. + +#### Fast Path + +For `torch.Tensor`, `bytes`, or pool-local `ManagedBuffer`: + +- sender returns `is_fast_path=True` +- receiver returns a `ManagedBuffer` +- caller is responsible for calling `release()` + +This avoids an unnecessary copy on the receiver side. + +#### Serialized Object Path + +For arbitrary Python objects: + +- sender serializes the payload +- receiver converts the receive buffer to bytes +- receiver deserializes the object +- receive buffer is released internally + +This preserves a uniform object-oriented API while still allowing optimized raw-data transport when possible. + +### 11. Failure Handling and Cleanup + +#### 11.1 Timeouts and Socket Recovery + +The receiver caches ZMQ REQ sockets per thread, but invalidates them after failures. This avoids reusing sockets that may be stuck in a bad state after timeout or receive errors. + +Timeout is scaled based on payload size: + +- a base timeout +- plus additional time for large payloads + +This is intended to reduce false timeouts for large remote writes. + +#### 11.2 Stale Buffer Reclamation + +The sender periodically reclaims old entries from `_local_buffers` using a TTL policy. This protects the memory pool from permanent leaks if a receiver crashes or never consumes a prepared payload. + +This is a practical recovery mechanism, although the code notes that TTL cleanup can still race with very long-running in-flight transfers. + +#### 11.3 Connector Shutdown + +`close()` is a full resource teardown routine. It: + +- stops the listener thread +- shuts down the worker executor +- releases all pending buffers +- closes cached sockets +- unregisters memory from the engine when supported +- terminates the ZMQ context +- drops the pool reference + +This makes `MooncakeTransferEngineConnector` the most lifecycle-aware connector in the current connector family. + +### 12. Current Implementation Constraints + +The current code documents several important topology constraints. + +#### 12.1 One Sender to One Receiver per Successful Transfer + +After a successful transfer, the sender-side buffer is cleaned up immediately. This means the same prepared payload is not retained for multiple independent receivers. + +#### 12.2 One Receiver to One Active Sender Endpoint + +The receiver only stores one `(sender_host, sender_zmq_port)` pair through `update_sender_info(...)`. So the metadata-query mode is currently single-sender at a time. + +#### 12.3 Explicit Buffer Ownership Matters + +When the connector allocates a pool slice internally, it is responsible for releasing it. When a caller passes an externally owned `ManagedBuffer`, the connector keeps it alive for transfer but does not assume ownership of its eventual release. + +These constraints are consistent with the current implementation and should be treated as design assumptions rather than incidental behavior. + +### 13. Data Flow in the Pipeline + +The end-to-end sender/receiver interaction is: + +```mermaid +sequenceDiagram + participant SenderStage + participant SenderConnector + participant ReceiverConnector + participant ReceiverStage + + SenderStage->>SenderConnector: put(from_stage, to_stage, put_key, data) + SenderConnector->>SenderConnector: place payload in local memory pool + SenderConnector-->>SenderStage: metadata(source_host, source_port, data_size, is_fast_path) + + ReceiverStage->>ReceiverConnector: get(..., metadata) + ReceiverConnector->>ReceiverConnector: allocate destination buffer + ReceiverConnector->>SenderConnector: ZMQ pull request with dst addr + SenderConnector->>ReceiverConnector: TransferEngine remote write + SenderConnector-->>ReceiverConnector: TRANS_DONE + ReceiverConnector-->>ReceiverStage: ManagedBuffer or deserialized object +``` + +For metadata-less polling, the flow simply adds a metadata query step before the pull request. + +### 14. Strengths and Trade-offs + +#### Strengths + +- Best remote-transfer design in the current connector stack for large payloads. +- Supports raw-data fast path. +- Keeps stage communication under the same connector abstraction. +- Includes real lifecycle and memory-pool management. +- Works for both stage payload transfer and KV transfer scenarios. + +#### Trade-offs + +- More complex than the store-based connector. +- Correctness depends on role injection and endpoint coordination. +- Caller must release fast-path receive buffers. +- Current implementation is optimized for single-consumer transfer semantics. + +### 15. Summary + +`MooncakeTransferEngineConnector` is the high-performance peer-to-peer transport in the OmniConnector system. Its design combines: + +- a registered memory pool +- a safe subrange allocator +- a ZMQ control plane +- a Mooncake transfer-engine data plane + +This allows the connector to support both: + +1. a **fast path** for raw tensors and bytes +2. a **generic object path** for arbitrary Python payloads + +Within vLLM-Omni, it is the connector that most directly targets performance-sensitive remote transfer, especially for large payloads and KV cache movement. Its additional complexity is deliberate: it is the connector that turns the generic OmniConnector abstraction into a transport capable of efficient remote memory movement rather than simple object storage. diff --git a/docs/design/feature/omni_connectors/shared_memory_connector.md b/docs/design/feature/omni_connectors/shared_memory_connector.md index eb91a889a..5b35014f2 100644 --- a/docs/design/feature/omni_connectors/shared_memory_connector.md +++ b/docs/design/feature/omni_connectors/shared_memory_connector.md @@ -7,9 +7,10 @@ auto-configured when no explicit connector is specified for an edge. ## How It Works -- Small payloads (< threshold): serialized and passed inline in metadata. -- Large payloads (>= threshold): stored in shared memory; the SHM name is - returned in metadata. +All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM +segment name is returned in metadata. The configuration exposes a +`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current +implementation always uses shared memory regardless of payload size. ## Configuration @@ -25,3 +26,234 @@ runtime: ## Notes - Auto-mode uses SharedMemoryConnector if no connector is declared for an edge. + +--- + +## Design + +### 1. Overview + +`SharedMemoryConnector` is the default same-node connector in `vllm_omni/distributed/omni_connectors`. It is designed for stage-to-stage transfer when producer and consumer processes run on the same host and can share `/dev/shm`. + +The connector provides a unified `put()` / `get()` API for arbitrary Python objects while keeping the control plane lightweight: + +- The payload is serialized by the connector. +- The serialized bytes are placed in shared memory. +- The queue/control plane only carries a small metadata handle. + +This makes `SharedMemoryConnector` the simplest connector in the OmniConnector family and the default fallback when an edge does not explicitly configure another backend. + +### 2. Relationship with the OmniConnector System + +`SharedMemoryConnector` implements `OmniConnectorBase`, so it follows the same lifecycle and API contract as the other connectors: + +- `put(from_stage, to_stage, put_key, data)` +- `get(from_stage, to_stage, get_key, metadata=None)` +- `cleanup(request_id)` +- `health()` +- `close()` + +Within the larger system: + +- `load_omni_transfer_config()` automatically fills missing edges with `SharedMemoryConnector`. +- Callers interact with the connector exclusively through the `put()` / `get()` / `cleanup()` contract β€” the connector does not require caller-specific logic. + +Compared with the remote Mooncake-based connectors, `SharedMemoryConnector` is intentionally minimal and local-only. + +### 3. Design Goals + +The connector is built around the following goals: + +- **Low-friction local transfer** for single-node multi-process pipelines. +- **Unified object semantics** for arbitrary Python payloads. +- **Small control-plane overhead** by passing only metadata through queues. +- **Zero external dependencies** beyond Python shared memory and the existing stage utilities. + +It is not intended to provide cross-node transfer, RDMA, or raw tensor zero-copy semantics across processes. + +### 4. Core Design + +#### 4.1 Serialization Model + +`SharedMemoryConnector` always starts from a Python object and serializes it through the shared Omni serializer: + +```python +payload = self.serialize_obj(data) +``` + +This keeps the connector behavior consistent with the rest of the connector stack: + +- producer code does not need connector-specific serialization logic +- consumer code always receives the original object after deserialization +- the connector can reuse the same serializer used by other backends + +#### 4.2 Shared Memory as the Data Plane + +The actual data plane is a shared-memory segment created by: + +- `shm_write_bytes(...)` +- `shm_read_bytes(...)` + +The connector stores a small metadata object such as: + +```python +{ + "shm": {"name": ..., "size": ...}, + "size": ... +} +``` + +This metadata is passed over the control plane and allows the downstream stage to locate the shared-memory segment. + +#### 4.3 Locking Model + +To avoid races between the producer and consumer, the connector uses a lock file per request: + +```text +/dev/shm/shm_{put_key}_lockfile.lock +``` + +Locking is done with `fcntl.flock`: + +- producer uses `LOCK_EX` +- consumer uses `LOCK_EX` + +Both sides acquire an exclusive lock. This ensures that the shared-memory segment is not read while it is still being written and makes the handoff safer in a multi-process environment. + +### 5. Put / Get Flow + +#### 5.1 Producer Flow: `put()` + +The producer-side flow is: + +1. Serialize the input object to bytes. +2. Compute the payload size. +3. Acquire the per-request lock file. +4. Write the bytes into shared memory. +5. Return lightweight metadata to the caller. + +The returned tuple is: + +```python +(success, serialized_size, metadata) +``` + +where `metadata` contains the shared-memory handle needed by the consumer. + +#### 5.2 Consumer Flow: `get(metadata=...)` + +The primary consumer path is metadata-driven: + +1. Extract the shared-memory handle from `metadata`. +2. Acquire the exclusive lock. +3. Read the raw bytes from shared memory. +4. Deserialize the bytes back into the original Python object. +5. Remove the lock file if it still exists. + +This is the path used by the current stage-to-stage connector flow. + +#### 5.3 Compatibility Flow: `get(metadata=None)` + +The connector also keeps a compatibility path for callers that only know the key: + +1. Attempt to open the shared-memory segment by name via `SharedMemory(name=get_key)`. +2. If the segment exists and has non-zero size, acquire the exclusive lock and read the bytes. +3. Deserialize the bytes and return the object. + +If the segment does not exist or any exception occurs, the call returns `None` immediately. There is no retry loop in this path -- it is a single-attempt open. + +This path is mainly for older code paths and is not the preferred mode for the current connector pipeline. + +### 6. Key Implementation Characteristics + +#### 6.1 Threshold Exists, but the Current Code Always Uses SHM + +The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses: + +```python +if True: + ... +``` + +inside `put()`, which means the current code path always writes to shared memory. + +So the design still suggests a future split between: + +- small payloads inline +- large payloads in shared memory + +but the current behavior is effectively: + +- all payloads go through shared memory + +This should be documented because it affects real runtime behavior. + +#### 6.2 Cleanup Is Currently Passive + +`cleanup()` is currently a no-op. The intended assumption is: + +- the consumer reads the segment +- the underlying shared-memory helpers unlink it + +If the consumer never executes `get()`, the shared-memory segment may remain allocated. This means the connector relies on the normal success path for resource reclamation. + +#### 6.3 Close Is Currently Minimal + +`close()` is also a no-op. There is no connector-owned background thread, socket, or memory pool to tear down, so the lifecycle is simple. The trade-off is that `close()` does not scan or recover leaked shared-memory resources. + +### 7. Data Flow in the Pipeline + +The typical flow with `SharedMemoryConnector` is: + +```mermaid +sequenceDiagram + participant SenderStage + participant SharedMemoryConnector + participant QueueOrControlPlane + participant ReceiverStage + + SenderStage->>SharedMemoryConnector: put(from_stage, to_stage, put_key, data) + SharedMemoryConnector->>SharedMemoryConnector: serialize object + SharedMemoryConnector->>SharedMemoryConnector: write bytes to /dev/shm + SharedMemoryConnector-->>SenderStage: metadata {shm, size} + SenderStage->>QueueOrControlPlane: forward connector metadata + QueueOrControlPlane->>ReceiverStage: task + connector metadata + ReceiverStage->>SharedMemoryConnector: get(from_stage, to_stage, get_key, metadata) + SharedMemoryConnector->>SharedMemoryConnector: read bytes from /dev/shm + SharedMemoryConnector->>SharedMemoryConnector: deserialize object + SharedMemoryConnector-->>ReceiverStage: (data, size) +``` + +This is a classic split-control-plane / data-plane design, but constrained to a single host. + +### 8. Strengths and Trade-offs + +#### Strengths + +- Very simple deployment model. +- No external service dependency. +- Fits naturally into the existing queue-driven orchestration flow. +- Good default for local multi-process pipelines. + +#### Trade-offs + +- Same-node only. +- Full object serialization and deserialization are still required. +- Resource cleanup depends on the normal consumer path. +- Shared memory capacity is limited by host configuration. + +### 9. Summary + +`SharedMemoryConnector` is the baseline local transport for the OmniConnector system. Its design is intentionally straightforward: + +- serialize object +- place bytes in shared memory +- pass metadata through the control plane +- deserialize on the receiving side + +It plays two important roles in vLLM-Omni: + +1. It is the simplest production-ready connector for same-node stage pipelines. +2. It serves as the automatic fallback connector when no explicit edge transport is configured. + +Although the current implementation is deliberately minimal, it provides the foundation for reliable local connector semantics and keeps the stage communication model uniform across the system. diff --git a/docs/design/feature/omni_connectors/yuanrong_connector.md b/docs/design/feature/omni_connectors/yuanrong_connector.md index 88989d31a..12127ba3e 100644 --- a/docs/design/feature/omni_connectors/yuanrong_connector.md +++ b/docs/design/feature/omni_connectors/yuanrong_connector.md @@ -94,8 +94,265 @@ stage_args: Parameters: - host: datasystem worker host. -- port: datasystem worker port. +- port: datasystem worker port (default: `35001` if omitted; the example above uses `31501` to match the worker startup command). - get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout). For more details, refer to the [Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem). + +--- + +## Design + +### 1. Overview + +`YuanrongConnector` is the Datasystem-based remote connector in `vllm_omni/distributed/omni_connectors`. It uses Yuanrong Datasystem's distributed key-value client as the transport backend and exposes the same `put()` / `get()` interface as the other OmniConnectors. + +Like `MooncakeStoreConnector`, it is a store-oriented remote connector rather than a direct peer-to-peer transport. Its role is to let stage payloads move across nodes through a deterministic key-based storage abstraction while keeping the rest of the pipeline on the common connector API. + +It is intended for deployments that already use Yuanrong Datasystem and want a remote connector that integrates with the existing OmniConnector configuration and orchestration model. + +### 2. Relationship with the OmniConnector System + +`YuanrongConnector` implements `OmniConnectorBase`, so it participates in the same connector lifecycle as the other backends: + +- `OmniConnectorFactory` constructs it from a `ConnectorSpec` +- stage edge configuration is resolved by `load_omni_transfer_config()` +- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract + +This means the connector is not exposed directly to stage logic. Stages only interact with the generic connector contract, and the backend choice remains a configuration concern. + +### 3. Design Goals + +The connector is built around the following goals: + +1. **Cross-node payload transfer through Datasystem** + Reuse Yuanrong Datasystem as the remote exchange medium for stage data. + +2. **Uniform object transfer semantics** + Allow arbitrary Python objects to be transmitted through the shared Omni serializer. + +3. **Minimal connector-specific control plane** + Use deterministic keys so that consumers can retrieve data without an extra transport metadata handoff. + +4. **Operational reuse of existing infrastructure** + Fit into environments that already deploy Yuanrong Datasystem workers and etcd. + +The connector is not designed for direct remote-memory writes or tensor-specific fast-path transfer. + +### 4. Core Design + +#### 4.1 Store-Oriented Transfer Model + +`YuanrongConnector` treats the transport backend as a distributed object store: + +1. serialize the Python payload +2. build a deterministic connector key +3. write the serialized bytes into Datasystem +4. read the bytes back on the receiver side +5. deserialize them into the original object + +This is the same broad architectural class as `MooncakeStoreConnector`, but implemented on top of Yuanrong Datasystem APIs instead of Mooncake store APIs. + +#### 4.2 Deterministic Keying + +Unlike the default `_make_key()` in `OmniConnectorBase`, `YuanrongConnector` defines its own key format: + +```text +{request_id}:{from_stage}_{to_stage} +``` + +This has two design implications: + +- the request identifier remains the primary lookup handle +- stage routing information is embedded in the key so that the same logical request ID can safely appear on different edges + +The explicit override also makes the key format easier to align with Datasystem-side debugging and operational inspection. + +#### 4.3 No Extra Metadata Hand-off + +`put()` returns: + +```python +(success, serialized_size, None) +``` + +and does not generate connector-specific metadata. + +This design works because the receiver can reconstruct the exact same key from: + +- `get_key` +- `from_stage` +- `to_stage` + +As a result, the connector does not require a separate side-channel metadata handoff. + +### 5. Initialization + +#### 5.1 Datasystem Client Dependency + +The connector requires the Datasystem Python bindings to expose: + +- `KVClient` +- `SetParam` +- `WriteMode` + +If any of these symbols are unavailable, connector construction fails immediately with `ImportError`. This keeps configuration errors explicit and avoids a partially initialized runtime. + +#### 5.2 Client Setup + +During `_init_client()`, the connector: + +1. reads `host` and `port` +2. creates `KVClient(host, port)` +3. calls `client.init()` + +At construction time it also creates a `SetParam` and fixes: + +```python +self.set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT +``` + +This means the connector has a stable write policy for all writes and does not currently expose write-mode selection as a higher-level connector option. + +### 6. Put / Get Flow + +#### 6.1 Producer Flow: `put()` + +The producer-side flow is: + +1. verify that the Datasystem client has been initialized +2. serialize the input object with the shared Omni serializer +3. build the Datasystem key using the connector-specific `_make_key()` +4. call `client.set(key, serialized_data, self.set_param.write_mode)` +5. update metrics and return success + +The returned metadata is always `None`, because the Datasystem key itself is the lookup contract between producer and consumer. + +#### 6.2 Consumer Flow: `get()` + +The consumer-side flow is: + +1. verify that the Datasystem client has been initialized +2. rebuild the same key with `from_stage`, `to_stage`, and `get_key` +3. call: + +```python +client.get([key], False, self.get_sub_timeout_ms) +``` + +4. take the first returned element if present +5. deserialize it and return `(data, payload_size)` + +If the returned list is empty or contains no data for the key, `get()` returns `None`. + +### 7. Timeout and Retrieval Semantics + +The connector uses `get_sub_timeout_ms` as its read timeout. Unlike `MooncakeStoreConnector`, which performs an explicit retry loop in Python, `YuanrongConnector` delegates waiting behavior more directly to the Datasystem client call. + +This leads to a slightly different retrieval model: + +- `MooncakeStoreConnector`: retry-oriented polling in connector code +- `YuanrongConnector`: single client call with backend-managed timeout behavior + +From the connector API perspective the result is the same, but operationally the waiting behavior is more dependent on Datasystem client semantics. + +### 8. Integration with Stage Communication + +All callers use the connector through the same `put()` / `get()` contract: + +- the sender calls `put()` to serialize and store the payload +- the receiver calls `get()` to retrieve and deserialize it +- no connector-specific metadata is required, since the Datasystem key is the rendezvous point + +Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, and there is no raw tensor fast path, which makes the connector functional but not optimized for the largest payloads. + +### 9. Data Flow in the Pipeline + +The end-to-end transfer model is: + +```mermaid +sequenceDiagram + participant SenderStage + participant YuanrongConnector + participant Datasystem + participant ReceiverStage + + SenderStage->>YuanrongConnector: put(from_stage, to_stage, put_key, data) + YuanrongConnector->>YuanrongConnector: serialize object + YuanrongConnector->>Datasystem: set(key, bytes) + YuanrongConnector-->>SenderStage: (success, size, None) + + ReceiverStage->>YuanrongConnector: get(from_stage, to_stage, get_key) + YuanrongConnector->>Datasystem: get([key], timeout) + YuanrongConnector->>YuanrongConnector: deserialize bytes + YuanrongConnector-->>ReceiverStage: (data, size) +``` + +This is a store-mediated remote connector design with deterministic key lookup and no explicit side-channel metadata exchange. + +### 10. Strengths and Trade-offs + +#### Strengths + +- Reuses existing Yuanrong Datasystem infrastructure. +- Keeps the connector contract simple and uniform. +- Requires no connector-specific metadata handoff. +- Suitable for remote stage transfer in Datasystem-based deployments. + +#### Trade-offs + +- Always pays full serialization and deserialization cost. +- Does not support raw bytes / tensor fast-path semantics. +- Depends on external Datasystem worker availability. +- Retrieval and timeout behavior are partly delegated to the backend client. + +### 11. Important Implementation Characteristics + +#### 11.1 Cleanup Is a No-op + +`cleanup()` only logs a debug message and does not explicitly remove data from Datasystem. The current design assumes backend-side lifecycle or garbage collection rather than request-scoped delete semantics inside the connector. + +This mirrors the same design trade-off seen in other store-based connectors: simplicity over explicit per-request data reclamation. + +#### 11.2 Close Only Releases the Local Client Handle + +`close()` does not perform a remote shutdown. It simply clears the local `client` reference and marks the connector as closed from the process perspective: + +```python +self.client = None +``` + +This is appropriate for a client-based store connector, but it also means that backend resource lifecycle remains outside the connector's control. + +#### 11.3 Error and Timeout Metrics Are Coarse-Grained + +The connector tracks: + +- `puts` +- `gets` +- `bytes_transferred` +- `errors` +- `timeouts` + +In the current code, failed `put()` increments `errors`, while `get()` exceptions increment `timeouts`. This is operationally useful, but it does not distinguish between: + +- backend timeout +- not-found result +- transport failure +- deserialization failure + +So the metrics should be read as high-level health indicators, not detailed root-cause diagnostics. + +### 12. Summary + +`YuanrongConnector` is the Datasystem-backed remote connector in the OmniConnector family. Its design is straightforward: + +- serialize payloads +- store them under a deterministic stage-qualified key +- retrieve them by the same key +- deserialize them on the receiving side + +Within vLLM-Omni, it provides a clean integration point for Yuanrong Datasystem-based deployments while preserving the same connector abstraction used by the rest of the pipeline. + +Its design priorities are simplicity, infrastructure reuse, and API consistency rather than direct-memory transport optimization. diff --git a/docs/design/feature/sequence_parallel.md b/docs/design/feature/sequence_parallel.md index 18477e61a..3b899ffe3 100644 --- a/docs/design/feature/sequence_parallel.md +++ b/docs/design/feature/sequence_parallel.md @@ -55,6 +55,8 @@ The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The - Tensor operations happen at `nn.Module` boundaries - Predictable sharding/gathering patterns +This is the ideal approach for integrating sequence parallelism into new models, as it is easier to maintain and ensure compatibility with other types of acceleration. + **How it works:** 1. Declare `_sp_plan` dict in your transformer class 2. Framework automatically applies hooks when `sequence_parallel_size > 1` @@ -201,6 +203,36 @@ class TransformerWithRoPE(nn.Module): } ``` +**Pattern 3: Shard RoPE for Dual Stream Attention** +In some cases, different streams in attention may need to handle sequence parallelism differently. For example, we may want to shard the image embeddings, while replicating the text embeddings to correctly configure joint attention. + +```python +class DualStreamTransformer(nn.Module): + """ + Dual-stream model where we need to replicate the text components, but shard + the image components to correctly handle sequence parallelism. + """ + _sp_plan = { + # In this case, the rope_preparer returns a tuple of len 4, where the + # first 2 items correspond to the text, and the second 2 correspond to + # visual inputs, so we only shard the second. + "rope_preparer": { + # Outputs 0, 1 (text) - NOT sharded (replicated) + # Outputs 2, 3 (image) - sharded + 2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_cos + 3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_sin + }, + # Shard transformer block INPUT + "transformer_blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Gather at output + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } +``` + +NOTE: be careful to test adequately when refactoring classes that take this style of plan, as changing the order of the return values will break sequence parallelism. + ### API Reference **SequenceParallelInput Parameters:** @@ -240,7 +272,7 @@ class TransformerWithRoPE(nn.Module): ## Approach 2: Intrusive Modification (For Complex Cases) -For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used. +For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. **When to use:** @@ -253,18 +285,15 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather def forward(self, hidden_states, ...): if self.parallel_config.sequence_parallel_size > 1: - # hidden_states = sp_shard(hidden_states, dim=1) # ... computation ... if self.parallel_config.sequence_parallel_size > 1: output = sp_gather(output, dim=1) - # return output ``` -Note that currently, `sp_shard` / `sp_gather` do *not* automatically manage the `_sp_shard_depth`; you need to be careful to manage it yourself. --- @@ -439,6 +468,7 @@ Complete examples in the codebase: | Model | Path | Pattern | Notes | |-------|------|---------|-------| +| **LongCat** | `vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py` | Dual-stream | Text components replicated, image components sharded | | **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Dual-stream + preprocessing | auto_pad, separate RoPE | | **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Dual-Transformer + RoPE | Video transformer | | **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Unified sequence | Concatenated input | @@ -453,7 +483,7 @@ Complete examples in the codebase: Adding Sequence Parallel support to a transformer: 1. βœ… **Choose approach** - Use `_sp_plan` for standard cases, intrusive modification for complex cases -2. βœ… **Identify sharding boundaries** - Where should tensors be split/gathered? +2. βœ… **Identify sharding boundaries** - Where should tensors be split/gathered? And which module boundaries need to be moved to facilitate this? 3. βœ… **Extract inline operations** - Move `torch.cat`, `pad_sequence`, etc. to submodules 4. βœ… **Define `_sp_plan`** - Declare shard/gather points as class attribute 5. βœ… **Use `auto_pad` for variable lengths** - Support non-uniform sequences diff --git a/docs/features/comfyui.md b/docs/features/comfyui.md index 3b7a9551e..ff766cbb3 100644 --- a/docs/features/comfyui.md +++ b/docs/features/comfyui.md @@ -47,7 +47,8 @@ If no, check your shell running the ComfyUI process. There may be some error mes This extension offers the following nodes based on the output modalities: - **Generate Image** for text-to-image and image-to-image tasks -- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks +- **Generate Video** for text-to-video and image-to-video tasks +- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks - **TTS** and **TTS Voice Clone** for TTS tasks This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**) @@ -59,7 +60,7 @@ To build a simple workflow yourself, - Drag a generation node onto the canvas. - Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio** -- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node. +- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**. - If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**. - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node. - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters. diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py index 4e840280b..b4c928936 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/docs/mkdocs/hooks/generate_examples.py @@ -132,9 +132,19 @@ class Example: ".dat", ".db", ".sqlite", # Data files + ".pyc", + ".pyo", + ".pyd", # Python compiled + ".npy", + ".npz", + ".pkl", + ".pickle", # Serialized data } + excluded_dirs = {"__pycache__", ".git", "node_modules", ".tox", ".mypy_cache"} def is_other_file(file: Path) -> bool: + if any(part in excluded_dirs for part in file.parts): + return False return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions return [file for file in self.path.rglob("*") if is_other_file(file)] diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 56b5a0f6b..4c709425e 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -49,6 +49,8 @@ th { |`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` | |`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | |`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | +|`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | +|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ## List of Supported Models for NPU @@ -64,6 +66,7 @@ th { |--------------|--------|-------------------| | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`| +| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 2f0695298..46b9154bb 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -1,6 +1,9 @@ # Speech API -vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models. +vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include: + +- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz. +- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz. Each server instance runs a single model (specified at startup via `vllm serve --omni`). @@ -9,13 +12,22 @@ Each server instance runs a single model (specified at startup via `vllm serve < ### Start the Server ```bash -# CustomVoice model (predefined speakers) +# Qwen3-TTS: CustomVoice model (predefined speakers) vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \ --omni \ --port 8091 \ --trust-remote-code \ --enforce-eager + +# Fish Speech S2 Pro +vllm-omni serve fishaudio/s2-pro \ + --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager \ + --gpu-memory-utilization 0.9 ``` ### Generate Speech @@ -164,6 +176,42 @@ curl -X POST http://localhost:8091/v1/audio/voices \ -F "name=custom_voice_1" ``` +## Streaming Text Input (WebSocket) + +The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and generates audio per sentence as boundaries are detected. + +> Note: text input is always streamed incrementally. Audio output remains sentence-scoped: +> use `stream_audio=false` for one binary frame per sentence, or `stream_audio=true` for one or more PCM chunks per sentence. + +### WebSocket Protocol + +Client -> Server: + +| Message | Description | +|---------|-------------| +| `{"type": "session.config", ...}` | Session configuration (sent once, first message) | +| `{"type": "input.text", "text": "..."}` | Text chunk | +| `{"type": "input.done"}` | End of input, flushes remaining buffer | + +Server -> Client: + +| Message | Description | +|---------|-------------| +| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence | +| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) | +| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence | +| `{"type": "session.done", "total_sentences": N}` | Session complete | +| `{"type": "error", "message": "..."}` | Non-fatal error | + +### Session Config Parameters + +All REST API parameters are supported, plus: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket | +| `split_granularity` | string | "sentence" | Text splitting granularity | + ```bash DELETE /v1/audio/voices/{name} @@ -281,6 +329,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \ ## Supported Models +### Qwen3-TTS + | Model | Task Type | Description | |-------|-----------|-------------| | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control | @@ -289,6 +339,14 @@ curl -X POST http://localhost:8091/v1/audio/speech \ | `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant | | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning | +### Fish Speech S2 Pro + +| Model | Description | +|-------|-------------| +| `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. | + +Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details. + ## Error Responses ### 400 Bad Request diff --git a/docs/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index bd2a58799..505f9627f 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -37,6 +37,7 @@ The following table shows which models are currently supported by parallelism me | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | βœ… | βœ… | N/A | | **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | βœ… | ❌ | N/A | | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | βœ… | βœ… | ❌ | N/A | +| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | βœ… | ❌ | N/A | | **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | βœ… | ❌ | βœ… | !!! note "TP Limitations for Diffusion Models" diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 462ab8b18..7f32e3eff 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -69,6 +69,7 @@ The following table shows which models are currently supported by each accelerat | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | βœ… | ❌ | ❌ | βœ… | βœ… | ❌ | | **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | βœ… | βœ… | ❌ | | **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | βœ… | ❌ | ❌ | ❌ | βœ… | ❌ | +| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | βœ… | ❌ | ❌ | ❌ | βœ… | ❌ | ### VideoGen diff --git a/docs/user_guide/examples/offline_inference/cosyvoice3.md b/docs/user_guide/examples/offline_inference/cosyvoice3.md new file mode 100644 index 000000000..d912f1c62 --- /dev/null +++ b/docs/user_guide/examples/offline_inference/cosyvoice3.md @@ -0,0 +1,74 @@ +# CosyVoice3 + +Source . + + +## Setup + +Install dependencies: +``` +uv pip install -e . +``` + +> **Note:** This includes required libraries such as `librosa`, `soundfile`, +> `onnxruntime`, `x-transformers`, and `einops` via +> `requirements/common.txt` and platform-specific requirements files. + +Download the model snapshot: +``` +from huggingface_hub import snapshot_download +snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B') +``` + +Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`: +```json +{ + "model_type": "cosyvoice3", + "architectures": [ + "CosyVoice3Model" + ] +} +``` + +> **Why `config.json` is required:** +> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping. +> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class. +> If no `config.json` is present, model type cannot be inferred automatically. +> If your downloaded checkpoint already includes a valid `config.json` with +> `model_type: "cosyvoice3"`, this manual step can be skipped. + +Run the offline verification script: +``` +python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \ + --model pretrained_models/Fun-CosyVoice3-0.5B \ + --tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN +``` + +## Implementation Overview + +CosyVoice3 runs as a **2-stage Omni pipeline**. + +- **Stage 0** (`talker`) converts text + prompt audio to speech tokens. +- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN. + +Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`. + +- `CosyVoice3MultiModalProcessor` builds the multimodal inputs: + - Tokenizes `prompt` and `prompt_text`. + - Extracts speech tokens and mel features from the prompt audio. + - Extracts a speaker embedding. +- `CosyVoice3Model` implements both stages: + - Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features. + - Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform. + +Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`. + +- Stage 0 emits latent speech tokens. +- Stage 1 consumes them via `custom_process_input_func` and outputs audio. + +## Example materials + +??? abstract "verify_e2e_cosyvoice.py" + ``````py + --8<-- "examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py" + `````` diff --git a/docs/user_guide/examples/offline_inference/fish_speech.md b/docs/user_guide/examples/offline_inference/fish_speech.md new file mode 100644 index 000000000..dd9066afd --- /dev/null +++ b/docs/user_guide/examples/offline_inference/fish_speech.md @@ -0,0 +1,97 @@ +# Fish Speech S2 Pro + +Source . + + +This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio. + +## Model Overview + +[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses: + +- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone) +- **Fast AR**: Predicts residual codebook tokens from semantic tokens +- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform + +The model produces high-quality speech with voice cloning capabilities. + +## Prerequisites + +Install the `fish-speech` package for the DAC codec: + +```bash +pip install fish-speech +``` + +## Setup + +Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. + +## Quick Start + +### Text-Only Synthesis + +```bash +python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system." +``` + +Generated audio files are saved to `output_audio/` by default. + +### Voice Cloning + +Provide a reference audio file and its transcript to clone a voice: + +```bash +python end2end.py \ + --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." +``` + +## Streaming Mode + +Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config): + +```bash +python end2end.py --text "Hello, this is a test." --streaming +``` + +Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes. + +## Configuration + +| Argument | Default | Description | +|----------|---------|-------------| +| `--text` | `"Hello, this is a test..."` | Text to synthesize | +| `--ref-audio` | None | Path to reference audio for voice cloning | +| `--ref-text` | None | Transcript of the reference audio | +| `--model` | `fishaudio/s2-pro` | HuggingFace model path | +| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML | +| `--output-dir` | `output_audio` | Output directory for WAV files | +| `--streaming` | False | Enable streaming via AsyncOmni | +| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) | + +## Notes + +- Output audio is 44.1 kHz mono WAV. +- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory. +- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message. + +## Example materials + +??? abstract "diag_codebook.py" + ``````py + --8<-- "examples/offline_inference/fish_speech/diag_codebook.py" + `````` +??? abstract "diag_reference_compare.py" + ``````py + --8<-- "examples/offline_inference/fish_speech/diag_reference_compare.py" + `````` +??? abstract "end2end.py" + ``````py + --8<-- "examples/offline_inference/fish_speech/end2end.py" + `````` +??? abstract "standalone_test.py" + ``````py + --8<-- "examples/offline_inference/fish_speech/standalone_test.py" + `````` diff --git a/docs/user_guide/examples/online_serving/fish_speech.md b/docs/user_guide/examples/online_serving/fish_speech.md new file mode 100644 index 000000000..2c76b1bca --- /dev/null +++ b/docs/user_guide/examples/online_serving/fish_speech.md @@ -0,0 +1,194 @@ +# Fish Speech S2 Pro + +Source . + + +## Model + +| Model | Description | +|-------|-------------| +| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) | + +## Gradio Demo + +An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support. + +```bash +# Option 1: Launch server + Gradio together +./run_gradio_demo.sh + +# Option 2: If server is already running +python gradio_demo.py --api-base http://localhost:8091 +``` + +Then open http://localhost:7860 in your browser. + +Features: + +- Text-to-speech synthesis +- Voice cloning from uploaded audio or URL +- Streaming mode (progressive PCM playback) + +## Launch the Server + +```bash +vllm-omni serve fishaudio/s2-pro \ + --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager \ + --gpu-memory-utilization 0.9 +``` + +Or use the convenience script: + +```bash +./run_server.sh +``` + +## Send TTS Request + +### Using curl + +```bash +# Basic TTS +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav" + }' --output output.wav +``` + +### Voice Cloning + +Provide a reference audio (URL or base64 data URL) and its transcript: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is a cloned voice.", + "voice": "default", + "ref_audio": "https://example.com/reference.wav", + "ref_text": "Transcript of the reference audio." + }' --output cloned.wav +``` + +### Using Python + +```python +import httpx + +# Basic TTS +response = httpx.post( + "http://localhost:8091/v1/audio/speech", + json={ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav", + }, + timeout=300.0, +) + +with open("output.wav", "wb") as f: + f.write(response.content) +``` + +### Using the CLI Client + +```bash +cd examples/online_serving/fish_speech + +# Basic TTS +python speech_client.py --text "Hello, how are you?" + +# Voice cloning +python speech_client.py \ + --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." + +# Streaming PCM output +python speech_client.py --text "Hello world" --stream --output output.pcm +``` + +The CLI client supports: + +- `--api-base`: API base URL (default: `http://localhost:8091`) +- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`) +- `--text`: Text to synthesize (required) +- `--ref-audio`: Reference audio for voice cloning (local path or URL) +- `--ref-text`: Transcript of the reference audio +- `--stream`: Enable streaming (PCM output) +- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav) +- `--output` (or `-o`): Output file path + +## Streaming + +Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "default", + "stream": true, + "response_format": "pcm" + }' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 - +``` + +**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz). + +## API Parameters + +Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation. + +Key parameters for Fish Speech: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `input` | string | **required** | Text to synthesize | +| `voice` | string | "default" | Voice name (use "default" for Fish Speech) | +| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus | +| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning | +| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) | +| `max_new_tokens` | int | 4096 | Maximum tokens to generate | +| `stream` | bool | false | Stream raw PCM chunks | + +## Prerequisites + +Install the `fish-speech` package for the DAC codec: + +```bash +pip install fish-speech +``` + +## Troubleshooting + +1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder +2. **Connection refused**: Ensure the server is running on the correct port +3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings +4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM + +## Example materials + +??? abstract "gradio_demo.py" + ``````py + --8<-- "examples/online_serving/fish_speech/gradio_demo.py" + `````` +??? abstract "run_gradio_demo.sh" + ``````sh + --8<-- "examples/online_serving/fish_speech/run_gradio_demo.sh" + `````` +??? abstract "run_server.sh" + ``````sh + --8<-- "examples/online_serving/fish_speech/run_server.sh" + `````` +??? abstract "speech_client.py" + ``````py + --8<-- "examples/online_serving/fish_speech/speech_client.py" + `````` diff --git a/examples/offline_inference/cosyvoice3/README.md b/examples/offline_inference/cosyvoice3/README.md new file mode 100644 index 000000000..895d3f660 --- /dev/null +++ b/examples/offline_inference/cosyvoice3/README.md @@ -0,0 +1,64 @@ +# CosyVoice3 + +## Setup + +Install dependencies: +``` +uv pip install -e . +``` + +> **Note:** This includes required libraries such as `librosa`, `soundfile`, +> `onnxruntime`, `x-transformers`, and `einops` via +> `requirements/common.txt` and platform-specific requirements files. + +Download the model snapshot: +``` +from huggingface_hub import snapshot_download +snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B') +``` + +Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`: +```json +{ + "model_type": "cosyvoice3", + "architectures": [ + "CosyVoice3Model" + ] +} +``` + +> **Why `config.json` is required:** +> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping. +> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class. +> If no `config.json` is present, model type cannot be inferred automatically. +> If your downloaded checkpoint already includes a valid `config.json` with +> `model_type: "cosyvoice3"`, this manual step can be skipped. + +Run the offline verification script: +``` +python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \ + --model pretrained_models/Fun-CosyVoice3-0.5B \ + --tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN +``` + +## Implementation Overview + +CosyVoice3 runs as a **2-stage Omni pipeline**. + +- **Stage 0** (`talker`) converts text + prompt audio to speech tokens. +- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN. + +Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`. + +- `CosyVoice3MultiModalProcessor` builds the multimodal inputs: + - Tokenizes `prompt` and `prompt_text`. + - Extracts speech tokens and mel features from the prompt audio. + - Extracts a speaker embedding. +- `CosyVoice3Model` implements both stages: + - Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features. + - Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform. + +Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`. + +- Stage 0 emits latent speech tokens. +- Stage 1 consumes them via `custom_process_input_func` and outputs audio. diff --git a/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py b/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py new file mode 100644 index 000000000..c962aac80 --- /dev/null +++ b/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import argparse +import os +from pathlib import Path + +import librosa +import numpy as np +import soundfile as sf +from vllm import SamplingParams +from vllm.assets.audio import AudioAsset + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config +from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer +from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token + + +def _ensure_mel_filters_asset() -> None: + repo_root = Path(__file__).resolve().parents[3] + filters_path = repo_root / "vllm_omni" / "model_executor" / "models" / "cosyvoice3" / "assets" / "mel_filters.npz" + if filters_path.exists(): + return + + source_url = "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz" + raise FileNotFoundError( + "Missing CosyVoice3 mel filter asset:\n" + f" {filters_path}\n" + "Download it with:\n" + f" mkdir -p {filters_path.parent} && " + f"curl -L {source_url} -o {filters_path}" + ) + + +def run_e2e(): + parser = argparse.ArgumentParser() + # ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512 + parser.add_argument( + "--model", + type=str, + required=True, + help="Path to CosyVoice3 model directory (e.g., pretrained_models/Fun-CosyVoice3-0.5B/).", + ) + parser.add_argument("--stage-config", type=str, default="vllm_omni/model_executor/stage_configs/cosyvoice3.yaml") + parser.add_argument("--prompt", type=str, default="Hello, this is a test of the CosyVoice system capability.") + parser.add_argument( + "--prompt-text", + type=str, + default="You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?", + ) + parser.add_argument("--audio-path", type=str, default="prompt.wav") + parser.add_argument( + "--tokenizer", + type=str, + required=True, + help="Path to tokenizer directory (e.g., /CosyVoice-BlankEN).", + ) + args = parser.parse_args() + _ensure_mel_filters_asset() + # Ensure tokenizer directory exists + if not os.path.exists(args.tokenizer): + raise FileNotFoundError(f"{args.tokenizer} does not exist!") + + # Ensure stage config exists + if not os.path.exists(args.stage_config): + raise FileNotFoundError(f"{args.stage_config} does not exist!") + + print(f"Initializing cosyvoice E2E with model={args.model}") + + # Initialize Omni + # This spins up the engine(s) based on the stage config + # We pass trust_remote_code=True same as Qwen examples + omni = Omni( + model=args.model, + stage_configs_path=args.stage_config, + trust_remote_code=True, + tokenizer=args.tokenizer, + log_stats=True, + ) + + sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1} + + print("Model initialized. Preparing inputs...") + if args.audio_path: + if not os.path.exists(args.audio_path): + raise FileNotFoundError(f"Audio file not found: {args.audio_path}") + # Load at native sample rate + audio_signal, sr = librosa.load(args.audio_path, sr=None) + + # Validate sample rate before processing (similar to original CosyVoice) + min_sr = 16000 + if sr < min_sr: + raise ValueError( + f"Audio sample rate {sr} Hz is too low. " + f"Minimum required: {min_sr} Hz. " + f"Please provide audio with sample rate >= {min_sr} Hz." + ) + + audio_data = (audio_signal.astype(np.float32), sr) + else: + audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate + + prompts = { + "prompt": args.prompt, + "multi_modal_data": { + "audio": audio_data, + }, + "mm_processor_kwargs": { + "prompt_text": args.prompt_text, + "sample_rate": audio_data[1], + }, + } + + print(f"Generating for prompt: {args.prompt}") + + config = CosyVoice3Config() + tokenizer = get_qwen_tokenizer( + token_path=args.tokenizer, + skip_special_tokens=config.skip_special_tokens, + version=config.version, + ) + _, text_token_len = extract_text_token(args.prompt, tokenizer, config.allowed_special) + base_len = int(text_token_len) + min_len = int(base_len * config.min_token_text_ratio) + max_len = int(base_len * config.max_token_text_ratio) + + # Build SamplingParams for each stage (GPT, S2Mel, Vocoder) + gpt_sampling = SamplingParams( + temperature=1.0, + top_p=sampling_cfg["top_p"], + top_k=sampling_cfg["top_k"], + repetition_penalty=2.0, + min_tokens=min_len, + max_tokens=max_len, + stop_token_ids=[sampling_cfg["eos_token_id"]], + # allowed_token_ids=list(range(6561+3)), + detokenize=False, + ) + # Not used + s2mel_sampling = SamplingParams( + temperature=1.0, + top_p=1.0, + top_k=-1, + repetition_penalty=2.0, + max_tokens=256, + detokenize=False, + ) + + sampling_params_list = [gpt_sampling, s2mel_sampling] + + # Start profiling (requires VLLM_TORCH_PROFILER_DIR env var) + if os.environ.get("VLLM_TORCH_PROFILER_DIR"): + print("Starting profiler...") + omni.start_profile() + + # Generate (Omni orchestrator requires a per-stage SamplingParams list) + outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2])) + + # Stop profiling and get results + if os.environ.get("VLLM_TORCH_PROFILER_DIR"): + print("Stopping profiler...") + profile_results = omni.stop_profile() + print(f"Profile traces saved to: {profile_results}") + + print(outputs) + # Verify outputs + print(f"Received {len(outputs)} outputs.") + for i, output in enumerate(outputs): + try: + ro_list = output.request_output or [] + if not ro_list: + print("No request_output found.") + continue + + for ro in ro_list: + # Multimodal output may be attached to RequestOutput or CompletionOutput. + mm = getattr(ro, "multimodal_output", None) + if not mm and ro.outputs: + mm = getattr(ro.outputs[0], "multimodal_output", None) + + if mm: + print(f"Multimodal output keys: {mm.keys()}") + if "audio" in mm: + audio_out = mm["audio"] + print(f"Generated Audio Shape: {audio_out.shape}") + out_path = f"output_{i}.wav" + sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050) + print(f"Saved audio to {out_path}") + else: + print("No multimodal output found.") + except Exception as e: + print(f"Error inspecting output: {e}") + omni.close() + + +if __name__ == "__main__": + run_e2e() diff --git a/examples/offline_inference/fish_speech/README.md b/examples/offline_inference/fish_speech/README.md new file mode 100644 index 000000000..9af853096 --- /dev/null +++ b/examples/offline_inference/fish_speech/README.md @@ -0,0 +1,75 @@ +# Fish Speech S2 Pro + +This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio. + +## Model Overview + +[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses: + +- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone) +- **Fast AR**: Predicts residual codebook tokens from semantic tokens +- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform + +The model produces high-quality speech with voice cloning capabilities. + +## Prerequisites + +Install the `fish-speech` package for the DAC codec: + +```bash +pip install fish-speech +``` + +## Setup + +Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. + +## Quick Start + +### Text-Only Synthesis + +```bash +python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system." +``` + +Generated audio files are saved to `output_audio/` by default. + +### Voice Cloning + +Provide a reference audio file and its transcript to clone a voice: + +```bash +python end2end.py \ + --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." +``` + +## Streaming Mode + +Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config): + +```bash +python end2end.py --text "Hello, this is a test." --streaming +``` + +Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes. + +## Configuration + +| Argument | Default | Description | +|----------|---------|-------------| +| `--text` | `"Hello, this is a test..."` | Text to synthesize | +| `--ref-audio` | None | Path to reference audio for voice cloning | +| `--ref-text` | None | Transcript of the reference audio | +| `--model` | `fishaudio/s2-pro` | HuggingFace model path | +| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML | +| `--output-dir` | `output_audio` | Output directory for WAV files | +| `--streaming` | False | Enable streaming via AsyncOmni | +| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) | + +## Notes + +- Output audio is 44.1 kHz mono WAV. +- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory. +- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message. diff --git a/examples/offline_inference/fish_speech/end2end.py b/examples/offline_inference/fish_speech/end2end.py new file mode 100644 index 000000000..9e95f7e5c --- /dev/null +++ b/examples/offline_inference/fish_speech/end2end.py @@ -0,0 +1,252 @@ +"""Offline inference demo for Fish Speech S2 Pro via vLLM Omni. + +Generates speech from text using the fishaudio/s2-pro model. +Supports both text-only synthesis and voice cloning with reference audio. + +Usage: + # Text-only synthesis + python end2end.py --text "Hello, this is a test." + + # Voice cloning with reference audio + python end2end.py --text "Hello, this is a test." --ref-audio ref.wav --ref-text "Reference text." + + # Streaming mode + python end2end.py --text "Hello, this is a test." --streaming +""" + +import asyncio +import logging +import os +import time + +import soundfile as sf +import torch + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import AsyncOmni, Omni + +logger = logging.getLogger(__name__) + +DEFAULT_MODEL = "fishaudio/s2-pro" +DEFAULT_STAGE_CONFIG = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "..", + "vllm_omni", + "model_executor", + "stage_configs", + "fish_speech_s2_pro.yaml", +) + + +def build_prompt( + text: str, + ref_audio_path: str | None = None, + ref_text: str | None = None, + model_name: str = DEFAULT_MODEL, +) -> dict: + """Build a prompt for Fish Speech S2 Pro. + + Uses Qwen3 chat template format with <|voice|> modality marker: + <|im_start|>user\n<|speaker:0|>text<|im_end|>\n<|im_start|>assistant\n<|voice|> + + The model then generates semantic tokens autoregressively. + """ + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + + # Build user message with speaker tag prefix. + user_text = f"<|speaker:0|>{text}" + + # Use chat template for user turn + assistant generation prompt. + messages = [{"role": "user", "content": user_text}] + prompt_ids = tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + ) + + # Append <|voice|> token (151673) to signal voice generation mode. + voice_token_id = tokenizer.encode("<|voice|>", add_special_tokens=False) + prompt_ids = prompt_ids + voice_token_id + + additional_information = { + "text": [text], + "max_new_tokens": [4096], + } + + if ref_audio_path: + additional_information["ref_audio"] = [ref_audio_path] + if ref_text: + additional_information["ref_text"] = [ref_text] + + return { + "prompt_token_ids": prompt_ids, + "additional_information": additional_information, + } + + +def _save_wav(output_dir: str, request_id: str, mm: dict) -> None: + """Concatenate audio chunks and write to a wav file.""" + audio_data = mm["audio"] + sr_raw = mm["sr"] + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) + audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data + out_wav = os.path.join(output_dir, f"output_{request_id}.wav") + sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV") + logger.info("Request %s: saved audio to %s (sr=%d)", request_id, out_wav, sr) + + +def main(args): + """Run offline inference with Omni.""" + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG + model_name = args.model or DEFAULT_MODEL + + inputs = [ + build_prompt( + text=args.text, + ref_audio_path=args.ref_audio, + ref_text=args.ref_text, + model_name=model_name, + ) + ] + + omni = Omni( + model=model_name, + stage_configs_path=stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + + t_start = time.perf_counter() + for stage_outputs in omni.generate(inputs): + for output in stage_outputs.request_output: + _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output) + t_end = time.perf_counter() + logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000) + + +async def main_streaming(args): + """Run offline inference with AsyncOmni for streaming audio chunks.""" + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG + model_name = args.model or DEFAULT_MODEL + + prompt = build_prompt( + text=args.text, + ref_audio_path=args.ref_audio, + ref_text=args.ref_text, + model_name=model_name, + ) + + omni = AsyncOmni( + model=model_name, + stage_configs_path=stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + + request_id = "0" + t_start = time.perf_counter() + t_prev = t_start + chunk_idx = 0 + + async for stage_output in omni.generate(prompt, request_id=request_id): + mm = stage_output.request_output.outputs[0].multimodal_output + if not stage_output.finished: + t_now = time.perf_counter() + audio = mm.get("audio") + n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1) + dt_ms = (t_now - t_prev) * 1000 + ttfa_ms = (t_now - t_start) * 1000 + if chunk_idx == 0: + logger.info("Request %s: chunk %d samples=%d TTFA=%.1fms", request_id, chunk_idx, n, ttfa_ms) + else: + logger.info("Request %s: chunk %d samples=%d inter_chunk=%.1fms", request_id, chunk_idx, n, dt_ms) + t_prev = t_now + chunk_idx += 1 + else: + t_end = time.perf_counter() + total_ms = (t_end - t_start) * 1000 + logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx) + _save_wav(output_dir, request_id, mm) + + +def parse_args(): + parser = FlexibleArgumentParser( + description="Fish Speech S2 Pro offline inference with vLLM Omni", + ) + parser.add_argument( + "--text", + type=str, + default="Hello, this is a test of the Fish Speech text to speech system.", + help="Text to synthesize.", + ) + parser.add_argument( + "--ref-audio", + type=str, + default=None, + help="Path to reference audio for voice cloning.", + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help="Reference text corresponding to the reference audio.", + ) + parser.add_argument( + "--model", + type=str, + default=None, + help=f"HuggingFace model path (default: {DEFAULT_MODEL}).", + ) + parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to stage configs YAML.", + ) + parser.add_argument( + "--output-dir", + default="output_audio", + help="Output directory for generated wav files.", + ) + parser.add_argument( + "--log-stats", + action="store_true", + default=False, + help="Enable writing detailed statistics.", + ) + parser.add_argument( + "--stage-init-timeout", + type=int, + default=600, + help="Timeout for initializing stages in seconds.", + ) + parser.add_argument( + "--streaming", + action="store_true", + default=False, + help="Stream audio chunks as they arrive via AsyncOmni.", + ) + return parser.parse_args() + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + args = parse_args() + if args.streaming: + asyncio.run(main_streaming(args)) + else: + main(args) diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index ff2dda010..2513627e8 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -106,6 +106,23 @@ python text_to_image.py \ --seed 42 ``` +### Flux.2-dev Models +To start Flux.2-dev with a single GPU, cpu-offload must be enabled because the total size of its weights exceeds the 80GB memory capacity of the GPU. +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model black-forest-labs/FLUX.2-dev \ + --prompt "a lovely bunny holding a sign that says 'vllm-omni'" \ + --seed 42 \ + --tensor-parallel-size 1 \ + --num-images-per-prompt 1 \ + --num-inference-steps 50 \ + --guidance-scale 4.0 \ + --height 1024 \ + --width 1024 \ + --enable-cpu-offload \ + --output flux2-dev.png +``` + ### Key Arguments **Common arguments:** diff --git a/examples/online_serving/fish_speech/README.md b/examples/online_serving/fish_speech/README.md new file mode 100644 index 000000000..ae968d3ba --- /dev/null +++ b/examples/online_serving/fish_speech/README.md @@ -0,0 +1,172 @@ +# Fish Speech S2 Pro + +## Model + +| Model | Description | +|-------|-------------| +| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) | + +## Gradio Demo + +An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support. + +```bash +# Option 1: Launch server + Gradio together +./run_gradio_demo.sh + +# Option 2: If server is already running +python gradio_demo.py --api-base http://localhost:8091 +``` + +Then open http://localhost:7860 in your browser. + +Features: + +- Text-to-speech synthesis +- Voice cloning from uploaded audio or URL +- Streaming mode (progressive PCM playback) + +## Launch the Server + +```bash +vllm-omni serve fishaudio/s2-pro \ + --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \ + --omni \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager \ + --gpu-memory-utilization 0.9 +``` + +Or use the convenience script: + +```bash +./run_server.sh +``` + +## Send TTS Request + +### Using curl + +```bash +# Basic TTS +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav" + }' --output output.wav +``` + +### Voice Cloning + +Provide a reference audio (URL or base64 data URL) and its transcript: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is a cloned voice.", + "voice": "default", + "ref_audio": "https://example.com/reference.wav", + "ref_text": "Transcript of the reference audio." + }' --output cloned.wav +``` + +### Using Python + +```python +import httpx + +# Basic TTS +response = httpx.post( + "http://localhost:8091/v1/audio/speech", + json={ + "input": "Hello, how are you?", + "voice": "default", + "response_format": "wav", + }, + timeout=300.0, +) + +with open("output.wav", "wb") as f: + f.write(response.content) +``` + +### Using the CLI Client + +```bash +cd examples/online_serving/fish_speech + +# Basic TTS +python speech_client.py --text "Hello, how are you?" + +# Voice cloning +python speech_client.py \ + --text "Hello, this is a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of the reference audio." + +# Streaming PCM output +python speech_client.py --text "Hello world" --stream --output output.pcm +``` + +The CLI client supports: + +- `--api-base`: API base URL (default: `http://localhost:8091`) +- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`) +- `--text`: Text to synthesize (required) +- `--ref-audio`: Reference audio for voice cloning (local path or URL) +- `--ref-text`: Transcript of the reference audio +- `--stream`: Enable streaming (PCM output) +- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav) +- `--output` (or `-o`): Output file path + +## Streaming + +Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, how are you?", + "voice": "default", + "stream": true, + "response_format": "pcm" + }' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 - +``` + +**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz). + +## API Parameters + +Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation. + +Key parameters for Fish Speech: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `input` | string | **required** | Text to synthesize | +| `voice` | string | "default" | Voice name (use "default" for Fish Speech) | +| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus | +| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning | +| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) | +| `max_new_tokens` | int | 4096 | Maximum tokens to generate | +| `stream` | bool | false | Stream raw PCM chunks | + +## Prerequisites + +Install the `fish-speech` package for the DAC codec: + +```bash +pip install fish-speech +``` + +## Troubleshooting + +1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder +2. **Connection refused**: Ensure the server is running on the correct port +3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings +4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM diff --git a/examples/online_serving/fish_speech/gradio_demo.py b/examples/online_serving/fish_speech/gradio_demo.py new file mode 100644 index 000000000..a8e51b274 --- /dev/null +++ b/examples/online_serving/fish_speech/gradio_demo.py @@ -0,0 +1,272 @@ +"""Gradio demo for Fish Speech S2 Pro online serving via /v1/audio/speech API. + +Supports: + - Text-to-speech synthesis + - Voice cloning from reference audio (upload or URL) + - Streaming (progressive PCM chunks) and non-streaming modes + +Usage: + # Start the server first (see run_server.sh), then: + python gradio_demo.py --api-base http://localhost:8091 + + # Or use run_gradio_demo.sh to start both server and demo together. +""" + +import argparse +import base64 +import io + +import gradio as gr +import httpx +import numpy as np +import soundfile as sf + +PCM_SAMPLE_RATE = 44100 +DEFAULT_API_BASE = "http://localhost:8091" + + +def encode_audio_to_base64(audio_data: tuple) -> str: + """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL.""" + sample_rate, audio_np = audio_data + if audio_np.dtype != np.int16: + if audio_np.dtype in (np.float32, np.float64): + audio_np = np.clip(audio_np, -1.0, 1.0) + audio_np = (audio_np * 32767).astype(np.int16) + else: + audio_np = audio_np.astype(np.int16) + buf = io.BytesIO() + sf.write(buf, audio_np, sample_rate, format="WAV") + wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8") + return f"data:audio/wav;base64,{wav_b64}" + + +def build_payload( + text: str, + ref_audio: tuple | None, + ref_audio_url: str, + ref_text: str, + response_format: str = "wav", + stream: bool = False, +) -> dict: + """Build the /v1/audio/speech request payload.""" + if not text or not text.strip(): + raise gr.Error("Please enter text to synthesize.") + + payload: dict = { + "input": text.strip(), + "voice": "default", + "response_format": "pcm" if stream else response_format, + "stream": stream, + } + + # Voice cloning: ref_audio takes priority over URL. + ref_url_stripped = ref_audio_url.strip() if ref_audio_url else "" + if ref_audio is not None: + payload["ref_audio"] = encode_audio_to_base64(ref_audio) + elif ref_url_stripped: + payload["ref_audio"] = ref_url_stripped + + if "ref_audio" in payload: + if not ref_text or not ref_text.strip(): + raise gr.Error("Voice cloning requires a transcript of the reference audio.") + payload["ref_text"] = ref_text.strip() + + return payload + + +def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format): + """Non-streaming: call /v1/audio/speech and return full audio.""" + payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False) + + try: + with httpx.Client(timeout=300.0) as client: + resp = client.post( + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") + + if resp.status_code != 200: + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + + content_type = resp.headers.get("content-type", "") + if "application/json" in content_type: + try: + raise gr.Error(f"Server error: {resp.json()}") + except ValueError: + pass + + try: + if response_format == "pcm": + audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0 + return (PCM_SAMPLE_RATE, audio_np) + audio_np, sample_rate = sf.read(io.BytesIO(resp.content)) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + return (sample_rate, audio_np.astype(np.float32)) + except Exception as e: + raise gr.Error(f"Failed to decode audio: {e}") + + +def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5): + """Streaming: yield PCM chunks for progressive playback.""" + payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True) + + min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds) + prebuffer_chunks = 2 + pending: list[np.ndarray] = [] + pending_len = 0 + chunks_yielded = 0 + prebuffer: list[np.ndarray] = [] + + try: + leftover = b"" + with httpx.Client(timeout=300.0) as client: + with client.stream( + "POST", + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) as resp: + if resp.status_code != 200: + resp.read() + raise gr.Error(f"Server error ({resp.status_code}): {resp.text}") + for chunk in resp.iter_bytes(): + if not chunk: + continue + raw = leftover + chunk + usable = len(raw) - (len(raw) % 2) + leftover = raw[usable:] + if usable == 0: + continue + float_samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0 + pending.append(float_samples) + pending_len += len(float_samples) + if pending_len >= min_chunk_samples: + audio_chunk = np.concatenate(pending) + pending.clear() + pending_len = 0 + if chunks_yielded < prebuffer_chunks: + prebuffer.append(audio_chunk) + chunks_yielded += 1 + if chunks_yielded == prebuffer_chunks: + yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer)) + prebuffer.clear() + else: + yield (PCM_SAMPLE_RATE, audio_chunk) + remaining = prebuffer + pending + if remaining: + yield (PCM_SAMPLE_RATE, np.concatenate(remaining)) + except httpx.TimeoutException: + raise gr.Error("Request timed out. The server may be busy.") + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") + + +def on_stream_change(stream: bool): + """Lock format to PCM when streaming.""" + if stream: + return gr.update(value="pcm", interactive=False) + return gr.update(value="wav", interactive=True) + + +def build_interface(api_base: str, stream_chunk_seconds: float = 0.5): + """Build the Gradio interface.""" + with gr.Blocks(title="Fish Speech S2 Pro Demo") as demo: + gr.Markdown("# Fish Speech S2 Pro - Text to Speech") + gr.Markdown(f"**Server:** `{api_base}` | **Model:** fishaudio/s2-pro | **Output:** 44.1kHz") + + with gr.Row(): + with gr.Column(scale=3): + text_input = gr.Textbox( + label="Text to Synthesize", + placeholder="Enter text here...", + lines=4, + ) + + with gr.Accordion("Voice Cloning (optional)", open=False): + gr.Markdown( + "Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice." + ) + ref_audio = gr.Audio( + label="Reference Audio", + type="numpy", + sources=["upload", "microphone"], + ) + ref_audio_url = gr.Textbox( + label="Reference Audio URL (alternative to upload)", + placeholder="https://example.com/reference.wav", + lines=1, + ) + ref_text = gr.Textbox( + label="Reference Audio Transcript (required for cloning)", + placeholder="Exact transcript of the reference audio...", + lines=2, + ) + + with gr.Row(): + response_format = gr.Dropdown( + choices=["wav", "mp3", "flac", "pcm"], + value="wav", + label="Audio Format", + scale=1, + ) + stream_checkbox = gr.Checkbox( + label="Stream output", + value=False, + info="Progressive PCM streaming", + scale=1, + ) + + generate_btn = gr.Button("Generate Speech", variant="primary", size="lg") + + with gr.Column(scale=2): + audio_output = gr.Audio( + label="Generated Audio", + interactive=False, + streaming=True, + autoplay=True, + ) + gr.Markdown( + "### About\n" + "- **Fish Speech S2 Pro** by FishAudio: 4B dual-AR model\n" + "- **Voice cloning**: upload 10-30s reference + transcript\n" + "- **Streaming**: real-time PCM output\n" + "- **44.1kHz** output via DAC codec" + ) + + stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format]) + + all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format] + + def dispatch(stream_enabled, *args): + if stream_enabled: + yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds) + else: + yield generate_speech(api_base, *args) + + generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output]) + demo.queue() + return demo + + +def main(): + parser = argparse.ArgumentParser(description="Gradio demo for Fish Speech S2 Pro") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})") + parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)") + parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)") + parser.add_argument("--share", action="store_true", help="Share publicly") + parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds per streaming chunk") + args = parser.parse_args() + + print(f"Connecting to vLLM server at: {args.api_base}") + demo = build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds) + demo.launch(server_name=args.host, server_port=args.port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/fish_speech/run_gradio_demo.sh b/examples/online_serving/fish_speech/run_gradio_demo.sh new file mode 100755 index 000000000..98a696644 --- /dev/null +++ b/examples/online_serving/fish_speech/run_gradio_demo.sh @@ -0,0 +1,48 @@ +#!/bin/bash +# Launch Fish Speech S2 Pro server + Gradio demo together. +# +# Usage: +# ./run_gradio_demo.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh + +set -e + +MODEL="${MODEL:-fishaudio/s2-pro}" +PORT="${PORT:-8091}" +GRADIO_PORT="${GRADIO_PORT:-7860}" +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)" + +echo "Starting Fish Speech S2 Pro server (port $PORT)..." +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +vllm-omni serve "$MODEL" \ + --stage-configs-path "$REPO_ROOT/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --gpu-memory-utilization 0.9 \ + --trust-remote-code \ + --enforce-eager \ + --omni & +SERVER_PID=$! + +cleanup() { + echo "Stopping server (PID $SERVER_PID)..." + kill $SERVER_PID 2>/dev/null + wait $SERVER_PID 2>/dev/null +} +trap cleanup EXIT + +# Wait for server to be ready. +echo "Waiting for server to start..." +for i in $(seq 1 120); do + if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then + echo "Server ready." + break + fi + sleep 2 +done + +echo "Starting Gradio demo (port $GRADIO_PORT)..." +python "$SCRIPT_DIR/gradio_demo.py" \ + --api-base "http://localhost:$PORT" \ + --port "$GRADIO_PORT" diff --git a/examples/online_serving/fish_speech/run_server.sh b/examples/online_serving/fish_speech/run_server.sh new file mode 100755 index 000000000..59c09c7fe --- /dev/null +++ b/examples/online_serving/fish_speech/run_server.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Launch vLLM-Omni server for Fish Speech S2 Pro +# +# Usage: +# ./run_server.sh +# CUDA_VISIBLE_DEVICES=0 ./run_server.sh + +set -e + +MODEL="${MODEL:-fishaudio/s2-pro}" +PORT="${PORT:-8091}" + +echo "Starting Fish Speech S2 Pro server with model: $MODEL" + +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +vllm-omni serve "$MODEL" \ + --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \ + --host 0.0.0.0 \ + --port "$PORT" \ + --gpu-memory-utilization 0.9 \ + --trust-remote-code \ + --enforce-eager \ + --omni diff --git a/examples/online_serving/fish_speech/speech_client.py b/examples/online_serving/fish_speech/speech_client.py new file mode 100644 index 000000000..e3ace37e5 --- /dev/null +++ b/examples/online_serving/fish_speech/speech_client.py @@ -0,0 +1,129 @@ +"""Client for Fish Speech S2 Pro via /v1/audio/speech endpoint. + +Examples: + # Basic TTS + python speech_client.py --text "Hello, how are you?" + + # Voice cloning + python speech_client.py --text "Hello, how are you?" \ + --ref-audio ref.wav --ref-text "This is the reference transcript." + + # Streaming PCM output + python speech_client.py --text "Hello world" --stream --output output.pcm +""" + +import argparse +import base64 +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" + + +def encode_audio_to_base64(audio_path: str) -> str: + """Encode a local audio file to base64 data URL.""" + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + ext = audio_path.lower().rsplit(".", 1)[-1] + mime_map = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"} + mime_type = mime_map.get(ext, "audio/wav") + with open(audio_path, "rb") as f: + audio_b64 = base64.b64encode(f.read()).decode("utf-8") + return f"data:{mime_type};base64,{audio_b64}" + + +def run_tts(args) -> None: + """Generate speech via /v1/audio/speech API.""" + payload = { + "model": args.model, + "input": args.text, + "voice": "default", + "response_format": args.response_format, + } + + # Voice cloning parameters. + if args.ref_audio: + if args.ref_audio.startswith(("http://", "https://")): + payload["ref_audio"] = args.ref_audio + else: + payload["ref_audio"] = encode_audio_to_base64(args.ref_audio) + if args.ref_text: + payload["ref_text"] = args.ref_text + + if args.stream: + payload["stream"] = True + payload["response_format"] = "pcm" + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + if args.ref_audio: + print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}") + print("Generating audio...") + + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + if args.stream: + output_path = args.output or "output.pcm" + with httpx.Client(timeout=300.0) as client: + with client.stream("POST", api_url, json=payload, headers=headers) as resp: + if resp.status_code != 200: + print(f"Error: {resp.status_code}") + print(resp.read().decode()) + return + total_bytes = 0 + with open(output_path, "wb") as f: + for chunk in resp.iter_bytes(): + f.write(chunk) + total_bytes += len(chunk) + print(f"Streamed {total_bytes} bytes to: {output_path}") + else: + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + + output_path = args.output or "output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="Fish Speech S2 Pro TTS client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default="fishaudio/s2-pro", help="Model name") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)") + parser.add_argument("--ref-text", default=None, help="Transcript of reference audio") + parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "mp3", "flac", "pcm", "aac", "opus"], + help="Audio format (default: wav)", + ) + parser.add_argument("--output", "-o", default=None, help="Output file path") + args = parser.parse_args() + run_tts(args) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/qwen3_tts/README.md b/examples/online_serving/qwen3_tts/README.md index 99442c29e..4709c9d42 100644 --- a/examples/online_serving/qwen3_tts/README.md +++ b/examples/online_serving/qwen3_tts/README.md @@ -349,6 +349,42 @@ curl -X POST http://localhost:8091/v1/audio/speech \ - `speed` adjustment is not supported when streaming. - Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`). +## Streaming Text Input (WebSocket) + +The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally, buffers and splits it at sentence boundaries, and generates audio per sentence. + +When `stream_audio=true`, each sentence is emitted as `audio.start`, one or more binary PCM frames, and `audio.done`. + +### Quick Start + +```bash +python streaming_speech_client.py \ + --text "Hello world. How are you? I am fine." + +python streaming_speech_client.py \ + --text "Hello world. How are you? I am fine." \ + --simulate-stt --stt-delay 0.1 +``` + +### WebSocket Protocol + +Client -> Server: + +```jsonc +{"type": "session.config", "voice": "Vivian", "task_type": "CustomVoice", "language": "Auto", "split_granularity": "sentence", "stream_audio": true, "response_format": "pcm"} +{"type": "input.text", "text": "Hello, how are you? "} +{"type": "input.done"} +``` + +Server -> Client: + +```jsonc +{"type": "audio.start", "sentence_index": 0, "sentence_text": "Hello, how are you?", "format": "pcm", "sample_rate": 24000} +// binary PCM frame(s) +{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false} +{"type": "session.done", "total_sentences": 1} +``` + ## Limitations - **Single request**: Batch processing is not yet optimized for online serving. diff --git a/examples/online_serving/qwen3_tts/streaming_speech_client.py b/examples/online_serving/qwen3_tts/streaming_speech_client.py new file mode 100644 index 000000000..785c6a0e8 --- /dev/null +++ b/examples/online_serving/qwen3_tts/streaming_speech_client.py @@ -0,0 +1,249 @@ +"""WebSocket client for streaming text-input TTS. + +Connects to the /v1/audio/speech/stream endpoint, sends text incrementally +(simulating real-time STT output), and saves per-sentence audio files. + +Usage: + # Send full text at once + python streaming_speech_client.py --text "Hello world. How are you? I am fine." + + # Simulate STT: send text word-by-word with delay + python streaming_speech_client.py \ + --text "Hello world. How are you? I am fine." \ + --simulate-stt --stt-delay 0.1 + + # VoiceDesign task + python streaming_speech_client.py \ + --text "Today is a great day. The weather is nice." \ + --task-type VoiceDesign \ + --instructions "A cheerful young female voice" + + # Base task (voice cloning) + python streaming_speech_client.py \ + --text "Hello world. How are you?" \ + --task-type Base \ + --ref-audio /path/to/reference.wav \ + --ref-text "Transcript of reference audio" + +Requirements: + pip install websockets +""" + +import argparse +import asyncio +import json +import os + +try: + import websockets +except ImportError: + print("Please install websockets: pip install websockets") + raise SystemExit(1) + + +async def stream_tts( + url: str, + text: str, + config: dict, + output_dir: str, + simulate_stt: bool = False, + stt_delay: float = 0.1, +) -> None: + """Connect to the streaming TTS endpoint and process audio responses.""" + os.makedirs(output_dir, exist_ok=True) + + async with websockets.connect(url) as ws: + # 1. Send session config + config_msg = {"type": "session.config", **config} + await ws.send(json.dumps(config_msg)) + print(f"Sent session config: {config}") + + # 2. Send text (either all at once or word-by-word) + async def send_text(): + if simulate_stt: + words = text.split(" ") + for i, word in enumerate(words): + chunk = word + (" " if i < len(words) - 1 else "") + await ws.send( + json.dumps( + { + "type": "input.text", + "text": chunk, + } + ) + ) + print(f" Sent: {chunk!r}") + await asyncio.sleep(stt_delay) + else: + await ws.send( + json.dumps( + { + "type": "input.text", + "text": text, + } + ) + ) + print(f"Sent full text: {text!r}") + + # 3. Signal end of input + await ws.send(json.dumps({"type": "input.done"})) + print("Sent input.done") + + # Run sender and receiver concurrently + sender_task = asyncio.create_task(send_text()) + + response_format = config.get("response_format", "wav") + current_sentence_index = 0 + current_chunks: list[bytes] = [] + + try: + while True: + message = await ws.recv() + + if isinstance(message, bytes): + current_chunks.append(message) + print(f" Received audio chunk for sentence {current_sentence_index}: {len(message)} bytes") + else: + # JSON frame + msg = json.loads(message) + msg_type = msg.get("type") + + if msg_type == "audio.start": + current_sentence_index = msg["sentence_index"] + current_chunks = [] + print(f" [sentence {msg['sentence_index']}] Generating: {msg['sentence_text']!r}") + elif msg_type == "audio.done": + filename = os.path.join( + output_dir, + f"sentence_{msg['sentence_index']:03d}.{response_format}", + ) + with open(filename, "wb") as f: + f.write(b"".join(current_chunks)) + print( + f" [sentence {msg['sentence_index']}] Done" + f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}" + f" error={msg.get('error', False)}" + f" -> {filename}" + ) + current_chunks = [] + elif msg_type == "session.done": + print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated") + break + elif msg_type == "error": + print(f" ERROR: {msg['message']}") + else: + print(f" Unknown message: {msg}") + finally: + sender_task.cancel() + try: + await sender_task + except asyncio.CancelledError: + pass # Task cancellation is expected during shutdown + + print(f"\nAudio files saved to: {output_dir}/") + + +def main(): + parser = argparse.ArgumentParser(description="Streaming text-input TTS client") + parser.add_argument( + "--url", + default="ws://localhost:8000/v1/audio/speech/stream", + help="WebSocket endpoint URL", + ) + parser.add_argument( + "--text", + required=True, + help="Text to synthesize", + ) + parser.add_argument( + "--output-dir", + default="streaming_tts_output", + help="Directory to save audio files (default: streaming_tts_output)", + ) + + # Session config options + parser.add_argument("--model", default=None, help="Model name") + parser.add_argument("--voice", default="Vivian", help="Speaker voice") + parser.add_argument( + "--task-type", + default="CustomVoice", + choices=["CustomVoice", "VoiceDesign", "Base"], + help="TTS task type", + ) + parser.add_argument("--language", default="Auto", help="Language") + parser.add_argument("--instructions", default=None, help="Voice style instructions") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "pcm", "flac", "mp3", "aac", "opus"], + help="Audio format", + ) + parser.add_argument( + "--stream-audio", + action="store_true", + help="Receive one or more PCM chunks per sentence (requires --response-format pcm)", + ) + parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)") + parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens") + + # Base task options + parser.add_argument("--ref-audio", default=None, help="Reference audio") + parser.add_argument("--ref-text", default=None, help="Reference text") + parser.add_argument( + "--x-vector-only-mode", + action="store_true", + default=False, + help="Speaker embedding only mode", + ) + + # STT simulation + parser.add_argument( + "--simulate-stt", + action="store_true", + help="Simulate STT by sending text word-by-word", + ) + parser.add_argument( + "--stt-delay", + type=float, + default=0.1, + help="Delay between words in STT simulation (seconds)", + ) + + args = parser.parse_args() + + # Build session config (only include non-None values) + config = {} + for key in [ + "model", + "voice", + "task_type", + "language", + "instructions", + "response_format", + "speed", + "max_new_tokens", + "ref_audio", + "ref_text", + ]: + val = getattr(args, key.replace("-", "_"), None) + if val is not None: + config[key] = val + if args.stream_audio: + config["stream_audio"] = True + if args.x_vector_only_mode: + config["x_vector_only_mode"] = True + + asyncio.run( + stream_tts( + url=args.url, + text=args.text, + config=config, + output_dir=args.output_dir, + simulate_stt=args.simulate_stt, + stt_delay=args.stt_delay, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/requirements/common.txt b/requirements/common.txt index 548f8f0bf..fc39c6c14 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -6,7 +6,7 @@ diffusers>=0.36.0 accelerate==1.12.0 gradio==5.50 soundfile>=0.13.1 -cache-dit==1.2.0 +cache-dit==1.3.0 tqdm>=4.66.0 torchsde>=0.2.6 openai-whisper>=20250625 diff --git a/tests/benchmarks/test_serve_cli.py b/tests/benchmarks/test_serve_cli.py index a48731934..4fc0b172c 100644 --- a/tests/benchmarks/test_serve_cli.py +++ b/tests/benchmarks/test_serve_cli.py @@ -3,13 +3,16 @@ from pathlib import Path import pytest +from tests.conftest import OmniServerParams from tests.utils import hardware_test models = ["Qwen/Qwen2.5-Omni-7B"] stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")] # Create parameter combinations for model and stage config -test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +test_params = [ + OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs +] @pytest.mark.core_model diff --git a/tests/comfyui/conftest.py b/tests/comfyui/conftest.py index 7c9770b8e..0b4565e94 100644 --- a/tests/comfyui/conftest.py +++ b/tests/comfyui/conftest.py @@ -63,6 +63,12 @@ def _setup_comfyui_test_environment(): mock_comfy_api_input.AudioInput = AudioInput mock_comfy_api_input.VideoInput = VideoInput mock_comfy_api.input = mock_comfy_api_input + mock_comfy_api_latest = MagicMock() + mock_comfy_api_latest.Types.VideoComponents = MagicMock(side_effect=lambda **kwargs: kwargs) + mock_comfy_api_latest.InputImpl.VideoFromComponents = MagicMock( + side_effect=lambda _: VideoInput(b"mock_video_from_components") + ) + mock_comfy_api.latest = mock_comfy_api_latest def mock_load(_: str | BinaryIO): """Mock nodes_audio.load that returns a waveform tensor (channels, samples) and sample rate.""" @@ -78,5 +84,6 @@ def _setup_comfyui_test_environment(): # Install mock modules BEFORE importing any comfyui_vllm_omni code sys.modules["comfy_api"] = mock_comfy_api sys.modules["comfy_api.input"] = mock_comfy_api_input + sys.modules["comfy_api.latest"] = mock_comfy_api_latest sys.modules["comfy_extras"] = mock_comfy_extras sys.modules["comfy_extras.nodes_audio"] = mock_nodes_audio diff --git a/tests/comfyui/test_comfyui_integration.py b/tests/comfyui/test_comfyui_integration.py index 463596326..d28818ab0 100644 --- a/tests/comfyui/test_comfyui_integration.py +++ b/tests/comfyui/test_comfyui_integration.py @@ -20,11 +20,12 @@ import torch from comfy_api.input import AudioInput, VideoInput from comfyui_vllm_omni.nodes import ( VLLMOmniGenerateImage, + VLLMOmniGenerateVideo, VLLMOmniTTS, VLLMOmniUnderstanding, VLLMOmniVoiceClone, ) -from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams +from comfyui_vllm_omni.utils.types import AutoregressionSamplingParams, DiffusionSamplingParams, WanModelSpecificParams from PIL import Image from vllm import SamplingParams from vllm.outputs import CompletionOutput, RequestOutput @@ -34,6 +35,8 @@ from vllm_omni.entrypoints.cli.serve import OmniServeCommand from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.outputs import OmniRequestOutput +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class ServerCase(NamedTuple): """Parametrizing the model to serve.""" @@ -49,6 +52,7 @@ class SamplingCase(NamedTuple): kind: "SamplingKind" sampling_params: dict | list[dict] | None + lora: dict | None = None class SamplingKind(StrEnum): @@ -58,11 +62,17 @@ class SamplingKind(StrEnum): UNDERSTANDING_AR_LIST = auto() TTS_NONE = auto() TTS_DIFFUSION_SINGLE = auto() + VIDEO_NONE = auto() + VIDEO_DIFFUSION_SINGLE = auto() # Pre-defined arguments to be used in function calls during the tests IMAGE_WIDTH = 64 IMAGE_HEIGHT = 64 +VIDEO_WIDTH = 32 +VIDEO_HEIGHT = 32 +VIDEO_FPS = 8 +VIDEO_NUM_FRAMES = 5 DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams( { "n": 2, @@ -102,6 +112,10 @@ AR_LIST_SAMPLING_PARAMS = [ ), ] +VIDEO_MODEL_PARAMS = WanModelSpecificParams({"guidance_scale_2": 5.0, "boundary_ratio": 0.98, "flow_shift": 12.0}) + +LORA_PARAMS = {"local_path": "test_lora_path", "name": "test_name", "scale": 0.7, "int_id": 10} + def _build_image_output(size: tuple[int, int] = (IMAGE_WIDTH, IMAGE_HEIGHT), color: str = "red") -> Image.Image: return Image.new("RGB", size, color=color) @@ -181,6 +195,16 @@ def _build_diffusion_image_output_for_images_endpoint() -> OmniRequestOutput: ) +def _build_diffusion_video_output() -> OmniRequestOutput: + # Small video: VIDEO_NUM_FRAMES frames of (VIDEO_HEIGHT x VIDEO_WIDTH) RGB, shape (F, H, W, C) + video_frames = torch.zeros((VIDEO_NUM_FRAMES, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32) + return OmniRequestOutput.from_diffusion( + request_id="test_req_video", + images=[video_frames], + final_output_type="video", + ) + + def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput: request_output = MagicMock() request_output.images = [_build_image_output(color="blue")] @@ -193,12 +217,40 @@ def _build_diffusion_image_output_for_chat_endpoint() -> OmniRequestOutput: ) -def _assert_sampling_param_values(received: OmniSamplingParams, expected: dict[str, Any]): +def _assert_sampling_param_values( + received: OmniSamplingParams, expected: dict[str, Any], expected_lora: dict | None = None +): for key, expected_value in expected.items(): - actual_value = getattr(received, key, None) + actual_value = getattr(received, key) assert actual_value == expected_value, ( f"Expected sampling param '{key}'={expected_value}, got {actual_value}. The received sampling params: {received}" ) + if expected_lora: + assert received.lora_request.lora_name == expected_lora["name"], ( + f"Expected lora name={(expected_lora['name'])}, got {received.lora_request.lora_name}. The received sampling params: {received}" + ) + assert received.lora_request.lora_int_id == expected_lora["int_id"], ( + f"Expected lora int_id={expected_lora['int_id']}, got {received.lora_request.lora_int_id}. The received sampling params: {received}" + ) + assert received.lora_request.lora_path == expected_lora["local_path"], ( + f"Expected lora path={expected_lora['local_path']}, got {received.lora_request.lora_path}. The received sampling params: {received}" + ) + assert received.lora_scale == expected_lora["scale"], ( + f"Expected lora scale={expected_lora['scale']}, got {received.lora_scale}. The received sampling params: {received}" + ) + + +def _assert_model_param_values(received: OmniSamplingParams, expected: dict): + for key, expected_value in expected.items(): + try: + actual_value = getattr(received, key) + expected_param_name = key + except AttributeError: + actual_value = received.extra_args.get(key, None) + expected_param_name = f'extra_args["{key}"]' + assert actual_value == expected_value, ( + f"Expected model param '{expected_param_name}'={expected_value}, got {actual_value}. The received sampling params: {received}" + ) def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: SamplingCase, server_case: ServerCase): @@ -235,6 +287,7 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam "height": IMAGE_HEIGHT, **expected, }, + LORA_PARAMS, ) elif sampling_case.kind is SamplingKind.UNDERSTANDING_NONE: assert len(received_sampling_params_list) == 3 @@ -244,6 +297,33 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam _assert_sampling_param_values(received_sampling_params_list[i], expected) elif sampling_case.kind in {SamplingKind.TTS_NONE, SamplingKind.TTS_DIFFUSION_SINGLE}: assert len(received_sampling_params_list) == 1 + elif sampling_case.kind is SamplingKind.VIDEO_NONE: + assert len(received_sampling_params_list) == 1 + _assert_sampling_param_values( + received_sampling_params_list[0], + { + "width": VIDEO_WIDTH, + "height": VIDEO_HEIGHT, + "num_frames": VIDEO_NUM_FRAMES, + "fps": VIDEO_FPS, + }, + ) + elif sampling_case.kind is SamplingKind.VIDEO_DIFFUSION_SINGLE: + assert len(received_sampling_params_list) == 1 + expected = DIFFUSION_SINGLE_SAMPLING_PARAMS.copy() + expected["num_outputs_per_prompt"] = expected.pop("n") # convert from n to num_outputs_per_prompt + _assert_sampling_param_values( + received_sampling_params_list[0], + { + "width": VIDEO_WIDTH, + "height": VIDEO_HEIGHT, + "num_frames": VIDEO_NUM_FRAMES, + "fps": VIDEO_FPS, + **expected, + }, + LORA_PARAMS, + ) + _assert_model_param_values(received_sampling_params_list[0], VIDEO_MODEL_PARAMS) else: raise AssertionError(f"Unknown sampling case: {sampling_case.kind}") @@ -401,9 +481,15 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni @pytest.mark.parametrize( "sampling_case", [ - pytest.param(SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None), id="no-sampling-params"), pytest.param( - SamplingCase(kind=SamplingKind.IMAGE_DIFFUSION_SINGLE, sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS), + SamplingCase(kind=SamplingKind.IMAGE_NONE, sampling_params=None, lora=None), id="no-sampling-params" + ), + pytest.param( + SamplingCase( + kind=SamplingKind.IMAGE_DIFFUSION_SINGLE, + sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS, + lora=LORA_PARAMS, + ), id="single-diffusion-sampling-params", ), ], @@ -423,7 +509,8 @@ async def test_image_generation_node(api_server: str, model: str, image_input: b kwargs["image"] = torch.zeros((1, IMAGE_WIDTH, IMAGE_HEIGHT, 3), dtype=torch.float32) if sampling_case.sampling_params is not None: kwargs["sampling_params"] = sampling_case.sampling_params - print(f"!!!!!! Calling {model} node.generate with kwargs: {sampling_case.sampling_params}") + if sampling_case.lora: + kwargs["lora"] = sampling_case.lora result = await node.generate(**kwargs) @@ -565,3 +652,75 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_ assert result[0]["sample_rate"] == 24000 assert isinstance(result[0]["waveform"], torch.Tensor) assert result[0]["waveform"].shape == (1, 1, 24000) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "server_case,model,image_input", + [ + pytest.param( + ServerCase( + served_model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", + stage_list=["diffusion"], + stage_configs=[{"stage_type": "diffusion"}], + outputs=[_build_diffusion_video_output()], + ), + "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + False, + id="text-to-video", + ), + pytest.param( + ServerCase( + served_model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", + stage_list=["diffusion"], + stage_configs=[{"stage_type": "diffusion"}], + outputs=[_build_diffusion_video_output()], + ), + "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + True, + id="image-to-video", + ), + ], + indirect=["server_case"], +) +@pytest.mark.parametrize( + "sampling_case", + [ + pytest.param(SamplingCase(kind=SamplingKind.VIDEO_NONE, sampling_params=None), id="no-sampling-params"), + pytest.param( + SamplingCase( + kind=SamplingKind.VIDEO_DIFFUSION_SINGLE, + sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS, + lora=LORA_PARAMS, + ), + id="single-diffusion-sampling-params", + ), + ], + indirect=["sampling_case"], +) +async def test_video_generation_node(api_server: str, model: str, image_input: bool, sampling_case: SamplingCase): + node = VLLMOmniGenerateVideo() + + kwargs = { + "url": api_server, + "model": model, + "prompt": "A beautiful sunset timelapse", + "negative_prompt": "", + "width": VIDEO_WIDTH, + "height": VIDEO_HEIGHT, + "fps": VIDEO_FPS, + "num_frames": VIDEO_NUM_FRAMES, + "model_params": VIDEO_MODEL_PARAMS, + } + if image_input: + kwargs["image"] = torch.zeros((1, VIDEO_HEIGHT, VIDEO_WIDTH, 3), dtype=torch.float32) + if sampling_case.sampling_params is not None: + kwargs["sampling_params"] = sampling_case.sampling_params + if sampling_case.lora: + kwargs["lora"] = sampling_case.lora + + result = await node.generate(**kwargs) + + assert isinstance(result, tuple) + assert len(result) == 1 + assert isinstance(result[0], VideoInput) diff --git a/tests/conftest.py b/tests/conftest.py index 619dae3b8..0a3c350d7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -18,16 +18,21 @@ import subprocess import sys import threading import time +from collections.abc import Generator from dataclasses import dataclass +from io import BytesIO from pathlib import Path -from typing import Any +from typing import Any, NamedTuple +import imageio.v3 as iio import numpy as np import psutil import pytest +import soundfile as sf import torch import yaml -from openai import OpenAI +from openai import OpenAI, omit +from PIL import Image from vllm import TextPrompt from vllm.distributed.parallel_state import cleanup_dist_env_and_memory from vllm.logger import init_logger @@ -44,6 +49,65 @@ PromptImageInput = list[Any] | Any | None PromptVideoInput = list[Any] | Any | None +class OmniServerParams(NamedTuple): + model: str + port: int | None = None + stage_config_path: str | None = None + server_args: list[str] | None = None + + +def assert_image_valid(image: Path | Image.Image, *, width: int | None = None, height: int | None = None): + """Assert the file is a loadable image with optional exact dimensions.""" + if isinstance(image, Path): + assert image.exists(), f"Image not found: {image}" + image = Image.open(image) + image.load() + assert image.width > 0 and image.height > 0 + if width is not None: + assert image.width == width, f"Expected width={width}, got {image.width} in {image.name}" + if height is not None: + assert image.height == height, f"Expected height={height}, got {image.height} in {image.name}" + return image + + +def assert_video_valid(frames: Path | np.ndarray, *, width: int, height: int, num_frames: int) -> None: + """Assert the MP4 has the expected resolution and exact frame count.""" + if isinstance(frames, Path): + assert frames.exists(), f"Video not found: {frames}" + frames = iio.imread(str(frames), plugin="pyav", index=None) + assert frames.shape[0] == num_frames, f"Expected {num_frames} frames, got {frames.shape[0]}" + assert frames.shape[1] == height, f"Expected height={height}, got {frames.shape[1]}" + assert frames.shape[2] == width, f"Expected width={width}, got {frames.shape[2]}" + + +def assert_audio_valid(path: Path, *, sample_rate: int, channels: int, duration_s: float) -> None: + """Assert the WAV has the expected sample rate, channel count, and duration.""" + + assert path.exists(), f"Audio not found: {path}" + info = sf.info(str(path)) + assert info.samplerate == sample_rate, f"Expected sample_rate={sample_rate}, got {info.samplerate}" + assert info.channels == channels, f"Expected {channels} channel(s), got {info.channels}" + expected_frames = int(duration_s * sample_rate) + assert info.frames == expected_frames, ( + f"Expected {expected_frames} frames ({duration_s}s @ {sample_rate} Hz), got {info.frames}" + ) + + +def decode_b64_image(b64: str): + img = Image.open(BytesIO(base64.b64decode(b64))) + img.load() + return img + + +@pytest.fixture(scope="session") +def model_prefix() -> str: + """Optional model-path prefix from MODEL_PREFIX env var. + Useful if models are downloaded to non-default local directories. + """ + prefix = os.environ.get("MODEL_PREFIX", "") + return f"{prefix.rstrip('/')}/" if prefix else "" + + @pytest.fixture(autouse=True) def default_vllm_config(): """Set a default VllmConfig for all tests. @@ -1151,7 +1215,9 @@ def pytest_addoption(parser): @pytest.fixture(scope="session") -def run_level(request): +def run_level(request) -> str: + """A command-line argument that specifies the level of tests to run in this session. + See https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/CI_5levels/""" return request.config.getoption("--run-level") @@ -1159,15 +1225,17 @@ _omni_server_lock = threading.Lock() @pytest.fixture(scope="module") -def omni_server(request, run_level): +def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]: """Start vLLM-Omni server as a subprocess with actual model weights. Uses session scope so the server starts only once for the entire test session. Multi-stage initialization can take 10-20+ minutes. """ with _omni_server_lock: - *port, model, stage_config_path = request.param - port = port[0] if port else None - if run_level == "advanced_model": + params: OmniServerParams = request.param + model = model_prefix + params.model + port = params.port + stage_config_path = params.stage_config_path + if run_level == "advanced_model" and stage_config_path is not None: stage_config_path = modify_stage_config( stage_config_path, deletes={ @@ -1179,7 +1247,11 @@ def omni_server(request, run_level): }, ) - server_args = ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"] + server_args = params.server_args or [] + server_args = ["--stage-init-timeout", "120", *server_args] + if stage_config_path is not None: + server_args += ["--stage-configs-path", stage_config_path] + with OmniServer(model, server_args, port=port) if port else OmniServer(model, server_args) as server: print("OmniServer started successfully") yield server @@ -1199,6 +1271,17 @@ class OmniResponse: error_message: str | None = None +@dataclass +class DiffusionResponse: + text_content: str | None = None + images: list[Image.Image] | None = None + audios: list[Any] | None = None + videos: list[Any] | None = None + e2e_latency: float | None = None + success: bool = False + error_message: str | None = None + + def assert_omni_response(response: OmniResponse, request_config: dict[str, Any], run_level): """ Validate response results. @@ -1247,6 +1330,51 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any], print(f"similarity is: {response.similarity}") +def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None): + """ + Validate diffusion response results. + + Args: + response: DiffusionResponse object. Any not-None content will be validated based on the request_config. + request_config: Request configuration dictionary containing parameters like model, messages, extra_body. + When validating a certain modality, the corresponding params in request_config['extra_body'] must present. + It will be used to check against the multimedia file in the response. + run_level: Test run level (e.g., "core_model", "advanced_model") + + Raises: + AssertionError: When the response does not meet validation criteria + KeyError: When the request_config does not contain necessary parameters for validation + """ + assert response.success, "The request failed." + + e2e_latency = response.e2e_latency + if e2e_latency is not None: + print(f"the avg e2e is: {e2e_latency}") + + extra_body = request_config.get("extra_body", {}) + + num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1) + + if response.images is not None: + assert len(response.images) > 0, "No images in response" + assert len(response.images) == num_outputs_per_prompt, ( + f"Expected {num_outputs_per_prompt} images, got {len(response.images)}" + ) + if run_level == "advanced_model": + expected_width = extra_body["width"] # intend to raise KeyError + expected_height = extra_body["height"] # intend to raise KeyError + for img in response.images: + assert_image_valid(img, width=expected_width, height=expected_height) + if response.videos is not None: + raise NotImplementedError( + "Video validation is not implemented yet" + ) # consider using assert_video_valid defined above + if response.audios is not None: + raise NotImplementedError( + "Audio validation is not implemented yet" + ) # consider using assert_audio_valid defined above + + class OpenAIClientHandler: """ OpenAI client handler class, encapsulating both streaming and non-streaming response processing logic. @@ -1269,7 +1397,7 @@ class OpenAIClientHandler: self.client = OpenAI(base_url=f"http://{host}:{port}/v1", api_key=api_key) self.run_level = run_level - def _process_stream_response(self, chat_completion) -> OmniResponse: + def _process_stream_omni_response(self, chat_completion) -> OmniResponse: """ Process streaming responses. @@ -1330,7 +1458,7 @@ class OpenAIClientHandler: return result - def _process_non_stream_response(self, chat_completion) -> OmniResponse: + def _process_non_stream_omni_response(self, chat_completion) -> OmniResponse: """ Process non-streaming responses. @@ -1384,7 +1512,46 @@ class OpenAIClientHandler: return result - def send_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: + def _process_diffusion_response(self, chat_completion) -> DiffusionResponse: + """ + Process diffusion responses (image generation/editing). + + Args: + chat_completion: OpenAI response object + + Returns: + DiffusionResponse: Processed response object + """ + result = DiffusionResponse() + start_time = time.perf_counter() + + try: + images = [] + # [TODO] reading video and audio output from API response for later validation + + for choice in chat_completion.choices: + if hasattr(choice.message, "content") and choice.message.content is not None: + content = choice.message.content + if isinstance(content, list): + for item in content: + if hasattr(item, "image_url") and item.image_url is not None: + image_url = item.image_url.url + if image_url.startswith("data:image"): + b64_data = image_url.split(",", 1)[1] + img = decode_b64_image(b64_data) + images.append(img) + + result.e2e_latency = time.perf_counter() - start_time + result.images = images if images else None + result.success = True + + except Exception as e: + result.error_message = f"Diffusion response processing error: {str(e)}" + print(f"Error: {result.error_message}") + + return result + + def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: """ Send OpenAI requests. @@ -1410,9 +1577,9 @@ class OpenAIClientHandler: ) if stream: - response = self._process_stream_response(chat_completion) + response = self._process_stream_omni_response(chat_completion) else: - response = self._process_non_stream_response(chat_completion) + response = self._process_non_stream_omni_response(chat_completion) assert_omni_response(response, request_config, run_level=self.run_level) responses.append(response) @@ -1438,25 +1605,82 @@ class OpenAIClientHandler: chat_completion = future.result() if stream: - response = self._process_stream_response(chat_completion) + response = self._process_stream_omni_response(chat_completion) else: - response = self._process_non_stream_response(chat_completion) + response = self._process_non_stream_omni_response(chat_completion) assert_omni_response(response, request_config, run_level=self.run_level) responses.append(response) return responses + def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: + """ + Send OpenAI requests for diffusion models. + + Args: + request_config: Request configuration dictionary containing parameters like model, messages + request_num: Number of requests to send concurrently, defaults to 1 (single request) + Returns: + List[OmniResponse]: List of response objects + """ + responses = [] + stream = request_config.get("stream", False) + modalities = request_config.get("modalities", omit) # Most diffusion models don't require modalities param + extra_body = request_config.get("extra_body", None) + + if stream: + raise NotImplementedError("Streaming is not currently implemented for diffusion model e2e test") + + if request_num == 1: + # Send single request + chat_completion = self.client.chat.completions.create( + model=request_config.get("model"), + messages=request_config.get("messages"), + extra_body=extra_body, + modalities=modalities, + ) + + response = self._process_diffusion_response(chat_completion) + assert_diffusion_response(response, request_config, run_level=self.run_level) + responses.append(response) + + else: + # Send concurrent requests + with concurrent.futures.ThreadPoolExecutor(max_workers=request_num) as executor: + futures = [] + + # Submit all request tasks + for _ in range(request_num): + future = executor.submit( + self.client.chat.completions.create, + model=request_config.get("model"), + messages=request_config.get("messages"), + modalities=modalities, + extra_body=extra_body, + ) + futures.append(future) + + # Process completed tasks + for future in concurrent.futures.as_completed(futures): + chat_completion = future.result() + response = self._process_diffusion_response(chat_completion) + assert_diffusion_response(response, request_config, run_level=self.run_level) + responses.append(response) + + return responses + @pytest.fixture -def openai_client(omni_server, run_level): - """Create OpenAIClientHandler fixture""" +def openai_client(omni_server: OmniServer, run_level: str): + """Create OpenAIClientHandler fixture to facilitate communication with OmniServer + with encapsulated request sending, concurrent requests, response handling, and validation.""" return OpenAIClientHandler(host=omni_server.host, port=omni_server.port, api_key="EMPTY", run_level=run_level) class OmniRunner: """ - Test runner for Omni models. + Offline test runner for Omni models. """ def __init__( @@ -1727,9 +1951,10 @@ class OmniRunner: @pytest.fixture(scope="module") -def omni_runner(request): +def omni_runner(request, model_prefix): with _omni_server_lock: model, stage_config_path = request.param + model = model_prefix + model with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, stage_init_timeout=300) as runner: print("OmniRunner started successfully") yield runner diff --git a/tests/diffusion/models/flux2/test_flux2_transformer_tp.py b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py new file mode 100644 index 000000000..a2d1fe6ab --- /dev/null +++ b/tests/diffusion/models/flux2/test_flux2_transformer_tp.py @@ -0,0 +1,244 @@ +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.diffusion.models.flux2.flux2_transformer import ( + Flux2PosEmbed, + Flux2Transformer2DModel, +) + + +# Initialize TP group before tests +@pytest.fixture(scope="function", autouse=True) +def setup_tp_group(): + """Set up TP group for each test function""" + with patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=2): + with patch("vllm.distributed.parallel_state.get_tp_group") as mock_get_tp_group: + mock_tp_group = MagicMock() + mock_tp_group.world_size = 2 + mock_get_tp_group.return_value = mock_tp_group + yield + + +class TestFlux2TransformerWeightLoading: + """Test Flux2Transformer weight loading functionality""" + + def test_weight_loading_tp2(self, setup_tp_group): + """Verify weights load correctly with TP=2""" + # Prepare test data + model = Flux2Transformer2DModel( + num_layers=2, + num_single_layers=1, + num_attention_heads=48, + attention_head_dim=128, + joint_attention_dim=15360, + ) + + # Mock TP=2 weight loading + mock_weights = [] + + # 1. Test regular weight loading + mock_weights.append(("x_embedder.weight", torch.randn(6144, 128))) + mock_weights.append(("context_embedder.weight", torch.randn(6144, 15360))) + mock_weights.append(("proj_out.weight", torch.randn(128, 6144))) + + # 2. Test stacked_params_mapping weight loading + # Full weights - load_weights handles sharding internally + to_qkv_weight = torch.randn(18432, 6144) # Full weights + mock_weights.append(("transformer_blocks.0.attn.to_qkv.weight", to_qkv_weight)) + + # Add_kv_proj weights + add_kv_proj_weight = torch.randn(18432, 6144) # Full weights + mock_weights.append(("transformer_blocks.0.attn.add_kv_proj.weight", add_kv_proj_weight)) + + # 3. Test single block weight loading + single_block_qkv_mlp_weight = torch.randn(18432 + 18432 * 3, 6144) # Full weights + mock_weights.append(("single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight", single_block_qkv_mlp_weight)) + + # Execute weight loading + loaded_params = model.load_weights(mock_weights) + + # Verify + assert len(loaded_params) > 0, "Parameters should be loaded" + + # Verify stacked_params_mapping is correctly set + assert model.stacked_params_mapping is not None + # Should have 6 mappings: 3 for to_qkv + 3 for add_kv_proj + assert len(model.stacked_params_mapping) == 6, "Should have 6 mappings" + + # Verify weights are correctly loaded to corresponding layers + assert hasattr(model.transformer_blocks[0].attn.to_qkv, "weight") + # With TP=2, weight dimension on each GPU should be 18432/2 = 9216 + assert model.transformer_blocks[0].attn.to_qkv.weight.shape == (9216, 6144), ( + f"With TP=2, to_qkv weight dimension should be (9216, 6144), got {model.transformer_blocks[0].attn.to_qkv.weight.shape}" + ) + + +class TestFlux2RopePositionEmbedding: + """Test Flux2 RoPE position embedding functionality""" + + def test_rope_position_embedding(self): + """Verify RoPE produces correct embeddings for 4D coordinates""" + # Prepare test data - use model default configuration + # axes_dims_rope default is (32, 32, 32, 32) + # get_1d_rotary_pos_embed outputs half the dimension for real/imag parts + # So actual output dimension should be (16+16+16+16) = 64 + axes_dims = (32, 32, 32, 32) # Use model default + rope_theta = 2000 # Model default is 2000 + pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims) + + # Create test IDs + seq_len = 10 + ids = torch.randint(0, 100, (seq_len, 4)) # [S, 4] + + # Forward pass + freqs_cos, freqs_sin = pos_embed(ids) + + # Verify output shape - based on model config, expected dimension is 64 + # Each axes_dim=32 outputs 16-dim real part, sum of 4 dimensions = 64 + expected_dim = sum(axes_dims) // 2 # 128/2 = 64 + assert freqs_cos.shape == (seq_len, expected_dim), ( + f"Expected shape {(seq_len, expected_dim)}, got {freqs_cos.shape}" + ) + assert freqs_sin.shape == (seq_len, expected_dim), ( + f"Expected shape {(seq_len, expected_dim)}, got {freqs_sin.shape}" + ) + + # Verify output type - NPU may use float32, other devices use float64 + assert freqs_cos.dtype in [torch.float32, torch.float64], "Should be float type" + assert freqs_sin.dtype in [torch.float32, torch.float64], "Should be float type" + + # Verify value range + assert torch.all(freqs_cos >= -1) and torch.all(freqs_cos <= 1), "cos values should be in [-1, 1]" + assert torch.all(freqs_sin >= -1) and torch.all(freqs_sin <= 1), "sin values should be in [-1, 1]" + + # Verify trigonometric relationship: cos^2 + sin^2 β‰ˆ 1 + cos_sq_sin_sq = freqs_cos**2 + freqs_sin**2 + assert torch.allclose(cos_sq_sin_sq, torch.ones_like(cos_sq_sin_sq), atol=1e-6), "cos^2 + sin^2 should β‰ˆ 1" + + # Verify different positions produce different embeddings + ids_diff = torch.randint(0, 100, (seq_len, 4)) + freqs_cos_diff, freqs_sin_diff = pos_embed(ids_diff) + assert not torch.allclose(freqs_cos, freqs_cos_diff), "Different positions should produce different embeddings" + + # Verify same positions produce same embeddings + ids_same = ids.clone() + freqs_cos_same, freqs_sin_same = pos_embed(ids_same) + assert torch.allclose(freqs_cos, freqs_cos_same), "Same positions should produce same embeddings" + assert torch.allclose(freqs_sin, freqs_sin_same), "Same positions should produce same embeddings" + + +class TestFlux2PackedModuleMapping: + """Test Flux2 packed module mapping functionality""" + + def test_packed_module_mapping(self, setup_tp_group): + """Verify to_qkv packing matches HF checkpoint""" + model = Flux2Transformer2DModel( + num_layers=1, + num_single_layers=0, + num_attention_heads=48, + attention_head_dim=128, + joint_attention_dim=15360, + ) + + # Verify stacked_params_mapping is correctly initialized + model.load_weights([]) # Trigger stacked_params_mapping initialization + assert model.stacked_params_mapping is not None + + # Verify mapping configuration + expected_mappings = [ + (".to_qkv.", ".to_q.", "q"), + (".to_qkv.", ".to_k.", "k"), + (".to_qkv.", ".to_v.", "v"), + (".add_kv_proj", ".add_q_proj", "q"), + (".add_kv_proj", ".add_k_proj", "k"), + (".add_kv_proj", ".add_v_proj", "v"), + ] + assert model.stacked_params_mapping == expected_mappings + + # Create mock HF checkpoint weights + hf_weights = [] + + # Mock HF separated Q/K/V weights + attn_block = model.transformer_blocks[0].attn + head_dim = 128 + num_heads = 48 + hidden_size = 6144 + # Full weight dimension + full_shard_size = num_heads * head_dim # 6144 + + # Q projection weights (full weights) + q_weight = torch.randn(full_shard_size, hidden_size) + hf_weights.append(("transformer_blocks.0.attn.to_q.weight", q_weight)) + + # K projection weights (full weights) + k_weight = torch.randn(full_shard_size, hidden_size) + hf_weights.append(("transformer_blocks.0.attn.to_k.weight", k_weight)) + + # V projection weights (full weights) + v_weight = torch.randn(full_shard_size, hidden_size) + hf_weights.append(("transformer_blocks.0.attn.to_v.weight", v_weight)) + + # Mock HF separated add_q/k/v projection weights + add_q_weight = torch.randn(full_shard_size, hidden_size) + hf_weights.append(("transformer_blocks.0.attn.add_q_proj.weight", add_q_weight)) + + add_k_weight = torch.randn(full_shard_size, hidden_size) + hf_weights.append(("transformer_blocks.0.attn.add_k_proj.weight", add_k_weight)) + + add_v_weight = torch.randn(full_shard_size, hidden_size) + hf_weights.append(("transformer_blocks.0.attn.add_v_proj.weight", add_v_weight)) + + # Execute weight loading + loaded_params = model.load_weights(hf_weights) + + # Verify weights are loaded + assert len(loaded_params) > 0 + + # Verify final QKV weights are correctly combined (considering TP sharding) + # With TP=2, dimension on each GPU should be half of full dimension + expected_qkv_shape = (full_shard_size * 3 // 2, hidden_size) # 18432/2 = 9216 + assert attn_block.to_qkv.weight.shape == expected_qkv_shape, ( + f"to_qkv weight dimension should be {expected_qkv_shape}, got {attn_block.to_qkv.weight.shape}" + ) + + expected_add_kv_shape = (full_shard_size * 3 // 2, hidden_size) + assert attn_block.add_kv_proj.weight.shape == expected_add_kv_shape, ( + f"add_kv_proj weight dimension should be {expected_add_kv_shape}, got {attn_block.add_kv_proj.weight.shape}" + ) + + def test_packed_mapping_edge_cases(self, setup_tp_group): + """Test edge cases for packed mapping""" + model = Flux2Transformer2DModel( + num_layers=1, + num_single_layers=1, + num_attention_heads=48, + attention_head_dim=128, + joint_attention_dim=15360, + ) + model.load_weights([]) + + # Test invalid weight names + invalid_weights = [("invalid.weight", torch.randn(10, 10))] + loaded_params = model.load_weights(invalid_weights) + assert len(loaded_params) == 0, "Should not load invalid weights" + + # Test to_out weight renaming + to_out_weight = torch.randn(6144, 6144) + weights = [("transformer_blocks.0.attn.to_out.0.weight", to_out_weight)] + loaded_params = model.load_weights(weights) + + # Check if to_out related weights are loaded + to_out_loaded = any("to_out" in p for p in loaded_params) + assert to_out_loaded, "to_out weights should be correctly renamed and loaded" + + # Test partial weight loading + partial_weights = [ + ("x_embedder.weight", torch.randn(6144, 128)), + ("transformer_blocks.0.attn.to_q.weight", torch.randn(6144, 6144)), # Full weights + ] + loaded_params = model.load_weights(partial_weights) + assert len(loaded_params) == 2, "Should load two weights" + assert "x_embedder.weight" in loaded_params + assert any("to_qkv" in p for p in loaded_params), "to_q should be mapped to to_qkv" diff --git a/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py b/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py new file mode 100644 index 000000000..2aa1adf44 --- /dev/null +++ b/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for HunyuanFusedMoE (Support HunyuanImage3 Diffusion Model, 5a779b4).""" + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestHunyuanFusedMoEPlatformDispatch: + """Test platform dispatch via platform qualname hooks.""" + + def test_default_platform_uses_default_impl_qualname(self, mocker): + """HunyuanFusedMoE should resolve the impl class from the platform hook.""" + import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + + mock_platform = mocker.MagicMock() + mock_platform.get_diffusion_model_impl_qualname.return_value = ( + "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault" + ) + + mocker.patch.object( + hunyuan_moe, + "current_omni_platform", + mock_platform, + ) + mock_resolve = mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname") + mock_impl = mocker.MagicMock() + mock_resolve.return_value = mock_impl + + from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import ( + HunyuanFusedMoE, + ) + + HunyuanFusedMoE(prefix="") + + mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("hunyuan_fused_moe") + mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe") + mock_resolve.assert_called_once_with( + "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault" + ) + mock_impl.assert_called_once_with(prefix="") + + +class TestHunyuanFusedMoEFactory: + """Test HunyuanFusedMoE factory __new__ and make_expert_params_mapping delegation.""" + + def test_new_delegates_to_impl_class(self, mocker): + """HunyuanFusedMoE(prefix=..., **kwargs) should instantiate and return impl instance.""" + import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + + class MockImpl: + def __init__(self, *, prefix: str = "", **kwargs): + self.prefix = prefix + self.kwargs = kwargs + + mock_platform = mocker.MagicMock() + mock_platform.get_diffusion_model_impl_qualname.return_value = "mock.impl.Qualname" + mocker.patch.object(hunyuan_moe, "current_omni_platform", mock_platform) + + mock_impl_class = mocker.MagicMock(return_value=MockImpl(prefix="test", a=1)) + mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class) + + from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import ( + HunyuanFusedMoE, + ) + + result = HunyuanFusedMoE(prefix="test", a=1) + + assert isinstance(result, MockImpl) + assert result.prefix == "test" + assert result.kwargs == {"a": 1} + mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("hunyuan_fused_moe") + mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe") + mock_impl_class.assert_called_once_with(prefix="test", a=1) + + def test_make_expert_params_mapping_delegates_to_impl(self, mocker): + """make_expert_params_mapping should delegate to impl class method.""" + import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + + expected_mapping = [("a", "b", 0, "c")] + mock_platform = mocker.MagicMock() + mock_platform.get_diffusion_model_impl_qualname.return_value = "mock.impl.Qualname" + mocker.patch.object(hunyuan_moe, "current_omni_platform", mock_platform) + + mock_impl_class = mocker.MagicMock() + mock_impl_class.make_expert_params_mapping = mocker.MagicMock(return_value=expected_mapping) + mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class) + + from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import ( + HunyuanFusedMoE, + ) + + result = HunyuanFusedMoE.make_expert_params_mapping( + model=None, + ckpt_gate_proj_name="gate", + ckpt_down_proj_name="down", + ckpt_up_proj_name="up", + num_experts=4, + num_redundant_experts=0, + ) + + assert result == expected_mapping + mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe") + mock_impl_class.make_expert_params_mapping.assert_called_once_with( + None, + ckpt_gate_proj_name="gate", + ckpt_down_proj_name="down", + ckpt_up_proj_name="up", + num_experts=4, + num_redundant_experts=0, + ) diff --git a/tests/diffusion/test_data_is_moe.py b/tests/diffusion/test_data_is_moe.py new file mode 100644 index 000000000..25fa59ef1 --- /dev/null +++ b/tests/diffusion/test_data_is_moe.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for OmniDiffusionConfig.is_moe (fix is_moe type and threshold, 6663c0b).""" + +import pytest + +from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestOmniDiffusionConfigIsMoE: + """Tests for OmniDiffusionConfig.is_moe property. + + Covers commit 6663c0b: fix is_moe type and threshold + - num_experts must be (list, tuple, int); otherwise return False. + - Threshold: is_moe is True when num_experts > 0 (not > 1). + """ + + def test_is_moe_missing_num_experts_returns_false(self): + """When num_experts is absent, is_moe should be False.""" + tf_config = TransformerConfig.from_dict({}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is False + + def test_is_moe_none_num_experts_returns_false(self): + """When num_experts is explicitly None (e.g. in params), is_moe should be False.""" + tf_config = TransformerConfig.from_dict({"num_experts": None}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is False + + def test_is_moe_non_allowed_type_returns_false(self): + """When num_experts is not int/list/tuple (e.g. str), is_moe should be False.""" + tf_config = TransformerConfig.from_dict({"num_experts": "2"}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is False + + def test_is_moe_int_zero_returns_false(self): + """num_experts int 0 should yield is_moe False (threshold > 0).""" + tf_config = TransformerConfig.from_dict({"num_experts": 0}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is False + + def test_is_moe_int_one_returns_true(self): + """num_experts int 1 should yield is_moe True (threshold > 0, not > 1).""" + tf_config = TransformerConfig.from_dict({"num_experts": 1}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is True + + def test_is_moe_int_gt_one_returns_true(self): + """num_experts int > 1 should yield is_moe True.""" + tf_config = TransformerConfig.from_dict({"num_experts": 2}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is True + + def test_is_moe_list_all_zero_returns_false(self): + """num_experts list with all <= 0 should yield is_moe False.""" + tf_config = TransformerConfig.from_dict({"num_experts": [0]}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is False + + def test_is_moe_list_has_positive_returns_true(self): + """num_experts list with any int > 0 should yield is_moe True.""" + tf_config = TransformerConfig.from_dict({"num_experts": [0, 1]}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is True + + def test_is_moe_tuple_has_positive_returns_true(self): + """num_experts tuple with any int > 0 should yield is_moe True.""" + tf_config = TransformerConfig.from_dict({"num_experts": (0, 2)}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is True + + def test_is_moe_list_non_int_ignored(self): + """num_experts list with only non-int entries should yield is_moe False.""" + tf_config = TransformerConfig.from_dict({"num_experts": ["a", 0.0]}) + config = OmniDiffusionConfig(model="test", tf_model_config=tf_config) + assert config.is_moe is False diff --git a/tests/distributed/omni_coordinator/test_load_balancer.py b/tests/distributed/omni_coordinator/test_load_balancer.py new file mode 100644 index 000000000..c54d24894 --- /dev/null +++ b/tests/distributed/omni_coordinator/test_load_balancer.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from time import time + +from vllm_omni.distributed.omni_coordinator import ( + InstanceInfo, + RandomBalancer, + StageStatus, +) + + +def test_load_balancer_select_returns_valid_index(): + """Verify RandomBalancer.select() returns a valid index for instances.""" + # Task structure mirrors async_omni; RandomBalancer ignores task contents. + task: dict = { + "request_id": "test", + "engine_inputs": None, + "sampling_params": None, + } + + now = time() + instances = [ + InstanceInfo( + input_addr="tcp://host:10001", + output_addr="tcp://host:10001-out", + stage_id=0, + status=StageStatus.UP, + queue_length=0, + last_heartbeat=now, + registered_at=now, + ), + InstanceInfo( + input_addr="tcp://host:10002", + output_addr="tcp://host:10002-out", + stage_id=0, + status=StageStatus.UP, + queue_length=1, + last_heartbeat=now, + registered_at=now, + ), + InstanceInfo( + input_addr="tcp://host:10003", + output_addr="tcp://host:10003-out", + stage_id=1, + status=StageStatus.UP, + queue_length=2, + last_heartbeat=now, + registered_at=now, + ), + ] + + balancer = RandomBalancer() + + index = balancer.select(task, instances) + + assert isinstance(index, int) + assert 0 <= index < len(instances) diff --git a/tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py b/tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py new file mode 100644 index 000000000..24b331923 --- /dev/null +++ b/tests/distributed/omni_coordinator/test_omni_coord_client_for_hub.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import time + +import pytest +import zmq + +from vllm_omni.distributed.omni_coordinator import ( + InstanceList, + OmniCoordClientForHub, +) + + +def _bind_pub() -> tuple[zmq.Context, zmq.Socket, str]: + ctx = zmq.Context.instance() + pub = ctx.socket(zmq.PUB) + pub.bind("tcp://127.0.0.1:*") + endpoint = pub.getsockopt(zmq.LAST_ENDPOINT).decode("ascii") + return ctx, pub, endpoint + + +def _wait_for_condition(cond, timeout: float = 2.0, interval: float = 0.01) -> bool: + start = time.time() + while time.time() - start < timeout: + if cond(): + return True + time.sleep(interval) + return False + + +def test_hub_client_caches_instance_list_from_pub(): + """Verify OmniCoordClientForHub receives instance list updates from OmniCoordinator and caches for get_instance_list().""" + ctx, pub, endpoint = _bind_pub() + + client = OmniCoordClientForHub(endpoint) + # ZMQ PUB/SUB slow-joiner: allow SUB to finish connecting before first send + time.sleep(0.2) + + now = time.time() + instances_payload = [ + { + "input_addr": "tcp://stage:10001", + "output_addr": "tcp://stage:10001-out", + "stage_id": 0, + "status": "up", + "queue_length": 0, + "last_heartbeat": now, + "registered_at": now, + }, + { + "input_addr": "tcp://stage:10002", + "output_addr": "tcp://stage:10002-out", + "stage_id": 0, + "status": "up", + "queue_length": 1, + "last_heartbeat": now, + "registered_at": now, + }, + { + "input_addr": "tcp://stage:10003", + "output_addr": "tcp://stage:10003-out", + "stage_id": 1, + "status": "error", + "queue_length": 5, + "last_heartbeat": now, + "registered_at": now, + }, + ] + + payload = {"instances": instances_payload, "timestamp": now} + pub.send(json.dumps(payload).encode("utf-8")) + + assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 3) + + inst_list = client.get_instance_list() + assert isinstance(inst_list, InstanceList) + assert len(inst_list.instances) == 3 + + for src, inst in zip(instances_payload, inst_list.instances, strict=True): + assert inst.input_addr == src["input_addr"] + assert inst.output_addr == src["output_addr"] + assert inst.stage_id == src["stage_id"] + assert inst.status.value == src["status"] + + stage0 = client.get_instances_for_stage(0) + stage1 = client.get_instances_for_stage(1) + + assert all(inst.stage_id == 0 for inst in stage0.instances) + assert all(inst.stage_id == 1 for inst in stage1.instances) + + # Send an updated list with fewer instances and verify cache refresh. + updated_payload = { + "instances": instances_payload[:2], + "timestamp": now + 1.0, + } + pub.send(json.dumps(updated_payload).encode("utf-8")) + + assert _wait_for_condition(lambda: len(client.get_instance_list().instances) == 2) + updated_list = client.get_instance_list() + assert len(updated_list.instances) == 2 + + client.close() + pub.close(0) + ctx.term() + + +def test_hub_client_close_closes_sub_socket(): + """Verify OmniCoordClientForHub.close() marks client as closed; second close raises.""" + ctx, pub, endpoint = _bind_pub() + client = OmniCoordClientForHub(endpoint) + client.close() + + with pytest.raises(RuntimeError, match="already closed"): + client.close() + + pub.close(0) + ctx.term() diff --git a/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py b/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py new file mode 100644 index 000000000..b74a48f49 --- /dev/null +++ b/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import zmq + +from vllm_omni.distributed.omni_coordinator import ( + OmniCoordClientForStage, + StageStatus, +) + + +def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]: + ctx = zmq.Context.instance() + router = ctx.socket(zmq.ROUTER) + router.bind("tcp://127.0.0.1:*") + endpoint = router.getsockopt(zmq.LAST_ENDPOINT).decode("ascii") + return ctx, router, endpoint + + +def _recv_event(router: zmq.Socket) -> dict: + frames = router.recv_multipart() + # ROUTER adds identity frame; the last frame is the payload. + payload = frames[-1] + return json.loads(payload.decode("utf-8")) + + +def test_stage_client_auto_register_on_init(): + """Verify OmniCoordClientForStage automatically sends initial registration/status-up event when created.""" + ctx, router, endpoint = _bind_router() + + input_addr = "tcp://stage:10001" + output_addr = "tcp://stage:10001-out" + stage_id = 0 + + client = OmniCoordClientForStage(endpoint, input_addr, output_addr, stage_id) + + event = _recv_event(router) + + assert event["event_type"] == "update" + assert event["status"] == StageStatus.UP.value + assert event["stage_id"] == stage_id + assert event["input_addr"] == input_addr + assert event["output_addr"] == output_addr + + client.close() + router.close(0) + ctx.term() + + +def test_stage_client_update_info_sends_correct_event(): + """Verify OmniCoordClientForStage.update_info() sends status/load update events with expected fields.""" + ctx, router, endpoint = _bind_router() + + input_addr = "tcp://stage:10002" + output_addr = "tcp://stage:10002-out" + stage_id = 1 + + client = OmniCoordClientForStage(endpoint, input_addr, output_addr, stage_id) + + # Discard initial registration event. + _recv_event(router) + + client.update_info(status=StageStatus.ERROR) + client.update_info(queue_length=10) + + first = _recv_event(router) + second = _recv_event(router) + + assert first["status"] == StageStatus.ERROR.value + assert first["stage_id"] == stage_id + assert first["input_addr"] == input_addr + assert first["output_addr"] == output_addr + + assert second["queue_length"] == 10 + assert second["stage_id"] == stage_id + assert second["input_addr"] == input_addr + assert second["output_addr"] == output_addr + + client.close() + router.close(0) + ctx.term() + + +def test_stage_client_close_sends_down_status(): + """Verify close() sends final status-down event before closing underlying socket.""" + ctx, router, endpoint = _bind_router() + + input_addr = "tcp://stage:10003" + output_addr = "tcp://stage:10003-out" + stage_id = 2 + + client = OmniCoordClientForStage(endpoint, input_addr, output_addr, stage_id) + + # Discard initial registration event. + _recv_event(router) + + client.close() + + event = _recv_event(router) + assert event["status"] == StageStatus.DOWN.value + assert event["stage_id"] == stage_id + assert event["input_addr"] == input_addr + assert event["output_addr"] == output_addr + + assert client._socket.closed # DEALER socket no longer usable after close + + router.close(0) + ctx.term() diff --git a/tests/distributed/omni_coordinator/test_omni_coordinator.py b/tests/distributed/omni_coordinator/test_omni_coordinator.py new file mode 100644 index 000000000..0c68e61bb --- /dev/null +++ b/tests/distributed/omni_coordinator/test_omni_coordinator.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import time + +import zmq +from vllm.v1.utils import get_engine_client_zmq_addr + +from vllm_omni.distributed.omni_coordinator import ( + OmniCoordClientForStage, + OmniCoordinator, + StageStatus, +) + + +def _recv_instance_list(sub: zmq.Socket, timeout_ms: int = 2000) -> dict | None: + """Receive InstanceList JSON from SUB socket. Returns None on timeout.""" + sub.setsockopt(zmq.RCVTIMEO, timeout_ms) + try: + data = sub.recv() + return json.loads(data.decode("utf-8")) + except zmq.Again: + return None + + +def _wait_for_instance_list( + sub: zmq.Socket, + expected_count: int, + timeout: float = 3.0, +) -> dict | None: + """Wait until received InstanceList with expected_count active instances.""" + start = time.time() + while time.time() - start < timeout: + msg = _recv_instance_list(sub, timeout_ms=500) + if msg is not None and len(msg.get("instances", [])) == expected_count: + return msg + return None + + +def test_omni_coordinator_registration_broadcast(): + """Verify that after multiple OmniCoordClientForStage instances register, + OmniCoordinator publishes an InstanceList containing all registered instances. + """ + router_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + pub_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + coordinator = OmniCoordinator( + router_zmq_addr=router_addr, + pub_zmq_addr=pub_addr, + heartbeat_timeout=1000.0, + ) + + sub_ctx = zmq.Context.instance() + sub = sub_ctx.socket(zmq.SUB) + sub.connect(pub_addr) + sub.setsockopt(zmq.SUBSCRIBE, b"") + + # ZMQ PUB/SUB slow-joiner: allow SUB to connect before clients register. + time.sleep(0.3) + + # Create 3 stage clients; each auto-registers on init. + clients = [ + OmniCoordClientForStage(router_addr, "tcp://stage:10001", "tcp://stage:10001-out", 0), + OmniCoordClientForStage(router_addr, "tcp://stage:10002", "tcp://stage:10002-out", 0), + OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1), + ] + + msg = _wait_for_instance_list(sub, expected_count=3) + assert msg is not None, "Expected InstanceList with 3 instances" + assert len(msg["instances"]) == 3 + assert isinstance(msg["timestamp"], (int, float)) + + input_addrs = {inst["input_addr"] for inst in msg["instances"]} + assert "tcp://stage:10001" in input_addrs + assert "tcp://stage:10002" in input_addrs + assert "tcp://stage:10003" in input_addrs + + for c in clients: + c.close() + coordinator.close() + sub.close(0) + sub_ctx.term() + + +def test_omni_coordinator_heartbeat_timeout_handling(): + """Verify that when a stage instance stops sending heartbeats, + OmniCoordinator marks it as unhealthy and excludes it from the active list. + """ + router_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + pub_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + coordinator = OmniCoordinator( + router_zmq_addr=router_addr, + pub_zmq_addr=pub_addr, + heartbeat_timeout=5.0, + ) + + sub_ctx = zmq.Context.instance() + sub = sub_ctx.socket(zmq.SUB) + sub.connect(pub_addr) + sub.setsockopt(zmq.SUBSCRIBE, b"") + + time.sleep(0.3) + + # A and B: real clients that send heartbeats every 5s. + client_a = OmniCoordClientForStage(router_addr, "tcp://stage:a", "tcp://stage:a-out", 0) + client_b = OmniCoordClientForStage(router_addr, "tcp://stage:b", "tcp://stage:b-out", 0) + + # C: raw DEALER that sends only registration, no heartbeat. + dealer_ctx = zmq.Context.instance() + dealer_c = dealer_ctx.socket(zmq.DEALER) + dealer_c.connect(router_addr) + reg_event = { + "input_addr": "tcp://stage:c", + "output_addr": "tcp://stage:c-out", + "stage_id": 0, + "event_type": "update", + "status": StageStatus.UP.value, + "queue_length": 0, + } + dealer_c.send(json.dumps(reg_event).encode("utf-8")) + + msg = _wait_for_instance_list(sub, expected_count=3) + assert msg is not None, "Expected initial 3 instances" + assert len(msg["instances"]) == 3 + + # Wait for heartbeat timeout (timeout=5s, check interval ~2.5s). + time.sleep(8.0) + + # Receive the update (C should be ERROR and excluded from active list). + msg_after_timeout = _wait_for_instance_list(sub, expected_count=2, timeout=5.0) + assert msg_after_timeout is not None, "Expected InstanceList with 2 instances after timeout" + instances = msg_after_timeout.get("instances", []) + input_addrs = {inst["input_addr"] for inst in instances} + + assert "tcp://stage:a" in input_addrs + assert "tcp://stage:b" in input_addrs + assert "tcp://stage:c" not in input_addrs + + client_a.close() + client_b.close() + dealer_c.close(0) + coordinator.close() + sub.close(0) + dealer_ctx.term() + sub_ctx.term() + + +def test_omni_coordinator_instance_shutdown_handling(): + """Verify that when a stage instance sends status='down', + OmniCoordinator removes it from the active list and broadcasts an updated list. + """ + router_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + pub_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + coordinator = OmniCoordinator( + router_zmq_addr=router_addr, + pub_zmq_addr=pub_addr, + heartbeat_timeout=1000.0, + ) + + sub_ctx = zmq.Context.instance() + sub = sub_ctx.socket(zmq.SUB) + sub.connect(pub_addr) + sub.setsockopt(zmq.SUBSCRIBE, b"") + + time.sleep(0.3) # PUB/SUB slow-joiner + + client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0) + + msg = _wait_for_instance_list(sub, expected_count=1) + assert msg is not None + assert len(msg["instances"]) == 1 + assert msg["instances"][0]["input_addr"] == "tcp://stage:shutdown" + + # Send down status (simulating graceful shutdown). + client.update_info(status=StageStatus.DOWN) + + # Receive updated list (should have 0 active instances). + msg = _wait_for_instance_list(sub, expected_count=0) + assert msg is not None + assert len(msg["instances"]) == 0 + + client.close() + coordinator.close() + sub.close(0) + sub_ctx.term() diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index 6e5f449c1..f3830f02e 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py @@ -6,7 +6,7 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.utils import GPUMemoryMonitor, hardware_test +from tests.utils import DeviceMemoryMonitor, hardware_test from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -23,9 +23,9 @@ models = ["riverclouds/qwen_image_random"] def inference(model_name: str, offload: bool = True): gc.collect() current_omni_platform.empty_cache() - device_index = torch.cuda.current_device() - torch.cuda.reset_peak_memory_stats(device=device_index) - monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02) + device_index = current_omni_platform.current_device() + current_omni_platform.reset_peak_memory_stats() + monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) monitor.start() m = Omni( model=model_name, @@ -33,7 +33,7 @@ def inference(model_name: str, offload: bool = True): # cache_backend="cache_dit", enable_cpu_offload=offload, ) - torch.cuda.reset_peak_memory_stats(device=device_index) + current_omni_platform.reset_peak_memory_stats() height = 256 width = 256 @@ -44,7 +44,7 @@ def inference(model_name: str, offload: bool = True): width=width, num_inference_steps=9, guidance_scale=0.0, - generator=torch.Generator("cuda").manual_seed(42), + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), ), ) peak = monitor.peak_used_mb @@ -70,6 +70,11 @@ def test_cpu_offload_diffusion_model(model_name: str): pytest.fail("Inference failed") print(f"Offload peak memory: {offload_peak_memory} MB") print(f"No offload peak memory: {no_offload_peak_memory} MB") - assert offload_peak_memory + 2500 < no_offload_peak_memory, ( - f"Offload peak memory {offload_peak_memory} MB should be less than no offload peak memory {no_offload_peak_memory} MB" + # Set platform-specific VRAM saving thresholds to account + # for varying runtime memory overhead and fragmentation between CUDA and ROCm. + is_rocm = torch.version.hip is not None + threshold = 2500 if not is_rocm else 2100 + assert offload_peak_memory + threshold < no_offload_peak_memory, ( + f"Offload peak memory {offload_peak_memory} MB should be less than " + f"no offload peak memory {no_offload_peak_memory} MB by {threshold} MB" ) diff --git a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py index beac14099..6132f1bd0 100644 --- a/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_layerwise_offload.py @@ -5,7 +5,7 @@ import pytest import torch from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.utils import GPUMemoryMonitor +from tests.utils import DeviceMemoryMonitor from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -28,11 +28,9 @@ def run_inference( layerwise_offload: bool = False, num_inference_steps: int = 3, ) -> float: - # For now, only support on GPU, so apply torch.cuda operations here - # NPU / ROCm platforms are expected to be detected and skipped this test function current_omni_platform.empty_cache() - device_index = torch.cuda.current_device() - monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02) + device_index = current_omni_platform.current_device() + monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) monitor.start() m = Omni( @@ -44,7 +42,7 @@ def run_inference( flow_shift=5.0, ) - torch.cuda.reset_peak_memory_stats(device=device_index) + current_omni_platform.reset_peak_memory_stats() # Refer to tests/e2e/offline_inference/test_t2v_model.py # Use minimal settings for testing @@ -57,7 +55,7 @@ def run_inference( OmniDiffusionSamplingParams( height=height, width=width, - generator=torch.Generator("cuda").manual_seed(42), + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), guidance_scale=1.0, num_inference_steps=num_inference_steps, num_frames=num_frames, diff --git a/tests/e2e/offline_inference/test_sequence_parallel.py b/tests/e2e/offline_inference/test_sequence_parallel.py index d66104aa1..0f0255b44 100644 --- a/tests/e2e/offline_inference/test_sequence_parallel.py +++ b/tests/e2e/offline_inference/test_sequence_parallel.py @@ -146,9 +146,6 @@ def _run_inference( # - warmup: whether to run warmup for this SP config # - is_perf_test: whether this is a performance test (show speedup metrics) SP_CONFIGS_L2 = [ - # Ulysses-2 - performance test - (2, 1, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True), - (1, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True), # Ring-2 - performance test # Hybrid - correctness only (2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False), ] diff --git a/tests/e2e/offline_inference/test_zimage_parallelism.py b/tests/e2e/offline_inference/test_zimage_parallelism.py index c43a5d224..63cb8743d 100644 --- a/tests/e2e/offline_inference/test_zimage_parallelism.py +++ b/tests/e2e/offline_inference/test_zimage_parallelism.py @@ -22,7 +22,7 @@ import torch from PIL import Image from vllm.distributed.parallel_state import cleanup_dist_env_and_memory -from tests.utils import GPUMemoryMonitor, hardware_test +from tests.utils import DeviceMemoryMonitor, hardware_test from vllm_omni import Omni from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -54,12 +54,13 @@ def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor: def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]: - """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1].""" + """Return (mean_abs_diff, p99_abs_diff) over RGB pixels in [0, 1].""" ta = _pil_to_float_rgb_tensor(a) tb = _pil_to_float_rgb_tensor(b) assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}" abs_diff = torch.abs(ta - tb) - return abs_diff.mean().item(), abs_diff.max().item() + p99_abs_diff = torch.quantile(abs_diff.flatten(), 0.99).item() + return abs_diff.mean().item(), p99_abs_diff def _extract_single_image(outputs) -> Image.Image: @@ -93,9 +94,9 @@ def _run_zimage_generate( if num_requests < 2: raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)") - torch.cuda.empty_cache() - device_index = torch.cuda.current_device() - monitor = GPUMemoryMonitor(device_index=device_index, interval=0.02) + current_omni_platform.empty_cache() + device_index = current_omni_platform.current_device() + monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02) monitor.start() m = Omni( model=_get_zimage_model(), @@ -161,8 +162,8 @@ def _run_zimage_generate( def test_zimage_tensor_parallel_tp2(tmp_path: Path): if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") - if not torch.cuda.is_available() or torch.cuda.device_count() < 2: - pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.") + if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: + pytest.skip("Z-Image TP=2 requires >= 2 devices.") enforce_eager = False @@ -196,18 +197,18 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): assert tp1_img.width == width and tp1_img.height == height assert tp2_img.width == width and tp2_img.height == height - mean_abs_diff, max_abs_diff = _diff_metrics(tp1_img, tp2_img) + mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img) mean_threshold = 3e-2 - max_threshold = 5e-1 + p99_threshold = 2.5e-1 print( "Z-Image TP image diff stats (TP=1 vs TP=2): " - f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; " - f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " + f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " f"tp1_img={tp1_path}, tp2_img={tp2_path}" ) - assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, ( - f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} " - f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e})" + assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( + f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" ) print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") @@ -219,59 +220,59 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): ) -@pytest.mark.integration -def test_zimage_vae_patch_parallel_tp2(tmp_path: Path): - if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): - pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA for now.") - if not torch.cuda.is_available() or torch.cuda.device_count() < 2: - pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 CUDA devices.") +# @pytest.mark.integration +# def test_zimage_vae_patch_parallel_tp2(tmp_path: Path): +# if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): +# pytest.skip("Z-Image VAE patch parallel e2e test is only supported on CUDA for now.") +# if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: +# pytest.skip("Z-Image VAE patch parallel TP=2 requires >= 2 devices.") - enforce_eager = False +# enforce_eager = False - # Use a larger image to ensure there are multiple VAE tiles. - height = 1152 - width = 1152 - num_inference_steps = 2 - seed = 42 +# # Use a larger image to ensure there are multiple VAE tiles. +# height = 1152 +# width = 1152 +# num_inference_steps = 2 +# seed = 42 - baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate( - tp_size=2, - height=height, - width=width, - num_inference_steps=num_inference_steps, - seed=seed, - enforce_eager=enforce_eager, - vae_use_tiling=True, - vae_patch_parallel_size=1, - num_requests=2, - ) - pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate( - tp_size=2, - height=height, - width=width, - num_inference_steps=num_inference_steps, - seed=seed, - enforce_eager=enforce_eager, - vae_use_tiling=True, - vae_patch_parallel_size=2, - num_requests=2, - ) +# baseline_img, _baseline_time_s, _baseline_peak_mem = _run_zimage_generate( +# tp_size=2, +# height=height, +# width=width, +# num_inference_steps=num_inference_steps, +# seed=seed, +# enforce_eager=enforce_eager, +# vae_use_tiling=True, +# vae_patch_parallel_size=1, +# num_requests=2, +# ) +# pp2_img, _pp2_time_s, _pp2_peak_mem = _run_zimage_generate( +# tp_size=2, +# height=height, +# width=width, +# num_inference_steps=num_inference_steps, +# seed=seed, +# enforce_eager=enforce_eager, +# vae_use_tiling=True, +# vae_patch_parallel_size=2, +# num_requests=2, +# ) - baseline_path = tmp_path / "zimage_tp2_vae_pp1.png" - pp2_path = tmp_path / "zimage_tp2_vae_pp2.png" - baseline_img.save(baseline_path) - pp2_img.save(pp2_path) +# baseline_path = tmp_path / "zimage_tp2_vae_pp1.png" +# pp2_path = tmp_path / "zimage_tp2_vae_pp2.png" +# baseline_img.save(baseline_path) +# pp2_img.save(pp2_path) - mean_abs_diff, max_abs_diff = _diff_metrics(baseline_img, pp2_img) - mean_threshold = 5e-3 - max_threshold = 1e-1 - print( - "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): " - f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; " - f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " - f"pp1_img={baseline_path}, pp2_img={pp2_path}" - ) - assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, ( - f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} " - f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e})" - ) +# mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img) +# mean_threshold = 5e-3 +# p99_threshold = 1e-1 +# print( +# "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): " +# f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " +# f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " +# f"pp1_img={baseline_path}, pp2_img={pp2_path}" +# ) +# assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( +# f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " +# f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" +# ) diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index fb9986592..9227928f3 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -8,6 +8,7 @@ from pathlib import Path import pytest from tests.conftest import ( + OmniServerParams, dummy_messages_from_mix_data, generate_synthetic_audio, generate_synthetic_image, @@ -51,7 +52,9 @@ else: stage_configs = [get_chunk_config()] # Create parameter combinations for model and stage config -test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +test_params = [ + OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs +] def get_system_prompt(): @@ -115,12 +118,11 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: "stream": True, "key_words": { "audio": ["test"], - "image": ["square", "quadrate"], }, } # Test single completion - openai_client.send_request(request_config) + openai_client.send_omni_request(request_config) @pytest.mark.advanced_model @@ -146,4 +148,4 @@ def test_text_to_text_001(omni_server, openai_client) -> None: "key_words": {"text": ["beijing"]}, } - openai_client.send_request(request_config, request_num=get_max_batch_size()) + openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index fd7e0704c..a28039888 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -12,6 +12,7 @@ from pathlib import Path import pytest from tests.conftest import ( + OmniServerParams, dummy_messages_from_mix_data, generate_synthetic_audio, generate_synthetic_image, @@ -50,7 +51,9 @@ def get_chunk_config(default_path): default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml") stage_configs = [default_path, get_chunk_config(default_path)] # Create parameter combinations for model and stage config -test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +test_params = [ + OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs +] def get_system_prompt(): @@ -106,7 +109,7 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: "key_words": {"text": ["beijing"]}, } - openai_client.send_request(request_config) + openai_client.send_omni_request(request_config) @pytest.mark.advanced_model @@ -128,7 +131,7 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None: "key_words": {"text": ["beijing"]}, } - openai_client.send_request(request_config, request_num=get_max_batch_size()) + openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @pytest.mark.advanced_model @@ -153,7 +156,7 @@ def test_image_to_text_001(omni_server, openai_client) -> None: "key_words": {"image": IMAGE_KEY}, } - openai_client.send_request(request_config) + openai_client.send_omni_request(request_config) @pytest.mark.advanced_model @@ -177,7 +180,7 @@ def test_image_to_audio_001(omni_server, openai_client) -> None: "key_words": {"image": IMAGE_KEY}, } - openai_client.send_request(request_config) + openai_client.send_omni_request(request_config) @pytest.mark.advanced_model @@ -201,7 +204,7 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None: "key_words": {"image": IMAGE_KEY}, } - openai_client.send_request(request_config, request_num=get_max_batch_size()) + openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @pytest.mark.advanced_model @@ -225,7 +228,7 @@ def test_video_to_text_001(omni_server, openai_client) -> None: "key_words": {"video": VIDEO_KEY}, } - openai_client.send_request(request_config) + openai_client.send_omni_request(request_config) @pytest.mark.advanced_model @@ -249,7 +252,7 @@ def test_video_to_audio_001(omni_server, openai_client) -> None: "key_words": {"video": VIDEO_KEY}, } - openai_client.send_request(request_config) + openai_client.send_omni_request(request_config) @pytest.mark.advanced_model @@ -273,7 +276,7 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None: "key_words": {"video": VIDEO_KEY}, } - openai_client.send_request(request_config, request_num=get_max_batch_size()) + openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @pytest.mark.advanced_model @@ -298,7 +301,7 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None: "key_words": {"audio": AUDIO_KEY}, } - openai_client.send_request(request_config) + openai_client.send_omni_request(request_config) @pytest.mark.advanced_model @@ -324,7 +327,7 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None: "key_words": {"image": IMAGE_KEY}, } - openai_client.send_request(request_config) + openai_client.send_omni_request(request_config) @pytest.mark.advanced_model @@ -351,7 +354,7 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None: "key_words": {"video": VIDEO_KEY}, } - openai_client.send_request(request_config) + openai_client.send_omni_request(request_config) @pytest.mark.skip(reason="There is a known issue with shape mismatch error.") @@ -383,4 +386,4 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: "stream": True, "key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY}, } - openai_client.send_request(request_config, request_num=get_max_batch_size()) + openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) diff --git a/tests/e2e/online_serving/test_qwen3_tts_websocket.py b/tests/e2e/online_serving/test_qwen3_tts_websocket.py new file mode 100644 index 000000000..df0514601 --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_tts_websocket.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E online tests for Qwen3-TTS WebSocket streaming speech. +""" + +import asyncio +import json +import os +from pathlib import Path + +import pytest +import websockets + +from tests.conftest import OmniServer +from tests.utils import hardware_test + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" + + +def get_stage_config() -> str: + return str( + Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml" + ) + + +@pytest.fixture(scope="module") +def omni_server(): + stage_config_path = get_stage_config() + + with OmniServer( + MODEL, + [ + "--stage-configs-path", + stage_config_path, + "--stage-init-timeout", + "120", + "--trust-remote-code", + "--enforce-eager", + "--disable-log-stats", + ], + env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1"}, + ) as server: + yield server + + +async def _run_ws_session(host: str, port: int) -> dict: + uri = f"ws://{host}:{port}/v1/audio/speech/stream" + starts: list[dict] = [] + dones: list[dict] = [] + chunk_lengths: dict[int, list[int]] = {} + session_done: dict | None = None + + async with websockets.connect(uri, max_size=None) as ws: + await ws.send( + json.dumps( + { + "type": "session.config", + "model": MODEL, + "voice": "vivian", + "language": "English", + "response_format": "pcm", + "stream_audio": True, + } + ) + ) + await ws.send( + json.dumps( + { + "type": "input.text", + "text": ( + "Hello, this is a websocket streaming test for Qwen three TTS, " + "and this sentence is intentionally long enough to produce audio chunks. " + "This is the second sentence." + ), + } + ) + ) + await ws.send(json.dumps({"type": "input.done"})) + + while True: + message = await asyncio.wait_for(ws.recv(), timeout=180) + if isinstance(message, bytes): + if not starts: + raise AssertionError("Received audio bytes before audio.start") + sentence_index = starts[-1]["sentence_index"] + chunk_lengths.setdefault(sentence_index, []).append(len(message)) + continue + + payload = json.loads(message) + msg_type = payload.get("type") + if msg_type == "audio.start": + starts.append(payload) + chunk_lengths.setdefault(payload["sentence_index"], []) + elif msg_type == "audio.done": + dones.append(payload) + elif msg_type == "session.done": + session_done = payload + break + elif msg_type == "error": + raise AssertionError(f"WebSocket error: {payload['message']}") + else: + raise AssertionError(f"Unexpected WebSocket message: {payload}") + + return { + "starts": starts, + "dones": dones, + "chunk_lengths": chunk_lengths, + "session_done": session_done, + } + + +class TestQwen3TTSWebSocket: + @pytest.mark.core_model + @pytest.mark.omni + @hardware_test(res={"cuda": "L4"}, num_cards=4) + def test_streaming_pcm_output(self, omni_server) -> None: + result = asyncio.run(_run_ws_session(omni_server.host, omni_server.port)) + + starts = result["starts"] + dones = result["dones"] + chunk_lengths = result["chunk_lengths"] + session_done = result["session_done"] + + assert session_done is not None + assert session_done["total_sentences"] == 2 + assert len(starts) == 2 + assert len(dones) == 2 + + for idx, start in enumerate(starts): + assert start["type"] == "audio.start" + assert start["sentence_index"] == idx + assert start["format"] == "pcm" + assert start["sample_rate"] == 24000 + assert start["sentence_text"] + + for done in dones: + sentence_index = done["sentence_index"] + total_bytes = done["total_bytes"] + assert done["error"] is False + assert total_bytes > 0 + assert chunk_lengths[sentence_index], f"Expected binary PCM frames for sentence {sentence_index}" + assert sum(chunk_lengths[sentence_index]) == total_bytes diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py new file mode 100644 index 000000000..4501569aa --- /dev/null +++ b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py @@ -0,0 +1,164 @@ +""" +Comprehensive tests of diffusion features that are available in online serving mode +and are supported by the following models: +- Qwen-Image-Edit: single image input +- Qwen-Image-Edit-2509: two image inputs +""" + +import pytest + +from tests.conftest import ( + OmniServer, + OmniServerParams, + OpenAIClientHandler, + dummy_messages_from_mix_data, + generate_synthetic_image, +) +from tests.utils import hardware_marks + +EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." +MULTI_EDIT_PROMPT = ( + "Transform the first image into a Dadaism collage art. " + "Transform the second image into a Vincent van Gogh style painting. " + "Then juxtapose the two transformed images into a single artwork for visual contrast." +) +NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist" +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) +PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) + + +# This test file targets two models, so I write a helper function. +# If a similar test only involves one model, one can just define a global list variable. +def _get_diffusion_feature_cases(model: str): + return [ + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--cache-backend", + "tea_cache", # [TODO] may consider changing to cache_dit after #1779 is resolved. Currently cache_dit and layerwise offload cannot work together. + "--enable-layerwise-offload", + ], + ), + id="single_card_001", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--cache-backend", + "cache_dit", + "--ulysses-degree", + "2", + ], + ), + id="parallel_001", + marks=PARALLEL_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--cache-backend", + "cache_dit", + "--ring", + "2", + ], + ), + id="parallel_002", + marks=PARALLEL_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--cache-backend", + "tea_cache", + "--cfg-parallel-size", + "2", + ], + ), + id="parallel_003", + marks=PARALLEL_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--cache-backend", + "cache_dit", + "--tensor-parallel-size", + "2", + "--vae-patch-parallel-size", + "2", + "--vae-use-tiling", + ], + ), + id="parallel_004", + marks=PARALLEL_FEATURE_MARKS, + ), + ] + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"), + indirect=True, +) +def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler): + """Test all diffusion features with Qwen-Image-Edit in regular end-user scenarios.""" + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + + messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT) + + # CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 2, + "negative_prompt": NEGATIVE_PROMPT, + "true_cfg_scale": 4.0, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"), + indirect=True, +) +def test_qwen_image_edit_2509(omni_server: OmniServer, openai_client: OpenAIClientHandler): + """Test all diffusion features with Qwen-Image-Edit-2509 in regular end-user scenarios.""" + image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + + messages = dummy_messages_from_mix_data( + image_data_url=[image_data_url_1, image_data_url_2], content_text=MULTI_EDIT_PROMPT + ) + + # CFG parallel is only activated when a negative prompt and true_cfg_scale > 1.0 are both present + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 2, + "negative_prompt": NEGATIVE_PROMPT, + "true_cfg_scale": 4.0, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 36c58b3d8..743e36a05 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -136,7 +136,6 @@ def mock_async_diffusion(mocker: MockerFixture): async def generate(**kwargs): # Return n PIL images wrapped in result object - print("!!!!!!!!!!!!!!!!!!!!! kwargs", kwargs) n = kwargs["sampling_params_list"][0].num_outputs_per_prompt mock.captured_sampling_params_list = kwargs["sampling_params_list"] mock.captured_prompt = kwargs["prompt"] diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 8b7177365..e90992317 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1,23 +1,29 @@ # tests/entrypoints/openai/test_serving_speech.py +import asyncio import logging import os +import struct from inspect import Signature, signature from pathlib import Path -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import numpy as np import pytest import torch -from fastapi import FastAPI, HTTPException, UploadFile +from fastapi import FastAPI, HTTPException, Request, UploadFile from fastapi.params import File, Form +from fastapi.responses import JSONResponse from fastapi.testclient import TestClient from pydantic import ValidationError from pytest_mock import MockerFixture +from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm_omni.entrypoints.openai import api_server as api_server_module from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio, OpenAICreateSpeechRequest from vllm_omni.entrypoints.openai.serving_speech import ( OmniOpenAIServingSpeech, + _create_wav_header, ) from vllm_omni.outputs import OmniRequestOutput @@ -916,17 +922,20 @@ class TestStreamingProtocolValidation: """Unit tests for the stream field validators in OpenAICreateSpeechRequest.""" def test_stream_validation_errors(self): - """stream=True requires response_format='pcm' and speed=1.0.""" - with pytest.raises(ValidationError, match="response_format='pcm'"): - OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav") + """stream=True requires response_format not in ('pcm', 'wav') and speed=1.0.""" + with pytest.raises(ValidationError, match="requires response_format not in \\('pcm', 'wav'\\)"): + OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="mp3") with pytest.raises(ValidationError, match="Speed adjustment is not supported"): OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm", speed=2.0) def test_stream_valid(self): - """stream=True + response_format='pcm' + speed=1.0 is accepted.""" + """stream=True + response_format in ('pcm', 'wav') + speed=1.0 is accepted.""" req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm") assert req.stream is True + req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav") + assert req.stream is True + def test_sse_stream_format_is_blocked(self): """stream_format='sse' is blocked.""" with pytest.raises(ValidationError, match="sse"): @@ -1017,3 +1026,256 @@ class TestStreamingResponse: response = client.post("/v1/audio/speech", json={"input": "Hello", "response_format": "wav"}) assert response.status_code == 200 assert "audio/wav" in response.headers["content-type"] + + +class TestAsyncOmniSupportedTasks: + """Test that AsyncOmni reports correct supported tasks based on output modalities.""" + + @pytest.mark.asyncio + async def test_tts_only_no_generate_task(self): + """TTS-only models (audio output, no text) should not include 'generate'.""" + from unittest.mock import MagicMock + + from vllm_omni.entrypoints.async_omni import AsyncOmni + + omni = AsyncOmni.__new__(AsyncOmni) + omni.output_modalities = [None, "audio"] + stage = MagicMock() + stage.is_comprehension = False + omni.stage_list = [stage] + tasks = await omni.get_supported_tasks() + assert "generate" not in tasks + assert "speech" in tasks + + @pytest.mark.asyncio + async def test_omni_model_includes_generate(self): + """Models with text output (e.g. Qwen3-Omni) should include 'generate'.""" + from unittest.mock import MagicMock + + from vllm_omni.entrypoints.async_omni import AsyncOmni + + omni = AsyncOmni.__new__(AsyncOmni) + omni.output_modalities = ["text", None, "audio"] + stage = MagicMock() + stage.is_comprehension = True + omni.stage_list = [stage] + tasks = await omni.get_supported_tasks() + assert "generate" in tasks + + +def test_api_server_create_speech_wraps_error_response_status(): + handler = MagicMock() + handler.create_speech = AsyncMock( + return_value=ErrorResponse( + error=ErrorInfo(message="bad request", type="BadRequestError", param=None, code=400), + ) + ) + + app = FastAPI() + app.state.openai_serving_speech = handler + scope = { + "type": "http", + "app": app, + "method": "POST", + "path": "/v1/audio/speech", + "headers": [], + "query_string": b"", + "client": ("127.0.0.1", 12345), + "server": ("testserver", 80), + "scheme": "http", + } + raw_request = Request(scope) + request = OpenAICreateSpeechRequest(input="Hello") + + response = asyncio.run(api_server_module.create_speech(request, raw_request)) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + + +class TestWAVHeaderGeneration: + """Unit tests for WAV header generation with placeholder values.""" + + def test_wav_header_basic_structure(self): + """Test basic WAV header structure with default parameters.""" + header = _create_wav_header(sample_rate=24000, num_channels=1, bits_per_sample=16) + + # Verify header length (should be 44 bytes) + assert len(header) == 44, f"Expected 44 bytes, got {len(header)}" + + # Parse and verify header structure + ( + chunk_id, + chunk_size, + format_type, + subchunk1_id, + subchunk1_size, + audio_format, + num_channels, + sample_rate, + byte_rate, + block_align, + bits_per_sample, + subchunk2_id, + subchunk2_size, + ) = struct.unpack("<4sI4s4sIHHIIHH4sI", header) + + # Verify RIFF header + assert chunk_id == b"RIFF", f"Expected RIFF, got {chunk_id}" + assert chunk_size == 0xFFFFFFFF, f"Expected placeholder 0xFFFFFFFF, got {chunk_size:#x}" + assert format_type == b"WAVE", f"Expected WAVE, got {format_type}" + + # Verify fmt chunk + assert subchunk1_id == b"fmt ", f"Expected 'fmt ', got {subchunk1_id}" + assert subchunk1_size == 16, f"Expected 16, got {subchunk1_size}" + assert audio_format == 1, f"Expected PCM (1), got {audio_format}" + assert num_channels == 1, f"Expected 1 channel, got {num_channels}" + assert sample_rate == 24000, f"Expected 24000 Hz, got {sample_rate}" + assert byte_rate == 48000, f"Expected 48000 byte/s, got {byte_rate}" + assert block_align == 2, f"Expected 2 bytes block align, got {block_align}" + assert bits_per_sample == 16, f"Expected 16 bits, got {bits_per_sample}" + + # Verify data chunk + assert subchunk2_id == b"data", f"Expected 'data', got {subchunk2_id}" + assert subchunk2_size == 0xFFFFFFFF, f"Expected placeholder 0xFFFFFFFF, got {subchunk2_size:#x}" + + def test_wav_header_different_sample_rates(self): + """Test WAV header with different sample rates.""" + test_cases = [ + (16000, 1, 16), + (22050, 1, 16), + (24000, 1, 16), + (44100, 1, 16), + (48000, 1, 16), + ] + + for sample_rate, num_channels, bits_per_sample in test_cases: + header = _create_wav_header(sample_rate, num_channels, bits_per_sample) + assert len(header) == 44, f"Header length mismatch for {sample_rate} Hz" + + # Parse sample rate from header + parsed_sample_rate = struct.unpack(" OmniRequestOutput: + chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000)) + + class MockCompletionOutput: + def __init__(self, index: int = 0): + self.index = index + self.text = "" + self.token_ids = [] + self.finish_reason = "stop" + self.stop_reason = None + self.logprobs = None + + class MockRequestOutput: + def __init__(self, audio_tensor: torch.Tensor): + self.request_id = "speech-wav-stream-test" + self.outputs = [MockCompletionOutput(index=0)] + self.multimodal_output = {"audio": audio_tensor, "sr": 24000} + self.finished = finished + self.prompt_token_ids = None + self.encoder_prompt_token_ids = None + self.num_cached_tokens = None + self.prompt_logprobs = None + self.kv_transfer_params = None + + return OmniRequestOutput( + stage_id=0, + final_output_type="audio", + request_output=MockRequestOutput(audio_tensor=chunk), + finished=finished, + ) + + async def mock_generate_streaming(*args, **kwargs): + yield _make_output(finished=False) + yield _make_output(finished=True) + + mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming) + mock_engine_client.default_sampling_params_list = [{}] + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + + speech_server = OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + + original_create_speech = speech_server.create_speech + sig = signature(original_create_speech) + new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"] + new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation) + + async def awaitable_create_speech(*args, **kwargs): + return await original_create_speech(*args, **kwargs) + + awaitable_create_speech.__signature__ = new_sig + speech_server.create_speech = awaitable_create_speech + + app = FastAPI() + app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None) + return app + + def test_wav_streaming_success(self, wav_streaming_app): + """Test WAV format streaming returns correct content type and includes WAV header.""" + client = TestClient(wav_streaming_app) + response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": "wav"}) + + assert response.status_code == 200 + assert "audio/wav" in response.headers["content-type"] + assert len(response.content) > 44 # Should have WAV header + audio data + + # Verify WAV header is present + header = response.content[:44] + chunk_id = header[0:4] + format_type = header[8:12] + assert chunk_id == b"RIFF", "Should start with RIFF" + assert format_type == b"WAVE", "Should contain WAVE format" + + def test_streaming_unsupported_format_rejected(self, wav_streaming_app): + """Test that unsupported formats are rejected for streaming.""" + client = TestClient(wav_streaming_app) + + unsupported_formats = ["mp3"] + for fmt in unsupported_formats: + response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": fmt}) + assert response.status_code == 422 diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py new file mode 100644 index 000000000..bd136ac72 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -0,0 +1,387 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import FastAPI, WebSocket +from starlette.testclient import TestClient +from starlette.websockets import WebSocketDisconnect + +from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _build_test_app(speech_service=None, *, idle_timeout=30.0, config_timeout=10.0): + if speech_service is None: + speech_service = MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav")) + speech_service._prepare_speech_generation = AsyncMock(return_value=("req-1", object(), {})) + + async def mock_generate_pcm_chunks(_generator, _request_id): + for chunk in (b"\x01\x02", b"\x03\x04\x05"): + yield chunk + + speech_service._generate_pcm_chunks = mock_generate_pcm_chunks + speech_service.engine_client = MagicMock() + speech_service.engine_client.abort = AsyncMock() + + handler = OmniStreamingSpeechHandler( + speech_service=speech_service, + idle_timeout=idle_timeout, + config_timeout=config_timeout, + ) + app = FastAPI() + + @app.websocket("/v1/audio/speech/stream") + async def ws_endpoint(websocket: WebSocket): + await handler.handle_session(websocket) + + return app, speech_service + + +class TestStreamingSpeechWebSocket: + def test_non_streaming_single_frame(self): + app, speech_service = _build_test_app() + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + + start = ws.receive_json() + assert start["type"] == "audio.start" + assert start["sentence_index"] == 0 + assert start["sentence_text"] == "Hello world." + assert start["format"] == "wav" + + audio = ws.receive_bytes() + assert audio.startswith(b"RIFF") + + done = ws.receive_json() + assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": len(audio), "error": False} + + ws.send_json({"type": "input.done"}) + session_done = ws.receive_json() + assert session_done == {"type": "session.done", "total_sentences": 1} + + assert speech_service._generate_audio_bytes.await_count == 1 + + def test_streaming_multiple_binary_frames(self): + captured_requests = [] + + speech_service = MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav")) + speech_service.engine_client = MagicMock() + speech_service.engine_client.abort = AsyncMock() + + async def mock_prepare_speech_generation(request): + captured_requests.append(request) + return "req-stream", object(), {} + + speech_service._prepare_speech_generation = mock_prepare_speech_generation + + async def mock_generate_pcm_chunks(_generator, _request_id): + for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"): + yield chunk + + speech_service._generate_pcm_chunks = mock_generate_pcm_chunks + app, _ = _build_test_app(speech_service) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "stream_audio": True, + "response_format": "pcm", + "initial_codec_chunk_frames": 12, + } + ) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + + start = ws.receive_json() + assert start["type"] == "audio.start" + assert start["format"] == "pcm" + assert start["sample_rate"] == 24000 + + assert ws.receive_bytes() == b"\x01\x02" + assert ws.receive_bytes() == b"\x03\x04\x05" + assert ws.receive_bytes() == b"\x06" + + done = ws.receive_json() + assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 6, "error": False} + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + + assert len(captured_requests) == 1 + assert captured_requests[0].stream is True + assert captured_requests[0].response_format == "pcm" + assert captured_requests[0].initial_codec_chunk_frames == 12 + assert speech_service._generate_audio_bytes.await_count == 0 + + def test_flush_on_input_done(self): + app, _ = _build_test_app() + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "Hello world without punctuation"}) + ws.send_json({"type": "input.done"}) + + assert ws.receive_json()["type"] == "audio.start" + assert ws.receive_bytes() + assert ws.receive_json() == { + "type": "audio.done", + "sentence_index": 0, + "total_bytes": 36, + "error": False, + } + assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + + def test_invalid_streaming_config(self): + app, _ = _build_test_app() + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "stream_audio": True, + "response_format": "wav", + } + ) + error = ws.receive_json() + assert error["type"] == "error" + assert "response_format='pcm'" in error["message"] + + def test_empty_input_text_emits_no_audio(self): + app, speech_service = _build_test_app() + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": ""}) + ws.send_json({"type": "input.done"}) + + assert ws.receive_json() == {"type": "session.done", "total_sentences": 0} + + assert speech_service._generate_audio_bytes.await_count == 0 + + def test_multiple_sentences_increment_indices(self): + app, _ = _build_test_app() + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "First sentence. Second sentence. "}) + + first_start = ws.receive_json() + assert first_start["sentence_index"] == 0 + ws.receive_bytes() + assert ws.receive_json() == { + "type": "audio.done", + "sentence_index": 0, + "total_bytes": 36, + "error": False, + } + + second_start = ws.receive_json() + assert second_start["sentence_index"] == 1 + ws.receive_bytes() + assert ws.receive_json() == { + "type": "audio.done", + "sentence_index": 1, + "total_bytes": 36, + "error": False, + } + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == {"type": "session.done", "total_sentences": 2} + + def test_unknown_message_type_keeps_session_open(self): + app, _ = _build_test_app() + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "unknown"}) + + error = ws.receive_json() + assert error == {"type": "error", "message": "Unknown message type: unknown"} + + ws.send_json({"type": "input.text", "text": "Hello world. "}) + assert ws.receive_json()["type"] == "audio.start" + ws.receive_bytes() + assert ws.receive_json() == { + "type": "audio.done", + "sentence_index": 0, + "total_bytes": 36, + "error": False, + } + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + + def test_config_timeout_closes_session(self): + app, _ = _build_test_app(config_timeout=0.01) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + error = ws.receive_json() + assert error == {"type": "error", "message": "Timeout waiting for session.config"} + + def test_generation_error_marks_audio_done(self): + speech_service = MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = AsyncMock(side_effect=RuntimeError("boom")) + speech_service._prepare_speech_generation = AsyncMock(return_value=("req-err", object(), {})) + speech_service._generate_pcm_chunks = AsyncMock() + speech_service.engine_client = MagicMock() + speech_service.engine_client.abort = AsyncMock() + app, _ = _build_test_app(speech_service) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + + assert ws.receive_json()["type"] == "audio.start" + assert ws.receive_json() == {"type": "error", "message": "Generation failed for sentence 0: boom"} + assert ws.receive_json() == {"type": "audio.done", "sentence_index": 0, "total_bytes": 0, "error": True} + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + + def test_streaming_generation_error_marks_audio_done(self): + speech_service = MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav")) + speech_service._prepare_speech_generation = AsyncMock(return_value=("req-stream-err", object(), {})) + speech_service.engine_client = MagicMock() + speech_service.engine_client.abort = AsyncMock() + + async def mock_generate_pcm_chunks(_generator, _request_id): + yield b"\x01\x02" + raise RuntimeError("stream boom") + + speech_service._generate_pcm_chunks = mock_generate_pcm_chunks + app, _ = _build_test_app(speech_service) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "stream_audio": True, + "response_format": "pcm", + } + ) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + + assert ws.receive_json()["type"] == "audio.start" + assert ws.receive_bytes() == b"\x01\x02" + assert ws.receive_json() == { + "type": "error", + "message": "Generation failed for sentence 0: stream boom", + } + assert ws.receive_json() == { + "type": "audio.done", + "sentence_index": 0, + "total_bytes": 2, + "error": True, + } + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + + def test_invalid_input_text_type_returns_validation_error(self): + app, speech_service = _build_test_app() + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": 123}) + + assert ws.receive_json() == { + "type": "error", + "message": "input.text requires a string value", + } + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == {"type": "session.done", "total_sentences": 0} + + assert speech_service._generate_audio_bytes.await_count == 0 + + def test_input_text_message_too_large(self, monkeypatch): + monkeypatch.setattr(streaming_speech_module, "_MAX_INPUT_TEXT_MESSAGE_SIZE", 32) + app, speech_service = _build_test_app() + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "x" * 128}) + + assert ws.receive_json() == { + "type": "error", + "message": "input.text message too large", + } + + ws.send_json({"type": "input.done"}) + assert ws.receive_json() == {"type": "session.done", "total_sentences": 0} + + assert speech_service._generate_audio_bytes.await_count == 0 + + def test_session_config_message_too_large(self, monkeypatch): + monkeypatch.setattr(streaming_speech_module, "_MAX_CONFIG_MESSAGE_SIZE", 64) + app, _ = _build_test_app() + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian", "ref_audio": "x" * 512}) + + assert ws.receive_json() == { + "type": "error", + "message": "session.config message too large", + } + + def test_disconnect_aborts_streaming_request(self): + speech_service = MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav")) + speech_service._prepare_speech_generation = AsyncMock(return_value=("req-abort", object(), {})) + speech_service.engine_client = MagicMock() + speech_service.engine_client.abort = AsyncMock() + + async def mock_generate_pcm_chunks(_generator, _request_id): + yield b"\x01\x02" + + speech_service._generate_pcm_chunks = mock_generate_pcm_chunks + handler = OmniStreamingSpeechHandler(speech_service=speech_service) + + websocket = MagicMock() + websocket.send_json = AsyncMock(side_effect=[None, WebSocketDisconnect()]) + websocket.send_bytes = AsyncMock(side_effect=WebSocketDisconnect()) + + config = MagicMock() + config.model = None + config.voice = "Vivian" + config.task_type = None + config.language = None + config.instructions = None + config.response_format = "pcm" + config.speed = 1.0 + config.max_new_tokens = None + config.initial_codec_chunk_frames = None + config.ref_audio = None + config.ref_text = None + config.x_vector_only_mode = None + config.stream_audio = True + + with pytest.raises(WebSocketDisconnect): + asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0)) + + speech_service.engine_client.abort.assert_awaited_once_with("req-abort") + assert websocket.send_json.await_count == 2 diff --git a/tests/entrypoints/openai_api/test_text_splitter.py b/tests/entrypoints/openai_api/test_text_splitter.py new file mode 100644 index 000000000..23d4d191f --- /dev/null +++ b/tests/entrypoints/openai_api/test_text_splitter.py @@ -0,0 +1,253 @@ +"""Tests for SentenceSplitter used in streaming TTS input.""" + +import pytest + +from vllm_omni.entrypoints.openai.text_splitter import SentenceSplitter + +pytestmark = [pytest.mark.openai, pytest.mark.speech] + + +class TestSentenceSplitterEnglish: + """Tests for English sentence splitting.""" + + def test_single_sentence_no_boundary(self): + splitter = SentenceSplitter() + result = splitter.add_text("Hello world") + assert result == [] + assert splitter.buffer == "Hello world" + + def test_single_sentence_with_boundary(self): + splitter = SentenceSplitter() + result = splitter.add_text("Hello world. How are you?") + assert len(result) == 1 + assert result[0] == "Hello world." + + def test_multiple_sentences(self): + splitter = SentenceSplitter() + result = splitter.add_text("Hello. How are you? I am fine! ") + assert len(result) == 3 + assert result[0] == "Hello." + assert result[1] == "How are you?" + assert result[2] == "I am fine!" + + def test_exclamation_mark(self): + splitter = SentenceSplitter() + result = splitter.add_text("Wow, that is great! Tell me more.") + assert len(result) == 1 + assert result[0] == "Wow, that is great!" + + def test_question_mark(self): + splitter = SentenceSplitter() + result = splitter.add_text("Can you hear me? I hope so.") + assert len(result) == 1 + assert result[0] == "Can you hear me?" + + +class TestSentenceSplitterChinese: + """Tests for CJK sentence splitting.""" + + def test_chinese_period(self): + splitter = SentenceSplitter() + result = splitter.add_text("δ½ ε₯½δΈ–η•Œγ€‚δ½ ε₯½ε—") + assert len(result) == 1 + assert result[0] == "δ½ ε₯½δΈ–η•Œγ€‚" + + def test_chinese_exclamation(self): + splitter = SentenceSplitter() + result = splitter.add_text("ε€ͺε₯½δΊ†οΌθ°’θ°’δ½ ") + assert len(result) == 1 + assert result[0] == "ε€ͺε₯½δΊ†οΌ" + + def test_chinese_question(self): + splitter = SentenceSplitter() + result = splitter.add_text("δ½ ζ˜―θ°οΌŸζˆ‘ζ˜―ε°ζ˜Ž") + assert len(result) == 1 + assert result[0] == "你是谁？" + + def test_chinese_comma_no_split(self): + """Chinese commas are clause-level and should not trigger a split.""" + splitter = SentenceSplitter() + result = splitter.add_text("δ½ ε₯½οΌŒδΈ–η•Œ") + assert result == [] + assert splitter.buffer == "δ½ ε₯½οΌŒδΈ–η•Œ" + + def test_chinese_semicolon_no_split(self): + """Chinese semicolons are clause-level and should not trigger a split.""" + splitter = SentenceSplitter() + result = splitter.add_text("η¬¬δΈ€η‚ΉοΌ›η¬¬δΊŒη‚Ή") + assert result == [] + assert splitter.buffer == "η¬¬δΈ€η‚ΉοΌ›η¬¬δΊŒη‚Ή" + + def test_chinese_multiple(self): + splitter = SentenceSplitter() + result = splitter.add_text("δ½ ε₯½οΌδ½ ε₯½ε—οΌŸζˆ‘εΎˆε₯½γ€‚") + assert len(result) == 3 + assert result[0] == "δ½ ε₯½οΌ" + assert result[1] == "δ½ ε₯½ε—οΌŸ" + assert result[2] == "ζˆ‘εΎˆε₯½γ€‚" + + +class TestSentenceSplitterMixed: + """Tests for mixed-language sentence splitting.""" + + def test_mixed_english_chinese(self): + splitter = SentenceSplitter() + result = splitter.add_text("HelloδΈ–η•Œγ€‚How are you? ") + assert len(result) == 2 + assert result[0] == "HelloδΈ–η•Œγ€‚" + assert result[1] == "How are you?" + + +class TestSentenceSplitterIncremental: + """Tests for incremental (multi-chunk) text input.""" + + def test_accumulation_across_chunks(self): + splitter = SentenceSplitter() + # First chunk: no boundary + result1 = splitter.add_text("Hello ") + assert result1 == [] + + # Second chunk: completes a sentence + result2 = splitter.add_text("world. How") + assert len(result2) == 1 + assert result2[0] == "Hello world." + assert splitter.buffer == "How" + + def test_word_by_word(self): + splitter = SentenceSplitter() + words = ["Hello, ", "how ", "are ", "you? ", "I ", "am ", "fine."] + all_sentences = [] + for word in words: + all_sentences.extend(splitter.add_text(word)) + + assert len(all_sentences) == 1 + assert all_sentences[0] == "Hello, how are you?" + # "I am fine." stays in buffer (no trailing whitespace after period) + + def test_three_chunks(self): + splitter = SentenceSplitter() + splitter.add_text("The quick brown ") + splitter.add_text("fox jumps. ") + result = splitter.add_text("Over the lazy dog. ") + # "The quick brown fox jumps." should have been returned on second chunk + # "Over the lazy dog." on third chunk + assert len(result) == 1 + assert result[0] == "Over the lazy dog." + + +class TestSentenceSplitterFlush: + """Tests for flush behavior.""" + + def test_flush_returns_remaining(self): + splitter = SentenceSplitter() + splitter.add_text("Hello world") + result = splitter.flush() + assert result == "Hello world" + assert splitter.buffer == "" + + def test_flush_empty_buffer(self): + splitter = SentenceSplitter() + result = splitter.flush() + assert result is None + + def test_flush_after_sentence(self): + splitter = SentenceSplitter() + splitter.add_text("Hello world. Remaining text") + result = splitter.flush() + assert result == "Remaining text" + + def test_flush_whitespace_only(self): + splitter = SentenceSplitter() + splitter.add_text("Hello. ") + # "Hello." extracted, buffer is " " + result = splitter.flush() + # Whitespace-only should return None + assert result is None + + def test_flush_clears_buffer(self): + splitter = SentenceSplitter() + splitter.add_text("some text") + splitter.flush() + assert splitter.buffer == "" + # Second flush should return None + assert splitter.flush() is None + + +class TestSentenceSplitterEdgeCases: + """Edge case tests.""" + + def test_empty_input(self): + splitter = SentenceSplitter() + result = splitter.add_text("") + assert result == [] + assert splitter.buffer == "" + + def test_none_like_empty(self): + """Empty string should not affect buffer.""" + splitter = SentenceSplitter() + splitter.add_text("Hello") + splitter.add_text("") + assert splitter.buffer == "Hello" + + def test_only_punctuation(self): + splitter = SentenceSplitter() + result = splitter.add_text(". ") + # "." is 1 char, below default min_sentence_length of 2 + # It will be carried forward + assert result == [] + + def test_min_sentence_length(self): + splitter = SentenceSplitter(min_sentence_length=10) + result = splitter.add_text("Hi. Hello world. ") + # "Hi." is 3 chars (< 10), so it gets carried to "Hello world." + assert len(result) == 1 + assert "Hi." in result[0] + assert "Hello world." in result[0] + + def test_short_segments_are_carried_until_long_enough(self): + splitter = SentenceSplitter(min_sentence_length=10) + result = splitter.add_text("Hi. Ok. Hello there. ") + assert result == ["Hi.Ok.Hello there."] + assert splitter.buffer == "" + + def test_min_sentence_length_zero(self): + splitter = SentenceSplitter(min_sentence_length=0) + result = splitter.add_text("A. B. ") + assert len(result) == 2 + + def test_no_boundary_then_flush(self): + splitter = SentenceSplitter() + result = splitter.add_text("Hello world how are you") + assert result == [] + flushed = splitter.flush() + assert flushed == "Hello world how are you" + + def test_consecutive_punctuation(self): + splitter = SentenceSplitter() + result = splitter.add_text("Really?! Yes, really. ") + assert len(result) >= 1 + + def test_reuse_after_flush(self): + """Splitter can be reused after flush.""" + splitter = SentenceSplitter() + splitter.add_text("First session.") + splitter.flush() + + result = splitter.add_text("Second session. More text") + assert len(result) == 1 + assert result[0] == "Second session." + assert splitter.buffer == "More text" + + +class TestSentenceSplitterBufferLimit: + """Tests for buffer overflow protection.""" + + def test_buffer_overflow_raises(self): + from vllm_omni.entrypoints.openai.text_splitter import _MAX_BUFFER_SIZE + + splitter = SentenceSplitter() + # Fill buffer just under the limit + splitter.add_text("x" * (_MAX_BUFFER_SIZE - 1)) + # One more char should exceed the limit + with pytest.raises(ValueError, match="exceeded maximum size"): + splitter.add_text("xx") diff --git a/tests/entrypoints/test_omni_diffusion.py b/tests/entrypoints/test_omni_diffusion.py index 90b0fd053..9e555fa85 100644 --- a/tests/entrypoints/test_omni_diffusion.py +++ b/tests/entrypoints/test_omni_diffusion.py @@ -345,7 +345,7 @@ def _setup_multiprocessing_mocks(monkeypatch: pytest.MonkeyPatch, mocker: Mocker fake_process_instance = mocker.MagicMock() fake_process_instance.start = mocker.MagicMock() fake_process_instance.join = mocker.MagicMock() - fake_process_instance.is_alive = mocker.MagicMock(return_value=False) + fake_process_instance.is_alive = mocker.MagicMock(return_value=True) fake_process_instance.terminate = mocker.MagicMock() fake_process_class.return_value = fake_process_instance diff --git a/tests/examples/online_serving/test_qwen2_5_omni.py b/tests/examples/online_serving/test_qwen2_5_omni.py index 8e08d5bc5..58b011f19 100644 --- a/tests/examples/online_serving/test_qwen2_5_omni.py +++ b/tests/examples/online_serving/test_qwen2_5_omni.py @@ -12,7 +12,7 @@ from pathlib import Path import pytest -from tests.conftest import convert_audio_file_to_text, cosine_similarity_text +from tests.conftest import OmniServerParams, convert_audio_file_to_text, cosine_similarity_text from tests.utils import hardware_test models = ["Qwen/Qwen2.5-Omni-7B"] @@ -22,7 +22,11 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen2_5_omni") # Create parameter combinations for model and stage config -test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs] +test_params = [ + OmniServerParams(model=model, port=8091, stage_config_path=stage_config) + for model in models + for stage_config in stage_configs +] def run_cmd(command): diff --git a/tests/examples/online_serving/test_qwen3_omni.py b/tests/examples/online_serving/test_qwen3_omni.py index c7354905e..568f1f9b4 100644 --- a/tests/examples/online_serving/test_qwen3_omni.py +++ b/tests/examples/online_serving/test_qwen3_omni.py @@ -12,7 +12,7 @@ from pathlib import Path import pytest -from tests.conftest import convert_audio_file_to_text, cosine_similarity_text +from tests.conftest import OmniServerParams, convert_audio_file_to_text, cosine_similarity_text from tests.utils import hardware_test models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] @@ -22,7 +22,11 @@ stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_config example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen3_omni") # Create parameter combinations for model and stage config -test_params = [(8091, model, stage_config) for model in models for stage_config in stage_configs] +test_params = [ + OmniServerParams(model=model, port=8091, stage_config_path=stage_config) + for model in models + for stage_config in stage_configs +] def run_cmd(command): diff --git a/tests/perf/scripts/run_benchmark.py b/tests/perf/scripts/run_benchmark.py index eb2ac5359..ed2fab83e 100644 --- a/tests/perf/scripts/run_benchmark.py +++ b/tests/perf/scripts/run_benchmark.py @@ -107,8 +107,14 @@ def omni_server(request): print("OmniServer stopped") -def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_prompt) -> Any: - """Generate synthetic image with random values.""" +def run_benchmark( + args: list, + test_name: str, + flow, + dataset_name: str, + num_prompt, +) -> Any: + """Run a single benchmark iteration and return the parsed result JSON.""" current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_{current_dt}.json" if "--result-filename" in args: @@ -117,10 +123,6 @@ def run_benchmark(args: list, test_name: str, flow, dataset_name: str, num_promp ["vllm", "bench", "serve", "--omni"] + args + [ - "--backend", - "openai-chat-omni", - "--endpoint", - "/v1/chat/completions", "--save-result", "--result-dir", os.environ.get("BENCHMARK_DIR", "tests"), @@ -196,7 +198,10 @@ def benchmark_params(request, omni_server): total = len(all_params) print(f"\n Running benchmark {current}/{total} for {test_name}") - return {"test_name": test_name, "params": all_params[param_index]} + return { + "test_name": test_name, + "params": all_params[param_index], + } def assert_result(result, params, num_prompt): @@ -266,7 +271,11 @@ def test_performance_benchmark(omni_server, benchmark_params): for qps, num_prompt in zip(qps_list, num_prompt_list): args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)] result = run_benchmark( - args=args, test_name=test_name, flow=qps, dataset_name=dataset_name, num_prompt=num_prompt + args=args, + test_name=test_name, + flow=qps, + dataset_name=dataset_name, + num_prompt=num_prompt, ) assert_result(result, params, num_prompt=num_prompt) @@ -274,6 +283,10 @@ def test_performance_benchmark(omni_server, benchmark_params): for concurrency, num_prompt in zip(max_concurrency_list, num_prompt_list): args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"] result = run_benchmark( - args=args, test_name=test_name, flow=concurrency, dataset_name=dataset_name, num_prompt=num_prompt + args=args, + test_name=test_name, + flow=concurrency, + dataset_name=dataset_name, + num_prompt=num_prompt, ) assert_result(result, params, num_prompt=num_prompt) diff --git a/tests/perf/stage_configs/qwen3_tts.yaml b/tests/perf/stage_configs/qwen3_tts.yaml new file mode 100644 index 000000000..4ba4e6e83 --- /dev/null +++ b/tests/perf/stage_configs/qwen3_tts.yaml @@ -0,0 +1,99 @@ +# Stage config for running Qwen3-TTS with 2-stage architecture +# Stage 0: Talker (text -> 8-layer RVQ codec codes) +# Stage 1: Code2Wav (codec codes -> audio waveform) +# +# The following config has been verified on 1x H100-80G GPU. +async_chunk: true +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 4 + engine_args: + model_stage: qwen3_tts + model_arch: Qwen3TTSTalkerForConditionalGeneration + hf_overrides: + architectures: [Qwen3TTSTalkerForConditionalGeneration] + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: false + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.3 + distributed_executor_backend: "mp" + max_num_batched_tokens: 512 + max_model_len: 4096 + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + detokenize: false + repetition_penalty: 1.05 + stop_token_ids: [2150] + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 4 + engine_args: + model_stage: code2wav + model_arch: Qwen3TTSCode2Wav + hf_overrides: + architectures: [Qwen3TTSCode2Wav] + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.2 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 32768 + engine_input_source: [0] + final_output: true + final_output_type: audio + input_connectors: + from_stage_0: connector_of_shared_memory + tts_args: + max_instructions_length: 500 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + detokenize: true + repetition_penalty: 1.0 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 4 + + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 25 + + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/tests/perf/tests/test.json b/tests/perf/tests/test.json index 88041c119..65ef6588d 100644 --- a/tests/perf/tests/test.json +++ b/tests/perf/tests/test.json @@ -8,6 +8,8 @@ "benchmark_params": [ { "dataset_name": "random", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", "num_prompts": [ 10, 40 @@ -28,6 +30,8 @@ }, { "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", "num_prompts": [ 10, 40 @@ -88,6 +92,8 @@ "benchmark_params": [ { "dataset_name": "random", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", "num_prompts": [ 10, 40 @@ -108,6 +114,8 @@ }, { "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", "num_prompts": [ 10, 40 @@ -140,5 +148,38 @@ } } ] + }, + { + "test_name": "test_qwen3_tts", + "server_params": { + "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", + "stage_config_name": "qwen3_tts.yaml" + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "num_prompts": [ + 10, + 40 + ], + "max_concurrency": [ + 1, + 4 + ], + "random_input_len": 100, + "random_output_len": 100, + "extra_body": { + "voice": "Vivian", + "language": "English" + }, + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_audio_ttfp_ms": 100000, + "mean_audio_rtf": 100000 + } + } + ] } ] diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py new file mode 100644 index 000000000..21ef22523 --- /dev/null +++ b/tests/test_config_factory.py @@ -0,0 +1,570 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Unit tests for StageConfigFactory and related classes. +""" + +from vllm_omni.config.stage_config import ( + ModelPipeline, + StageConfig, + StageConfigFactory, + StageType, +) + + +class TestStageType: + """Tests for StageType enum.""" + + def test_stage_type_values(self): + """Test StageType enum values.""" + assert StageType.LLM.value == "llm" + assert StageType.DIFFUSION.value == "diffusion" + + def test_stage_type_from_string(self): + """Test creating StageType from string.""" + assert StageType("llm") == StageType.LLM + assert StageType("diffusion") == StageType.DIFFUSION + + +class TestStageConfig: + """Tests for StageConfig dataclass.""" + + def test_minimal_config(self): + """Test creating StageConfig with minimal required fields.""" + config = StageConfig(stage_id=0, model_stage="thinker") + assert config.stage_id == 0 + assert config.model_stage == "thinker" + assert config.stage_type == StageType.LLM + assert config.input_sources == [] + assert config.final_output is False + assert config.worker_type is None + + def test_full_config(self): + """Test creating StageConfig with all fields.""" + config = StageConfig( + stage_id=1, + model_stage="talker", + stage_type=StageType.LLM, + input_sources=[0], + custom_process_input_func="module.path.func", + final_output=True, + final_output_type="audio", + worker_type="ar", + scheduler_cls="path.to.Scheduler", + hf_config_name="talker_config", + is_comprehension=False, + ) + assert config.stage_id == 1 + assert config.model_stage == "talker" + assert config.input_sources == [0] + assert config.final_output_type == "audio" + assert config.worker_type == "ar" + + def test_to_omegaconf_basic(self): + """Test converting StageConfig to OmegaConf format.""" + config = StageConfig( + stage_id=0, + model_stage="thinker", + stage_type=StageType.LLM, + worker_type="ar", + final_output=True, + final_output_type="text", + ) + omega_config = config.to_omegaconf() + + assert omega_config.stage_id == 0 + assert omega_config.stage_type == "llm" + assert omega_config.engine_args.model_stage == "thinker" + assert omega_config.engine_args.worker_type == "ar" + assert omega_config.final_output is True + assert omega_config.final_output_type == "text" + # Legacy field name for backward compatibility + assert omega_config.engine_input_source == [] + + def test_to_omegaconf_with_runtime_overrides(self): + """Test that runtime overrides are applied to OmegaConf output.""" + config = StageConfig( + stage_id=0, + model_stage="thinker", + runtime_overrides={ + "gpu_memory_utilization": 0.9, + "tensor_parallel_size": 2, + "devices": "0,1", + "max_batch_size": 64, + }, + ) + omega_config = config.to_omegaconf() + + assert omega_config.engine_args.gpu_memory_utilization == 0.9 + assert omega_config.engine_args.tensor_parallel_size == 2 + assert omega_config.runtime.devices == "0,1" + assert omega_config.runtime.max_batch_size == 64 + + +class TestModelPipeline: + """Tests for ModelPipeline class.""" + + def test_valid_linear_dag(self): + """Test validation of a valid linear DAG.""" + stages = [ + StageConfig(stage_id=0, model_stage="thinker", input_sources=[]), + StageConfig(stage_id=1, model_stage="talker", input_sources=[0]), + StageConfig(stage_id=2, model_stage="code2wav", input_sources=[1]), + ] + pipeline = ModelPipeline(model_type="test", stages=stages) + errors = pipeline.validate_pipeline() + assert errors == [], f"Unexpected errors: {errors}" + + def test_valid_branching_dag(self): + """Test validation of a valid branching DAG.""" + stages = [ + StageConfig(stage_id=0, model_stage="input", input_sources=[]), + StageConfig(stage_id=1, model_stage="branch_a", input_sources=[0]), + StageConfig(stage_id=2, model_stage="branch_b", input_sources=[0]), + ] + pipeline = ModelPipeline(model_type="test", stages=stages) + errors = pipeline.validate_pipeline() + assert errors == [], f"Unexpected errors: {errors}" + + def test_missing_entry_point(self): + """Test that missing entry point is detected.""" + stages = [ + StageConfig(stage_id=0, model_stage="stage_a", input_sources=[1]), + StageConfig(stage_id=1, model_stage="stage_b", input_sources=[0]), + ] + pipeline = ModelPipeline(model_type="test", stages=stages) + errors = pipeline.validate_pipeline() + assert any("entry point" in e.lower() for e in errors) + + def test_missing_dependency(self): + """Test that missing stage reference is detected.""" + stages = [ + StageConfig(stage_id=0, model_stage="input", input_sources=[]), + StageConfig(stage_id=1, model_stage="output", input_sources=[99]), # Invalid + ] + pipeline = ModelPipeline(model_type="test", stages=stages) + errors = pipeline.validate_pipeline() + assert any("non-existent" in e.lower() for e in errors) + + def test_duplicate_stage_ids(self): + """Test that duplicate stage IDs are detected.""" + stages = [ + StageConfig(stage_id=0, model_stage="stage_a", input_sources=[]), + StageConfig(stage_id=0, model_stage="stage_b", input_sources=[]), # Duplicate + ] + pipeline = ModelPipeline(model_type="test", stages=stages) + errors = pipeline.validate_pipeline() + assert any("duplicate" in e.lower() for e in errors) + + def test_self_reference(self): + """Test that self-references are detected.""" + stages = [ + StageConfig(stage_id=0, model_stage="entry", input_sources=[]), + StageConfig(stage_id=1, model_stage="self_ref", input_sources=[1]), # Self + ] + pipeline = ModelPipeline(model_type="test", stages=stages) + errors = pipeline.validate_pipeline() + assert any("itself" in e.lower() for e in errors) + + def test_get_stage_by_id(self): + """Test getting stage by ID.""" + stages = [ + StageConfig(stage_id=0, model_stage="thinker", input_sources=[]), + StageConfig(stage_id=1, model_stage="talker", input_sources=[0]), + ] + pipeline = ModelPipeline(model_type="test", stages=stages) + + stage = pipeline.get_stage(1) + assert stage is not None + assert stage.model_stage == "talker" + + missing = pipeline.get_stage(99) + assert missing is None + + def test_empty_pipeline(self): + """Test validation of empty pipeline.""" + pipeline = ModelPipeline(model_type="test", stages=[]) + errors = pipeline.validate_pipeline() + assert any("no stages" in e.lower() for e in errors) + + +class TestStageConfigFactory: + """Tests for StageConfigFactory class.""" + + def test_default_diffusion_no_yaml(self): + """Test single-stage diffusion works without YAML config (@ZJY0516).""" + kwargs = { + "cache_backend": "none", + "cache_config": None, + "dtype": "bfloat16", + } + configs = StageConfigFactory.create_default_diffusion(kwargs) + + assert len(configs) == 1 + cfg = configs[0] + assert cfg["stage_id"] == 0 + assert cfg["stage_type"] == "diffusion" + assert cfg["final_output"] is True + assert cfg["final_output_type"] == "image" + + def test_default_diffusion_with_parallel_config(self): + """Test diffusion config calculates devices from parallel_config.""" + + class MockParallelConfig: + world_size = 4 + + kwargs = { + "parallel_config": MockParallelConfig(), + "cache_backend": "tea_cache", + } + configs = StageConfigFactory.create_default_diffusion(kwargs) + + assert configs[0]["runtime"]["devices"] == "0,1,2,3" + + def test_per_stage_override_precedence(self): + """Test that --stage-0-gpu-memory-utilization overrides global.""" + stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[]) + cli_overrides = { + "gpu_memory_utilization": 0.5, # Global + "stage_0_gpu_memory_utilization": 0.9, # Per-stage override + } + + overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides) + + # Per-stage should override global + assert overrides["gpu_memory_utilization"] == 0.9 + + def test_cli_override_forwards_engine_registered_args(self): + """Test that any engine-registered CLI arg is forwarded (@wuhang2014).""" + stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[]) + cli_overrides = { + "gpu_memory_utilization": 0.9, # Well-known param + "custom_engine_flag": True, # Not in _INTERNAL_KEYS, so forwarded + } + + overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides) + + assert overrides["gpu_memory_utilization"] == 0.9 + assert overrides["custom_engine_flag"] is True + + def test_cli_override_excludes_internal_keys(self): + """Test that internal/orchestrator keys are not forwarded.""" + stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[]) + cli_overrides = { + "gpu_memory_utilization": 0.9, + "model": "some_model", # Internal + "stage_configs_path": "/path", # Internal + "batch_timeout": 10, # Internal + } + + overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides) + + assert overrides["gpu_memory_utilization"] == 0.9 + assert "model" not in overrides + assert "stage_configs_path" not in overrides + assert "batch_timeout" not in overrides + + def test_per_stage_override_excludes_internal_keys(self): + """Test that per-stage overrides also skip internal keys.""" + stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[]) + cli_overrides = { + "stage_0_gpu_memory_utilization": 0.9, + "stage_0_model": "override_model", # Internal, should be skipped + "stage_0_batch_timeout": 5, # Internal, should be skipped + } + + overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides) + + assert overrides["gpu_memory_utilization"] == 0.9 + assert "model" not in overrides + assert "batch_timeout" not in overrides + + +class TestPipelineYamlParsing: + """Tests for pipeline YAML file parsing (@ZJY0516).""" + + def test_parse_qwen3_omni_moe_yaml(self, tmp_path): + """Test parsing the qwen3_omni_moe pipeline YAML.""" + yaml_content = """\ +model_type: qwen3_omni_moe +async_chunk: false + +stages: + - stage_id: 0 + model_stage: thinker + stage_type: llm + input_sources: [] + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + hf_config_name: thinker_config + final_output: true + final_output_type: text + is_comprehension: true + + - stage_id: 1 + model_stage: talker + stage_type: llm + input_sources: [0] + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + hf_config_name: talker_config + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker + + - stage_id: 2 + model_stage: code2wav + stage_type: llm + input_sources: [1] + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + hf_config_name: thinker_config + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav + final_output: true + final_output_type: audio +""" + yaml_file = tmp_path / "qwen3_omni_moe.yaml" + yaml_file.write_text(yaml_content) + + pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "qwen3_omni_moe") + + assert pipeline.model_type == "qwen3_omni_moe" + assert len(pipeline.stages) == 3 + + # Stage 0: thinker + s0 = pipeline.stages[0] + assert s0.stage_id == 0 + assert s0.model_stage == "thinker" + assert s0.stage_type == StageType.LLM + assert s0.input_sources == [] + assert s0.worker_type == "ar" + assert s0.final_output is True + assert s0.final_output_type == "text" + assert s0.is_comprehension is True + + # Stage 1: talker + s1 = pipeline.stages[1] + assert s1.stage_id == 1 + assert s1.input_sources == [0] + assert s1.custom_process_input_func == ( + "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker" + ) + assert s1.final_output is False + + # Stage 2: code2wav + s2 = pipeline.stages[2] + assert s2.stage_id == 2 + assert s2.input_sources == [1] + assert s2.worker_type == "generation" + assert s2.final_output_type == "audio" + + def test_parse_yaml_with_legacy_engine_input_source(self, tmp_path): + """Test backward compatibility with engine_input_source field.""" + yaml_content = """\ +model_type: legacy_model + +stages: + - stage_id: 0 + model_stage: entry + stage_type: llm + - stage_id: 1 + model_stage: downstream + stage_type: llm + engine_input_source: [0] +""" + yaml_file = tmp_path / "legacy.yaml" + yaml_file.write_text(yaml_content) + + pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_model") + assert pipeline.stages[1].input_sources == [0] + + def test_parse_yaml_with_connectors_and_edges(self, tmp_path): + """Test parsing pipeline with optional connectors and edges.""" + yaml_content = """\ +model_type: test_model + +stages: + - stage_id: 0 + model_stage: entry + stage_type: llm + input_sources: [] + +connectors: + type: ray + +edges: + - from: 0 + to: 1 +""" + yaml_file = tmp_path / "with_connectors.yaml" + yaml_file.write_text(yaml_content) + + pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_model") + assert pipeline.connectors == {"type": "ray"} + assert pipeline.edges == [{"from": 0, "to": 1}] + + def test_parsed_pipeline_passes_validation(self, tmp_path): + """Test that a well-formed YAML produces a valid pipeline.""" + yaml_content = """\ +model_type: valid_model + +stages: + - stage_id: 0 + model_stage: entry + stage_type: llm + input_sources: [] + final_output: true + final_output_type: text + - stage_id: 1 + model_stage: next + stage_type: llm + input_sources: [0] +""" + yaml_file = tmp_path / "valid.yaml" + yaml_file.write_text(yaml_content) + + pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "valid_model") + errors = pipeline.validate_pipeline() + assert errors == [], f"Unexpected validation errors: {errors}" + + def test_parse_diffusion_stage_type(self, tmp_path): + """Test parsing a diffusion stage type from YAML.""" + yaml_content = """\ +model_type: diff_model + +stages: + - stage_id: 0 + model_stage: dit + stage_type: diffusion + input_sources: [] + final_output: true + final_output_type: image +""" + yaml_file = tmp_path / "diffusion.yaml" + yaml_file.write_text(yaml_content) + + pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "diff_model") + assert pipeline.stages[0].stage_type == StageType.DIFFUSION + + def test_parse_mimo_audio_yaml(self, tmp_path): + """Test parsing the MiMo Audio pipeline YAML.""" + yaml_content = """\ +model_type: mimo_audio + +stages: + - stage_id: 0 + model_stage: fused_thinker_talker + stage_type: llm + input_sources: [] + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + is_comprehension: true + final_output: true + final_output_type: text + + - stage_id: 1 + model_stage: code2wav + stage_type: llm + input_sources: [0] + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav + final_output: true + final_output_type: audio +""" + yaml_file = tmp_path / "mimo_audio.yaml" + yaml_file.write_text(yaml_content) + + pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "mimo_audio") + + assert pipeline.model_type == "mimo_audio" + assert len(pipeline.stages) == 2 + + # Stage 0: fused_thinker_talker + s0 = pipeline.stages[0] + assert s0.stage_id == 0 + assert s0.model_stage == "fused_thinker_talker" + assert s0.input_sources == [] + assert s0.worker_type == "ar" + assert s0.is_comprehension is True + assert s0.final_output is True + assert s0.final_output_type == "text" + + # Stage 1: code2wav + s1 = pipeline.stages[1] + assert s1.stage_id == 1 + assert s1.input_sources == [0] + assert s1.worker_type == "generation" + assert s1.custom_process_input_func == ( + "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav" + ) + assert s1.final_output is True + assert s1.final_output_type == "audio" + + # Pipeline validation + errors = pipeline.validate_pipeline() + assert errors == [], f"Unexpected errors: {errors}" + + +class TestAsyncChunk: + """Tests for async_chunk pipeline flag.""" + + def test_model_pipeline_async_chunk_default(self): + """Test ModelPipeline defaults async_chunk to False.""" + stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])] + pipeline = ModelPipeline(model_type="test", stages=stages) + assert pipeline.async_chunk is False + + def test_model_pipeline_async_chunk_set(self): + """Test ModelPipeline with async_chunk=True.""" + stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])] + pipeline = ModelPipeline(model_type="test", stages=stages, async_chunk=True) + assert pipeline.async_chunk is True + + def test_parse_async_chunk_from_yaml(self, tmp_path): + """Test that async_chunk is parsed from pipeline YAML.""" + yaml_content = """\ +model_type: test_async +async_chunk: true + +stages: + - stage_id: 0 + model_stage: entry + stage_type: llm + input_sources: [] +""" + yaml_file = tmp_path / "async.yaml" + yaml_file.write_text(yaml_content) + + pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_async") + assert pipeline.async_chunk is True + + def test_parse_missing_async_chunk_defaults_false(self, tmp_path): + """Test that missing async_chunk defaults to False.""" + yaml_content = """\ +model_type: test_no_async + +stages: + - stage_id: 0 + model_stage: entry + stage_type: llm + input_sources: [] +""" + yaml_file = tmp_path / "no_async.yaml" + yaml_file.write_text(yaml_content) + + pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_no_async") + assert pipeline.async_chunk is False + + +class TestArchitectureFallback: + """Tests for architecture-based model detection fallback.""" + + def test_architecture_models_mapping_exists(self): + """Test that _ARCHITECTURE_MODELS contains expected entries.""" + assert "MiMoAudioForConditionalGeneration" in StageConfigFactory._ARCHITECTURE_MODELS + assert StageConfigFactory._ARCHITECTURE_MODELS["MiMoAudioForConditionalGeneration"] == "mimo_audio" + assert "HunyuanImage3ForCausalMM" in StageConfigFactory._ARCHITECTURE_MODELS + assert StageConfigFactory._ARCHITECTURE_MODELS["HunyuanImage3ForCausalMM"] == "hunyuan_image3" + + def test_mimo_audio_in_pipeline_models(self): + """Test that mimo_audio is registered in PIPELINE_MODELS.""" + assert "mimo_audio" in StageConfigFactory.PIPELINE_MODELS diff --git a/tests/utils.py b/tests/utils.py index 3219821b6..e8a9d5c05 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -20,6 +20,8 @@ from typing_extensions import ParamSpec from vllm.platforms import current_platform from vllm.utils.torch_utils import cuda_device_count_stateless +from vllm_omni.platforms import current_omni_platform + _P = ParamSpec("_P") if current_platform.is_rocm(): @@ -441,27 +443,11 @@ def npu_marks(*, res: str, num_cards: int): return [mark for mark in [test_platform, test_resource, test_distributed] if mark is not None] -def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): +def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): """ - Decorate a test for multiple hardware platforms with a single call. - Automatically wraps the test with @create_new_process_for_each_test() for distributed tests. - - Args: - res: Mapping from platform to resource type. Supported platforms/resources: - - cuda: L4, H100 - - rocm: MI325 - - npu: A2, A3 - num_cards: Number of cards required. Can be: - - int: same card count for all platforms (default: 1) - - dict: per-platform card count, e.g., {"cuda": 2, "rocm": 2} - - Example: - @hardware_test( - res={"cuda": "L4", "rocm": "MI325", "npu": "A2"}, - num_cards={"cuda": 2, "rocm": 2, "npu": 2}, - ) - def test_multi_platform(): - ... + Get a collection of pytest marks to apply for `@hardware_test`, + including CUDA, ROCm, and NPU, + based on the specified platforms and resources. """ # Validate platforms # Don't validate platform details in this decorator @@ -484,7 +470,7 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): num_cards_dict[platform] = 1 # Collect marks from all platforms - all_marks: list[Callable[[Callable[_P, None]], Callable[_P, None]]] = [] + all_marks: list[pytest.MarkDecorator] = [] for platform, resource in res.items(): cards = num_cards_dict[platform] if platform == "cuda" or platform == "rocm": @@ -494,6 +480,32 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): else: raise ValueError(f"Unsupported platform: {platform}") all_marks.extend(marks) + return all_marks + + +def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): + """ + Decorate a test for multiple hardware platforms with a single call. + Automatically wraps the test with @create_new_process_for_each_test() for distributed tests. + + Args: + res: Mapping from platform to resource type. Supported platforms/resources: + - cuda: L4, H100 + - rocm: MI325 + - npu: A2, A3 + num_cards: Number of cards required. Can be: + - int: same card count for all platforms (default: 1) + - dict: per-platform card count, e.g., {"cuda": 2, "rocm": 2} + + Example: + @hardware_test( + res={"cuda": "L4", "rocm": "MI325", "npu": "A2"}, + num_cards={"cuda": 2, "rocm": 2, "npu": 2}, + ) + def test_multi_platform(): + ... + """ + all_marks = hardware_marks(res=res, num_cards=num_cards) def wrapper(f: Callable[_P, None]) -> Callable[_P, None]: func = f @@ -504,8 +516,8 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): return wrapper -class GPUMemoryMonitor: - """Poll global device memory usage via CUDA APIs.""" +class DeviceMemoryMonitor: + """Poll global device memory usage.""" def __init__(self, device_index: int, interval: float = 0.05): self.device_index = device_index @@ -518,8 +530,8 @@ class GPUMemoryMonitor: def monitor_loop() -> None: while not self._stop_event.is_set(): try: - with torch.cuda.device(self.device_index): - free_bytes, total_bytes = torch.cuda.mem_get_info() + with current_omni_platform.device(self.device_index): + free_bytes, total_bytes = current_omni_platform.mem_get_info() used_mb = (total_bytes - free_bytes) / (1024**2) self._peak_used_mb = max(self._peak_used_mb, used_mb) except Exception: @@ -537,8 +549,8 @@ class GPUMemoryMonitor: @property def peak_used_mb(self) -> float: - fallback_alloc = torch.cuda.max_memory_allocated(device=self.device_index) / (1024**2) - fallback_reserved = torch.cuda.max_memory_reserved(device=self.device_index) / (1024**2) + fallback_alloc = current_omni_platform.max_memory_allocated(device=self.device_index) / (1024**2) + fallback_reserved = current_omni_platform.max_memory_reserved(device=self.device_index) / (1024**2) return max(self._peak_used_mb, fallback_alloc, fallback_reserved) def __del__(self): diff --git a/tools/nightly/buildkite_testcase_statistics.py b/tools/nightly/buildkite_testcase_statistics.py new file mode 100644 index 000000000..5f3217981 --- /dev/null +++ b/tools/nightly/buildkite_testcase_statistics.py @@ -0,0 +1,808 @@ +#!/usr/bin/env python3 +""" +Parse pytest commands from Buildkite test-ready.yml, test-merge.yml, test-nightly.yml; +collect test cases (including parametrized) via pytest --collect-only -q and produce an HTML report. + +Usage (run from repo root): + python scripts/buildkite_testcase_statistics.py -o buildkite_testcase_statistics.html + +Requires: PyYAML (pip install pyyaml) +""" + +from __future__ import annotations + +import argparse +import ast +import html +import re +import subprocess +import sys +from pathlib import Path + +import yaml + +# Repo root (parent of the directory containing this script) +REPO_ROOT = Path(__file__).resolve().parent.parent.parent +BUILDKITE_DIR = REPO_ROOT / ".buildkite" +PIPELINE_FILES = ["test-ready.yml", "test-merge.yml", "test-nightly.yml"] + + +def load_yaml(path: Path) -> dict: + with open(path, encoding="utf-8") as f: + return yaml.safe_load(f) + + +PYTEST_CMD_RE = re.compile(r"(?:timeout\s+\S+\s+)?(?:python3? -m\s+)?pytest\s+[^\n&|;]*") + + +def normalize_commands(step: dict) -> list[str]: + raw = step.get("commands") or [] + if isinstance(raw, str): + raw = [raw] + text = "\n".join((c.strip() if isinstance(c, str) else "") for c in raw if c) + text = text.replace("$$", "$") + result = [] + for m in PYTEST_CMD_RE.finditer(text): + line = m.group(0).strip() + line_start = text.rfind("\n", 0, m.start()) + 1 + before_match = text[line_start : m.start()] + if before_match.lstrip().startswith("#"): + continue + if line: + result.append(line) + return result + + +def extract_pytest_targets_from_line(line: str) -> list[str]: + """ + Parse pytest test targets (file path or marker expression) from a single command line. + Returns e.g. ['tests/xxx.py']. + When there is no explicit test file but only -m/--run-level, returns + ['-m expr'] or ['-m expr --run-level level'] with the full marker args. + """ + if "pytest" not in line: + return [] + # Strip leading timeout and anything before pytest + idx = line.find("pytest") + if idx == -1: + return [] + rest = line[idx:].strip() + if not rest.startswith("pytest"): + return [] + # Simple split on whitespace (no handling of quoted spaces) + parts = re.split(r"\s+", rest, maxsplit=1) + args = parts[1].strip() if len(parts) > 1 else "" + if not args: + return [] + targets = [] + # Match tests/.../*.py + for m in re.finditer(r"tests/[^\s'\"#]+\.py", args): + path = m.group(0) + if path not in targets: + targets.append(path) + if not targets: + # May be -m/--run-level only; record full marker arguments instead of a generic marker. + extra = _parse_extra_args_from_line(line) + if extra: + # Join into a single string, e.g. "-m core and cpu --run-level nightly" + targets.append(" ".join(extra)) + return targets + + +def get_pytest_targets_from_step(step: dict) -> list[tuple[str, str]]: + """ + Extract all pytest-related targets from a step. + Returns [(target, raw_line), ...] where target is tests/xxx.py or '(marker/run-level)'. + """ + lines = normalize_commands(step) + out = [] + seen = set() + for line in lines: + for t in extract_pytest_targets_from_line(line): + if t not in seen: + seen.add(t) + out.append((t, line)) + return out + + +def _parse_extra_args_from_line(raw_line: str) -> list[str]: + """Extract -m and --run-level args from a raw pytest command line. Strips extra spaces and quotes. + Correctly handles quoted expressions with spaces, e.g. -m 'core_model and cpu'. + """ + extra: list[str] = [] + + # -m: match -m followed by '...', "..." or a single token (so -m expression is one argv) + m = re.search(r"-m\s+(?:'([^']*)'|\"([^\"]*)\"|(\S+))", raw_line) + if m: + value = (m.group(1) or m.group(2) or m.group(3) or "").strip() + if value: + extra.extend(["-m", value]) + + r = re.search(r"--run-level\s+(?:'([^']*)'|\"([^\"]*)\"|(\S+))", raw_line) + if r: + value = (r.group(1) or r.group(2) or r.group(3) or "").strip() + if value: + extra.extend(["--run-level", value]) + + return extra + + +def _parse_collect_only_stdout(stdout: str, *, raise_on_empty: bool = True, stderr: str = "") -> list[str]: + """ + Parse pytest --collect-only -q stdout into a list of node ids. + Normalizes \\r\\n and \\r to \\n before splitting. + If raise_on_empty and no node ids found, raises RuntimeError; otherwise returns []. + """ + raw = (stdout or "").replace("\r\n", "\n").replace("\r", "\n") + out: list[str] = [] + for line in raw.splitlines(): + line = line.strip() + if not line or "::" not in line or "test_" not in line: + continue + out.append(line) + if not out and raise_on_empty: + raise RuntimeError( + f"Failed to parse pytest --collect-only output: no node ids found.\nstdout:\n{stdout}\n\nstderr:\n{stderr}" + ) + return out + + +def _resolve_pytest_target(target: str, repo_root: Path, extra: list[str], raw_line: str) -> tuple[list[str], str, int]: + """ + Resolve target + repo_root + extra to (path_args, fallback_path, timeout). + Raises FileNotFoundError/RuntimeError on invalid config. + """ + if target.endswith(".py"): + # Target extracted from buildkite yaml may include glob patterns like + # tests/e2e/online_serving/test_*_expansion.py. Pytest does not reliably + # expand globs across platforms, so expand them here. + rel = Path(target) + if any(ch in target for ch in ["*", "?", "["]): + matches = sorted(repo_root.glob(rel.as_posix())) + if not matches: + raise FileNotFoundError(f"Pytest target glob did not match any files: {repo_root / target}") + return [str(p.resolve()) for p in matches], target.replace("\\", "/"), 90 + + path = (repo_root / rel).resolve() + if not path.exists(): + raise FileNotFoundError(f"Pytest target path does not exist: {path}") + return [str(path)], target.replace("\\", "/"), 60 + if not extra: + raise RuntimeError(f"Failed to parse -m/--run-level from pytest line: {raw_line!r}") + if not (repo_root / "tests").exists(): + raise FileNotFoundError("tests/ directory not found under repo root") + return ["tests/"], "tests/", 120 + + +def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]: + """ + Run pytest --collect-only -q for a target and return test node id list (incl. parametrized). + target: path like tests/foo.py or marker string (e.g. "-m expr"). raw_line used for -m/--run-level. + """ + extra = _parse_extra_args_from_line(raw_line) + path_args, _fallback, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line) + + cmd_quiet = [sys.executable, "-m", "pytest", *path_args, "--collect-only", "-q"] + cmd_quiet.extend(extra) + result = subprocess.run( + cmd_quiet, + cwd=repo_root, + capture_output=True, + text=True, + timeout=timeout_quiet, + ) + if result.returncode != 0: + raise RuntimeError(f"pytest --collect-only failed for {path_args} with args {extra}: {result.stderr.strip()}") + print(f"pytest --collect-only success for {path_args} with args {extra}: {result.stdout.strip()}") + return _parse_collect_only_stdout(result.stdout or "", stderr=result.stderr or "") + + +def _extra_with_skip_marker(extra: list[str]) -> list[str]: + """Build extra args that add -m skip: if extra has -m expr then use 'expr and skip', else -m skip.""" + out: list[str] = [] + i = 0 + has_m = False + while i < len(extra): + if extra[i] == "-m" and i + 1 < len(extra): + out.extend(["-m", extra[i + 1] + " and skip"]) + has_m = True + i += 2 + elif extra[i] == "--run-level" and i + 1 < len(extra): + out.extend(["--run-level", extra[i + 1]]) + i += 2 + else: + i += 1 + if not has_m: + out = ["-m", "skip"] + out + return out + + +def get_skip_status(target: str, raw_line: str, repo_root: Path) -> set[str]: + """ + Use pytest -m skip --collect-only -q to collect tests marked with skip; return their node ids. + Does not run tests; only collects. Supports both pytest xxx.py and pytest -m "marker". + """ + extra = _parse_extra_args_from_line(raw_line) + try: + path_args, _, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line) + except (FileNotFoundError, RuntimeError): + return set() + + extra_skip = _extra_with_skip_marker(extra) + cmd = [sys.executable, "-m", "pytest", *path_args, "--collect-only", "-q"] + cmd.extend(extra_skip) + try: + result = subprocess.run( + cmd, + cwd=repo_root, + capture_output=True, + text=True, + timeout=timeout_quiet, + ) + except (subprocess.TimeoutExpired, FileNotFoundError): + return set() + if result.returncode != 0: + return set() + ids = _parse_collect_only_stdout(result.stdout or "", raise_on_empty=False) + return set(ids) + + +def get_docstring_for_node_id(repo_root: Path, node_id: str) -> str: + """ + Parse the docstring of the test_ function/method from source by pytest node id. + Supports path::test_func, path::TestClass::test_method, path::test_func[param], etc. + """ + if not node_id or "::" not in node_id or node_id.startswith("["): + return "" + parts = node_id.split("::") + # Path accepts forward slashes; no platform-specific replacement needed + path = (repo_root / parts[0]).resolve() + if not path.exists(): + return "" + # Strip parametrize suffix, e.g. test_foo[param] -> test_foo + func_part = parts[-1] + if "[" in func_part: + func_part = func_part.split("[", 1)[0] + try: + source = path.read_text(encoding="utf-8", errors="replace") + tree = ast.parse(source) + except Exception: + return "" + # parts: [file, maybe_class, func] or [file, func] + if len(parts) == 2: + # path::test_func + target_func = func_part + target_class = None + else: + # path::TestClass::test_method + target_class = parts[1] + target_func = func_part + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == target_func and target_class is None: + doc = ast.get_docstring(node) + return (doc or "").strip() + if isinstance(node, ast.ClassDef) and node.name == target_class: + for item in node.body: + if isinstance(item, ast.FunctionDef) and item.name == target_func: + doc = ast.get_docstring(item) + return (doc or "").strip() + break + return "" + + +def get_test_file_for_node_id(node_id: str) -> str: + """Extract test file path from a pytest node id.""" + if not node_id: + return "" + if "::" in node_id: + return node_id.split("::", 1)[0] + if node_id.endswith(".py"): + return node_id + return "(marker-only / no collection)" + + +def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Path | None = None) -> None: + """Write stats to HTML with styling and responsive tables.""" + root = repo_root if repo_root is not None else REPO_ROOT + pipeline_badges = { + "test-ready": "ready", + "test-merge": "merge", + "test-nightly": "nightly", + } + + def esc(s: str) -> str: + return html.escape(str(s)) + + def tr_class(i: int) -> str: + return "odd" if i % 2 == 0 else "even" + + pipeline_totals: dict[str, int] = {} + for item in all_stats: + _, pipeline_name, _, _, _, _, count, _, _, _ = item + pipeline_totals[pipeline_name] = pipeline_totals.get(pipeline_name, 0) + count + pipeline_summary_rows = [] + for idx, (pipeline_name, pipeline_count) in enumerate(sorted(pipeline_totals.items())): + badge = pipeline_badges.get(pipeline_name, "") + badge_span = f'{esc(pipeline_name)}' if badge else esc(pipeline_name) + pipeline_summary_rows.append( + f'{badge_span}{pipeline_count}' + ) + + # Summary rows by module (item: ..., names, docstrings, skipped_ids) + summary_rows = [] + for i, item in enumerate(all_stats): + _, pipeline_name, label, test_files, _, marker_targets, count, _, _, _ = item + if test_files: + files_html = "
".join(esc(f) for f in test_files) + elif marker_targets: + # Display full marker / run-level arguments when no explicit test files. + files_html = "
".join(esc(m) for m in marker_targets) + else: + files_html = "No explicit tests or markers" + badge = pipeline_badges.get(pipeline_name, "") + badge_span = f'{esc(pipeline_name)}' if badge else esc(pipeline_name) + summary_rows.append( + f'' + f'{badge_span}{esc(label)}{count}' + f'{files_html}' + ) + + # Detail rows grouped by Pipeline, then by Test Suite (two-level collapsible) + case_details_thead = """ + + Pipeline + Test Suite + Test File + No. + Test Name + Status + Description + """ + # Per-item: (pipeline_name, label, count, badge_span, suite_rows_html) + pipeline_suites: dict[str, list[tuple[str, int, str, str]]] = {} # pipeline -> [(label, count, summary, body)] + for item in all_stats: + _, pipeline_name, label, _, _, _, count, names, docstrings, skipped_ids = item + badge = pipeline_badges.get(pipeline_name, "") + badge_span = f'{esc(pipeline_name)}' if badge else esc(pipeline_name) + suite_summary = f"{esc(label)} ({count} cases)" + file_sections: list[str] = [] + if not names: + suite_rows = [ + f'' + f"{badge_span}{esc(label)}" + f'(marker-only / no collection)0' + f'No collection or only -m/--run-level' + f'' + ] + file_sections.append( + '
' + "(marker-only / no collection)" + f'
{case_details_thead}' + f"{''.join(suite_rows)}
" + ) + else: + file_groups: dict[str, list[tuple[str, str, bool]]] = {} + for idx, name in enumerate(names): + desc = (docstrings[idx] if idx < len(docstrings) else "") or "" + if not desc and "::" in name and not name.startswith("["): + desc = get_docstring_for_node_id(root, name) + file_key = get_test_file_for_node_id(name) + file_groups.setdefault(file_key, []).append((name, desc, name in skipped_ids)) + + for file_name in sorted(file_groups.keys()): + suite_rows = [] + for i, (name, desc, is_skipped) in enumerate(file_groups[file_name], 1): + desc_html = esc(desc).replace("\n", "
") if desc else "" + status_html = 'Skipped' if is_skipped else "β€”" + suite_rows.append( + f'' + f'{badge_span}{esc(label)}{esc(file_name)}' + f'{i}' + f'{esc(name)}' + f'{status_html}{desc_html}' + ) + file_sections.append( + f'
' + f"{esc(file_name)} ({len(file_groups[file_name])} cases)" + f'
{case_details_thead}' + f"{''.join(suite_rows)}
" + ) + body = "".join(file_sections) + if pipeline_name not in pipeline_suites: + pipeline_suites[pipeline_name] = [] + pipeline_suites[pipeline_name].append((label, count, suite_summary, body)) + + # Build two-level HTML: Pipeline (outer) -> Test Suite (inner) -> table + detail_sections_html_parts = [] + for pipeline_name in sorted(pipeline_badges.keys()): + if pipeline_name not in pipeline_suites: + continue + suites = pipeline_suites[pipeline_name] + pipeline_total = sum(c for _, c, _, _ in suites) + badge = pipeline_badges.get(pipeline_name, "") + badge_span = f'{esc(pipeline_name)}' if badge else esc(pipeline_name) + pipeline_summary = f"{badge_span} ({pipeline_total} cases)" + inner_html = "".join( + f'
{suite_summary}' + f'
{body}
' + for _, _, suite_summary, body in suites + ) + detail_sections_html_parts.append( + f'
{pipeline_summary}' + f'
{inner_html}
' + ) + # Include pipelines that might be in all_stats but not in pipeline_badges (e.g. custom names) + for pipeline_name in sorted(pipeline_suites.keys()): + if pipeline_name in pipeline_badges: + continue + suites = pipeline_suites[pipeline_name] + pipeline_total = sum(c for _, c, _, _ in suites) + pipeline_summary = f"{esc(pipeline_name)} ({pipeline_total} cases)" + inner_html = "".join( + f'
{suite_summary}' + f'
{body}
' + for _, _, suite_summary, body in suites + ) + detail_sections_html_parts.append( + f'
{pipeline_summary}' + f'
{inner_html}
' + ) + detail_sections_html = "".join(detail_sections_html_parts) + + html_content = f""" + + + + + Buildkite Test Case Statistics + + + +
+

Buildkite Pytest Case Statistics

+

Test Suites: {len(all_stats)} Β· Total cases (sum of steps): {total} Β·

+ +
+

Pipeline Summary

+
+ + + + + + {"".join(pipeline_summary_rows)} + +
PipelineTotal Cases
+
+
+ +
+

Summary by Test Suite

+
+ + + + + + {"".join(summary_rows)} + +
PipelineTest SuiteCasesTest Files
+
+
+ +
+

Case Details

+

Expand by Pipeline, then by Test Suite, then by Test File.

+
+ + + Filter applies to the Test File fold sections below. +
+ {detail_sections_html} +
+
+ ↑ Top + + + +""" + out_path.write_text(html_content, encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Count pytest cases from Buildkite pipelines") + parser.add_argument( + "--buildkite-dir", + type=Path, + default=BUILDKITE_DIR, + help="Path to .buildkite directory", + ) + default_out = REPO_ROOT / "buildkite_testcase_statistics.html" + parser.add_argument( + "-o", + "--output", + type=Path, + default=default_out, + help=f"Output HTML file path (default: {default_out})", + ) + args = parser.parse_args() + + all_stats = [] + + for filename in PIPELINE_FILES: + yml_path = args.buildkite_dir / filename + if not yml_path.exists(): + print(f"Skipped (not found): {yml_path}", file=sys.stderr) + continue + data = load_yaml(yml_path) + steps = data.get("steps") or [] + pipeline_name = filename.replace(".yml", "") + + for step in steps: + if not isinstance(step, dict): + continue + label = step.get("label") or "(no label)" + # Skip commented-out steps (usually no label or commands) + if not step.get("commands"): + continue + targets = get_pytest_targets_from_step(step) + if not targets: + continue + + module_key = f"{pipeline_name} / {label}" + # test_files: only real .py files; marker-only steps will have an empty list. + test_files = [t for t, _ in targets if t.endswith(".py")] + # marker_targets: synthetic marker targets (e.g. "-m expr", "-m expr --run-level level"). + marker_targets = [t for t, _ in targets if not t.endswith(".py")] + # marker_only: all targets are marker expressions (no explicit .py files). + marker_only = bool(targets) and all(not t.endswith(".py") for t, _ in targets) + + names = [] + skipped_ids: set[str] = set() + for target, raw in targets: + collected = collect_test_names(target, REPO_ROOT, raw) + skipped_ids.update(get_skip_status(target, raw, REPO_ROOT)) + names.extend(collected) + count = len(names) - len(skipped_ids) + docstrings = [ + get_docstring_for_node_id(REPO_ROOT, n) if "::" in n and not n.startswith("[") else "" for n in names + ] + all_stats.append( + ( + module_key, + pipeline_name, + label, + test_files, + marker_only, + marker_targets, + count, + names, + docstrings, + skipped_ids, + ) + ) + + out_path = args.output.resolve() + out_path.parent.mkdir(parents=True, exist_ok=True) + total = sum(item[6] for item in all_stats) + write_html(all_stats, out_path, total, REPO_ROOT) + print(f"Written: {out_path}") + print(f"Test Suites: {len(all_stats)}, total cases (sum of steps): {total}") + + +if __name__ == "__main__": + main() diff --git a/tools/nightly/send_nightly_email.py b/tools/nightly/send_nightly_email.py new file mode 100644 index 000000000..723a5c3d5 --- /dev/null +++ b/tools/nightly/send_nightly_email.py @@ -0,0 +1,351 @@ +#!/usr/bin/env python3 +""" +Send the nightly report by email to DAILY_EMAIL_LIST. + +Reads SMTP and recipient config from environment variables. Use --dry-run to +validate config and print subject/body without sending. +""" + +from __future__ import annotations + +import argparse +import logging +import mimetypes +import os +import smtplib +import sys +from datetime import date +from email.header import Header +from email.mime.application import MIMEApplication +from email.mime.multipart import MIMEMultipart +from email.mime.text import MIMEText + +LOGGER = logging.getLogger(__name__) + +# Env keys for SMTP and recipients. +ENV_SMTP_HOST = "SMTP_HOST" +ENV_SMTP_PORT = "SMTP_PORT" +ENV_SMTP_USERNAME = "SMTP_USERNAME" +ENV_SMTP_PASSWORD = "SMTP_PASSWORD" +ENV_DAILY_EMAIL_LIST = "DAILY_EMAIL_LIST" +ENV_EMAIL_SENDER = "EMAIL_SENDER" +ENV_EMAIL_SUBJECT_PREFIX = "EMAIL_SUBJECT_PREFIX" +ENV_BUILD_URL = "BUILDKITE_BUILD_URL" +ENV_COMMIT = "BUILDKITE_COMMIT" + +DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR") + +# Do not attach file if size >= this (bytes); body will suggest downloading from build URL. +MAX_ATTACHMENT_BYTES = 20 * 1024 * 1024 + +SMTP_RETRIES = 3 +SMTP_RETRY_DELAY_SEC = 5 + + +def _strip_quoted(s: str) -> str: + """Remove one layer of surrounding quote characters from env values. + + Handles both ASCII quotes ('", ''') and common Unicode quotes (β€œ ” β€˜ ’). + """ + if not s: + return s + s = s.strip() + if not s: + return s + quotes = {'"', "'", "β€œ", "”", "β€˜", "’"} + if len(s) >= 2 and s[0] in quotes and s[-1] in quotes: + s = s[1:-1].strip() + return s + + +def _get_required_env() -> dict[str, str]: + """Read required env vars; raise SystemExit with clear message if any missing.""" + required = { + ENV_SMTP_HOST: os.environ.get(ENV_SMTP_HOST), + ENV_SMTP_PORT: os.environ.get(ENV_SMTP_PORT), + ENV_SMTP_USERNAME: os.environ.get(ENV_SMTP_USERNAME), + ENV_SMTP_PASSWORD: os.environ.get(ENV_SMTP_PASSWORD), + ENV_DAILY_EMAIL_LIST: os.environ.get(ENV_DAILY_EMAIL_LIST), + } + missing = [k for k, v in required.items() if not (v and str(v).strip())] + if missing: + raise SystemExit(f"Missing required env vars: {', '.join(missing)}. Set them (e.g. in Buildkite secrets).") + return {k: _strip_quoted(str(v)) for k, v in required.items()} + + +def _get_latest_file(path: str) -> str | list[str]: + """ + Resolve report file patterns into concrete files. + + - If path is a directory: return all files in that directory (list). + - If path contains wildcards (e.g. *.xlsx, *.html): expand using glob and + return all matching files (list). + - If path is a plain filename or file path: treat as (dir, filename), + search that directory for files with the same basename and return the + latest modified one (str). + + Note: higher-level callers will turn a single str into [str] when they + want to attach multiple files. + """ + import glob + + # Support comma-/semicolon-separated list: e.g. "*.xlsx, *.html; reports/foo.xlsx" + parts = [p.strip() for p in path.split(";") for p in p.split(",") if p and p.strip()] + if len(parts) > 1: + files: list[str] = [] + for part in parts: + files_part = _get_latest_file(part) + if isinstance(files_part, str): + files.append(files_part) + else: + files.extend(files_part) + if not files: + raise SystemExit(f"No files resolved from list: {path}") + return sorted(set(files)) + + # Single expression after splitting + path = parts[0] if parts else path + + # Directory: return all files inside + if os.path.isdir(path): + files = [os.path.join(path, f) for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))] + if not files: + raise SystemExit(f"No files found in {path}") + return sorted(files) + + # Wildcard pattern: expand via glob + if any(ch in path for ch in ("*", "?", "[")) and not os.path.exists(path): + matches = [p for p in glob.glob(path) if os.path.isfile(p)] + if not matches: + raise SystemExit(f"No files match pattern {path!r}") + return sorted(matches) + + # Plain filename or explicit file path: search its directory for same basename + parent = os.path.dirname(path) or "." + base = os.path.basename(path) + if not base: + raise SystemExit(f"Invalid path (no filename): {path}") + if not os.path.isdir(parent): + raise SystemExit(f"Directory not found: {parent}") + candidates = [ + p + for p in (os.path.join(parent, f) for f in os.listdir(parent)) + if os.path.basename(p) == base and os.path.isfile(p) + ] + if not candidates: + raise SystemExit(f"No file named {base!r} found in {parent}") + return max(candidates, key=os.path.getmtime) + + +def _recipients_list(comma_separated: str) -> list[str]: + """Parse DAILY_EMAIL_LIST into a list of addresses.""" + return [_strip_quoted(a.strip()) for a in comma_separated.split(",") if a and a.strip()] + + +def _build_body( + date_str: str, + commit_sha: str | None, + build_url: str | None, + attachment_skipped: bool, +) -> str: + """Plain-text body with build metadata; note when attachment was skipped.""" + lines = [ + f"Nightly report date: {date_str}", + "", + f"Commit: {commit_sha or 'N/A'}", + f"Build: {build_url or 'N/A'}", + "", + ] + if attachment_skipped: + lines.append( + "Some report file(s) were too large to attach. Please download from the build artifacts (Build URL above)." + ) + lines.append("") + return "\n".join(lines) + + +def _build_subject(prefix: str | None, date_str: str) -> str: + """Subject line: optional prefix + date.""" + base = f"Nightly Report {date_str}" + if prefix and prefix.strip(): + return f"{prefix.strip()} {base}" + return base + + +def _mime_type_and_subtype(file_path: str) -> tuple[str, str]: + """Return (maintype, subtype) for the file; fallback to application/octet-stream.""" + full_type, _ = mimetypes.guess_type(file_path, strict=False) + if not full_type: + return "application", "octet-stream" + parts = full_type.split("/", 1) + if len(parts) != 2: + return "application", "octet-stream" + return parts[0], parts[1] + + +def _attach_file(msg: MIMEMultipart, file_path: str, max_bytes: int) -> bool: + """ + Attach one file to msg with correct MIME type. Return True if attached, False if skipped (e.g. size). + """ + if not os.path.isfile(file_path): + return False + size = os.path.getsize(file_path) + if size <= 0 or size >= max_bytes: + return False + maintype, subtype = _mime_type_and_subtype(file_path) + name = os.path.basename(file_path) + try: + if maintype == "text": + with open(file_path, encoding="utf-8", errors="replace") as f: + part = MIMEText(f.read(), subtype, "utf-8") + else: + with open(file_path, "rb") as f: + part = MIMEApplication(f.read(), _subtype=subtype) + part.add_header("Content-Disposition", "attachment", filename=name) + msg.attach(part) + return True + except OSError: + return False + + +def _send_mail( + report_files: list[str], + date_str: str, + dry_run: bool, +) -> None: + """Load config, build message (attach each file with correct MIME type), and send (or dry-run).""" + cfg = _get_required_env() + recipients = _recipients_list(cfg[ENV_DAILY_EMAIL_LIST]) + if not recipients: + raise SystemExit("DAILY_EMAIL_LIST is empty after parsing.") + + commit_sha = os.environ.get(ENV_COMMIT) + build_url = os.environ.get(ENV_BUILD_URL) + # Allow EMAIL_SENDER / EMAIL_SUBJECT_PREFIX to be quoted; strip quotes if present. + raw_sender = os.environ.get(ENV_EMAIL_SENDER) + sender = _strip_quoted(raw_sender) if raw_sender else cfg[ENV_SMTP_USERNAME] + raw_prefix = os.environ.get(ENV_EMAIL_SUBJECT_PREFIX) + prefix = _strip_quoted(raw_prefix) if raw_prefix else None + + attachment_skipped = any(os.path.isfile(p) and os.path.getsize(p) >= MAX_ATTACHMENT_BYTES for p in report_files) + body = _build_body( + commit_sha=commit_sha, build_url=build_url, date_str=date_str, attachment_skipped=attachment_skipped + ) + subject = _build_subject(prefix=prefix, date_str=date_str) + + msg = MIMEMultipart() + msg["Subject"] = Header(subject, "utf-8") + msg["From"] = sender + msg["To"] = ", ".join(recipients) + msg.attach(MIMEText(body, "plain", "utf-8")) + + attached_count = 0 + for file_path in report_files: + if os.path.isfile(file_path) and os.path.getsize(file_path) >= MAX_ATTACHMENT_BYTES: + attachment_skipped = True + else: + if _attach_file(msg, file_path, MAX_ATTACHMENT_BYTES): + attached_count += 1 + elif os.path.isfile(file_path): + attachment_skipped = True + + if dry_run: + LOGGER.info("dry-run: not sending mail") + print("To:", recipients, file=sys.stderr) + print("Subject:", subject, file=sys.stderr) + print("Attachments:", attached_count, "(skipped large:", attachment_skipped, ")", file=sys.stderr) + print("Body preview:", body[:300] + ("..." if len(body) > 300 else ""), file=sys.stderr) + return + + host = cfg[ENV_SMTP_HOST] + port = int(cfg[ENV_SMTP_PORT], 10) + use_ssl = port == 465 + last_err: Exception | None = None + for attempt in range(SMTP_RETRIES): + try: + if use_ssl: + with smtplib.SMTP_SSL(host, port=port, timeout=30) as smtp: + smtp.login(cfg[ENV_SMTP_USERNAME], cfg[ENV_SMTP_PASSWORD]) + smtp.sendmail(sender, recipients, msg.as_string()) + else: + with smtplib.SMTP(host, port=port, timeout=30) as smtp: + smtp.starttls() + smtp.login(cfg[ENV_SMTP_USERNAME], cfg[ENV_SMTP_PASSWORD]) + smtp.sendmail(sender, recipients, msg.as_string()) + LOGGER.info("sent nightly email to %d recipient(s)", len(recipients)) + return + except Exception as e: + last_err = e + LOGGER.warning("SMTP attempt %d/%d failed: %s", attempt + 1, SMTP_RETRIES, e) + if attempt < SMTP_RETRIES - 1: + import time + + time.sleep(SMTP_RETRY_DELAY_SEC) + hint = "" + if last_err and "getaddrinfo" in str(last_err).lower(): + hint = f" (check SMTP_HOST={host!r} is correct and reachable, e.g. nslookup {host})" + raise SystemExit(f"Failed to send email after {SMTP_RETRIES} attempts.{hint}") from last_err + + +def _vllm_omni_root() -> str: + """Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools').""" + path = os.path.dirname(os.path.abspath(__file__)) + while path and path != os.path.dirname(path): + if os.path.isdir(os.path.join(path, "tests")): + return path + path = os.path.dirname(path) + return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..")) + + +def _default_output_dir() -> str: + """Default: vllm-omni root / DEFAULT_OUTPUT_DIR (where nightly report files live).""" + root = _vllm_omni_root() + subdir = DEFAULT_OUTPUT_DIR if (DEFAULT_OUTPUT_DIR and DEFAULT_OUTPUT_DIR.strip()) else "." + return os.path.join(root, subdir) + + +def parse_args() -> argparse.Namespace: + """Parse command-line arguments.""" + parser = argparse.ArgumentParser( + description="Send nightly report by email (config from env).", + ) + parser.add_argument( + "--report-file", + type=str, + default=_default_output_dir(), + help="Folder/file path to the sent email file; default is DEFAULT_OUTPUT_DIR.", + ) + parser.add_argument( + "--date", + type=str, + default=None, + help="Date string for subject/body (default: inferred from report filename).", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print recipient, subject, and body; do not send.", + ) + return parser.parse_args() + + +def main() -> None: + """Entrypoint.""" + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", + ) + args = parse_args() + + result = _get_latest_file(args.report_file) + report_files = result if isinstance(result, list) else [result] + report_files = [p for p in report_files if os.path.isfile(p)] + if not report_files: + raise SystemExit("No report file(s) found.") + + date_str = args.date or date.today().strftime("%Y-%m-%d") + _send_mail(report_files=report_files, date_str=date_str, dry_run=args.dry_run) + + +if __name__ == "__main__": + main() diff --git a/tools/nightly/send_nightly_perf_email.py b/tools/nightly/send_nightly_perf_email.py deleted file mode 100644 index 9b4bf7ece..000000000 --- a/tools/nightly/send_nightly_perf_email.py +++ /dev/null @@ -1,234 +0,0 @@ -#!/usr/bin/env python3 -""" -Send the nightly performance Excel report by email to DAILY_EMAIL_LIST. - -Reads SMTP and recipient config from environment variables. Use --dry-run to -validate config and print subject/body without sending. -""" - -from __future__ import annotations - -import argparse -import logging -import os -import smtplib -import sys -from email.header import Header -from email.mime.application import MIMEApplication -from email.mime.multipart import MIMEMultipart -from email.mime.text import MIMEText - -LOGGER = logging.getLogger(__name__) - -# Env keys for SMTP and recipients. -ENV_SMTP_HOST = "SMTP_HOST" -ENV_SMTP_PORT = "SMTP_PORT" -ENV_SMTP_USERNAME = "SMTP_USERNAME" -ENV_SMTP_PASSWORD = "SMTP_PASSWORD" -ENV_DAILY_EMAIL_LIST = "DAILY_EMAIL_LIST" -ENV_EMAIL_SENDER = "EMAIL_SENDER" -ENV_EMAIL_SUBJECT_PREFIX = "EMAIL_SUBJECT_PREFIX" -ENV_BUILD_URL = "BUILDKITE_BUILD_URL" -ENV_COMMIT = "BUILDKITE_COMMIT" - -DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR") - -# Do not attach Excel if size >= this (bytes); body will suggest downloading from build URL. -MAX_ATTACHMENT_BYTES = 20 * 1024 * 1024 - -SMTP_RETRIES = 3 -SMTP_RETRY_DELAY_SEC = 5 - - -def _get_required_env() -> dict[str, str]: - """Read required env vars; raise SystemExit with clear message if any missing.""" - required = { - ENV_SMTP_HOST: os.environ.get(ENV_SMTP_HOST), - ENV_SMTP_PORT: os.environ.get(ENV_SMTP_PORT), - ENV_SMTP_USERNAME: os.environ.get(ENV_SMTP_USERNAME), - ENV_SMTP_PASSWORD: os.environ.get(ENV_SMTP_PASSWORD), - ENV_DAILY_EMAIL_LIST: os.environ.get(ENV_DAILY_EMAIL_LIST), - } - missing = [k for k, v in required.items() if not (v and str(v).strip())] - if missing: - raise SystemExit(f"Missing required env vars: {', '.join(missing)}. Set them (e.g. in Buildkite secrets).") - return {k: str(v).strip() for k, v in required.items()} - - -def _get_latest_file(folder_path: str) -> str: - """Get the latest modified file from the folder path.""" - files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(".xlsx")] - if not files: - raise SystemExit(f"No Excel files found in {folder_path}") - return max(files, key=os.path.getmtime) - - -def _recipients_list(comma_separated: str) -> list[str]: - """Parse DAILY_EMAIL_LIST into a list of addresses.""" - return [a.strip() for a in comma_separated.split(",") if a.strip()] - - -def _build_body( - date_str: str, - commit_sha: str | None, - build_url: str | None, - attachment_skipped: bool, -) -> str: - """Plain-text body with build metadata; note when attachment was skipped.""" - lines = [ - f"Nightly performance report date: {date_str}", - "", - f"Commit: {commit_sha or 'N/A'}", - f"Build: {build_url or 'N/A'}", - "", - ] - if attachment_skipped: - lines.append( - "Report file was too large to attach. Please download the Excel from the build artifacts (Build URL above)." - ) - lines.append("") - return "\n".join(lines) - - -def _build_subject(prefix: str | None, date_str: str) -> str: - """Subject line: optional prefix + date.""" - base = f"Nightly Perf {date_str}" - if prefix and prefix.strip(): - return f"{prefix.strip()} {base}" - return base - - -def _send_mail( - report_file: str, - date_str: str, - dry_run: bool, -) -> None: - """Load config, build message, and send (or dry-run).""" - cfg = _get_required_env() - recipients = _recipients_list(cfg[ENV_DAILY_EMAIL_LIST]) - if not recipients: - raise SystemExit("DAILY_EMAIL_LIST is empty after parsing.") - - commit_sha = os.environ.get(ENV_COMMIT) - build_url = os.environ.get(ENV_BUILD_URL) - sender = os.environ.get(ENV_EMAIL_SENDER) or cfg[ENV_SMTP_USERNAME] - prefix = os.environ.get(ENV_EMAIL_SUBJECT_PREFIX) - - size = os.path.getsize(report_file) if os.path.isfile(report_file) else 0 - attach_excel = size < MAX_ATTACHMENT_BYTES and size > 0 - attachment_skipped = os.path.isfile(report_file) and not attach_excel - - body = _build_body( - commit_sha=commit_sha, build_url=build_url, date_str=date_str, attachment_skipped=attachment_skipped - ) - subject = _build_subject(prefix=prefix, date_str=date_str) - - msg = MIMEMultipart() - msg["Subject"] = Header(subject, "utf-8") - msg["From"] = sender - msg["To"] = ", ".join(recipients) - msg.attach(MIMEText(body, "plain", "utf-8")) - - if attach_excel: - with open(report_file, "rb") as f: - part = MIMEApplication(f.read(), _subtype="vnd.openxmlformats-officedocument.spreadsheetml.sheet") - part.add_header("Content-Disposition", "attachment", filename=os.path.basename(report_file)) - msg.attach(part) - - if dry_run: - LOGGER.info("dry-run: not sending mail") - print("To:", recipients, file=sys.stderr) - print("Subject:", subject, file=sys.stderr) - print("Attachment:", "yes" if attach_excel else "no (size limit)", file=sys.stderr) - print("Body preview:", body[:300] + ("..." if len(body) > 300 else ""), file=sys.stderr) - return - - port = int(cfg[ENV_SMTP_PORT], 10) - last_err: Exception | None = None - for attempt in range(SMTP_RETRIES): - try: - with smtplib.SMTP(cfg[ENV_SMTP_HOST], port=port, timeout=30) as smtp: - smtp.starttls() - smtp.login(cfg[ENV_SMTP_USERNAME], cfg[ENV_SMTP_PASSWORD]) - smtp.sendmail(sender, recipients, msg.as_string()) - LOGGER.info("sent nightly perf email to %d recipient(s)", len(recipients)) - return - except Exception as e: - last_err = e - LOGGER.warning("SMTP attempt %d/%d failed: %s", attempt + 1, SMTP_RETRIES, e) - if attempt < SMTP_RETRIES - 1: - import time - - time.sleep(SMTP_RETRY_DELAY_SEC) - raise SystemExit(f"Failed to send email after {SMTP_RETRIES} attempts.") from last_err - - -def _date_from_filename(path: str) -> str: - """Try to derive a date string from report filename (e.g. nightly_perf_20260211-020704.xlsx -> 2026-02-11).""" - base = os.path.splitext(os.path.basename(path))[0] - if base.startswith("nightly_perf_") and len(base) > 13: - raw = base.replace("nightly_perf_", "")[:8] - if len(raw) == 8 and raw.isdigit(): - return f"{raw[:4]}-{raw[4:6]}-{raw[6:8]}" - return base or "unknown" - - -def _vllm_omni_root() -> str: - """Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools').""" - path = os.path.dirname(os.path.abspath(__file__)) - while path and path != os.path.dirname(path): - if os.path.isdir(os.path.join(path, "tests")): - return path - path = os.path.dirname(path) - return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..")) - - -def _default_output_dir() -> str: - """Default: vllm-omni root / DEFAULT_OUTPUT_DIR (where performance .xlsx files live).""" - root = _vllm_omni_root() - return os.path.join(root, DEFAULT_OUTPUT_DIR) - - -def parse_args() -> argparse.Namespace: - """Parse command-line arguments.""" - parser = argparse.ArgumentParser( - description="Send nightly performance Excel report by email (config from env).", - ) - parser.add_argument( - "--report-file", - type=str, - default=_default_output_dir(), - help="Folder/file path to the nightly_perf_*.xlsx file; default is DEFAULT_OUTPUT_DIR.", - ) - parser.add_argument( - "--date", - type=str, - default=None, - help="Date string for subject/body (default: inferred from report filename).", - ) - parser.add_argument( - "--dry-run", - action="store_true", - help="Print recipient, subject, and body; do not send.", - ) - return parser.parse_args() - - -def main() -> None: - """Entrypoint.""" - logging.basicConfig( - level=logging.INFO, - format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", - ) - args = parse_args() - - report_file = _get_latest_file(args.report_file) if os.path.isdir(args.report_file) else args.report_file - if not os.path.isfile(report_file): - raise SystemExit(f"Report file not found: {report_file}") - - date_str = args.date or _date_from_filename(report_file) - _send_mail(report_file=report_file, date_str=date_str, dry_run=args.dry_run) - - -if __name__ == "__main__": - main() diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 78792e81c..84bcfa93a 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -208,10 +208,11 @@ async def async_request_openai_chat_omni_completions( async def async_request_openai_audio_speech( request_func_input: RequestFuncInput, session: aiohttp.ClientSession, pbar: tqdm | None = None ) -> MixRequestFuncOutput: - """Non-streaming request to /v1/audio/speech endpoint. + """Streaming request to /v1/audio/speech endpoint. - The endpoint returns raw audio bytes (e.g. WAV). Pass voice, instructions, - and other TTS-specific fields via ``extra_body``. + Sends ``stream=true`` with ``response_format=pcm`` so the server returns + raw PCM chunks as they are decoded. This allows measuring TTFP (time to + first audio packet) separately from E2EL. """ api_url = request_func_input.api_url _validate_api_url(api_url, "OpenAI Audio Speech API", "audio/speech") @@ -219,16 +220,11 @@ async def async_request_openai_audio_speech( payload = { "model": request_func_input.model_name if request_func_input.model_name else request_func_input.model, "input": request_func_input.prompt, + "stream": True, + "response_format": "pcm", } _update_payload_common(payload, request_func_input) - response_format = payload.get("response_format", "wav") - if response_format == "pcm": - raise ValueError( - "pcm response format is not supported yet. \ - Please use other formats like wav, mp3, etc. instead." - ) - headers = { "Content-Type": "application/json", "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", @@ -238,41 +234,38 @@ async def async_request_openai_audio_speech( output = MixRequestFuncOutput() output.prompt_len = request_func_input.prompt_len + # PCM format: 16-bit signed, 24 kHz, mono + sample_rate = 24000 + sample_width = 2 # 16-bit = 2 bytes + channels = 1 + st = time.perf_counter() output.start_time = st + total_pcm_bytes = 0 try: async with session.post(url=api_url, json=payload, headers=headers) as response: if response.status == 200: - audio_bytes = await response.read() + async for chunk in response.content.iter_any(): + if not chunk: + continue + timestamp = time.perf_counter() + if output.audio_ttfp == 0.0: + output.audio_ttfp = timestamp - st + output.ttft = output.audio_ttfp + total_pcm_bytes += len(chunk) + end_time = time.perf_counter() output.latency = end_time - st - # ttft = latency since this is a non-streaming request - # hence there is no distinction between first and last token/audio - output.ttft = output.latency - output.audio_ttfp = output.latency - try: - audio_segment = AudioSegment.from_file(io.BytesIO(audio_bytes)) - output.audio_duration = len(audio_segment) / 1000.0 - frame_width = audio_segment.frame_width - if frame_width > 0: - output.audio_frames = len(audio_segment.raw_data) // frame_width - else: - output.audio_frames = 0 - logger.warning("Audio frame width is zero") - if output.audio_duration > 0: - # rtf = audio_generate_time / audio_duration and - # audio_generate_time = latency since this is a non-streaming request - # so the time to receive last portion of audio is the latency - output.audio_rtf = output.latency / output.audio_duration - else: - output.audio_rtf = 0 - logger.warning("Audio duration is zero") - output.success = True - except Exception as e: - output.success = False - output.error = f"Failed to parse audio response: {e}" - logger.error(f"ERROR: Failed to parse audio response: {e}") + total_samples = total_pcm_bytes // (sample_width * channels) + output.audio_duration = total_samples / sample_rate + output.audio_frames = total_samples + if output.audio_duration > 0: + output.audio_rtf = output.latency / output.audio_duration + else: + output.audio_rtf = 0 + logger.warning("Audio duration is zero") + output.success = True else: output.error = response.reason or "" output.success = False diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index e2db6f427..2aa236e69 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -4,8 +4,28 @@ Configuration module for vLLM-Omni. from vllm_omni.config.lora import LoRAConfig from vllm_omni.config.model import OmniModelConfig +from vllm_omni.config.stage_config import ( + ModelPipeline, + StageConfig, + StageConfigFactory, + StageType, +) +from vllm_omni.config.yaml_util import ( + create_config, + load_yaml_config, + merge_configs, + to_dict, +) __all__ = [ "OmniModelConfig", "LoRAConfig", + "StageConfig", + "StageConfigFactory", + "ModelPipeline", + "StageType", + "create_config", + "load_yaml_config", + "merge_configs", + "to_dict", ] diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index fa5483b31..f837b8d45 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -48,7 +48,7 @@ class OmniModelConfig(ModelConfig): stage_id: int = 0 async_chunk: bool = False model_stage: str = "thinker" - model_arch: str = "Qwen2_5OmniForConditionalGeneration" + model_arch: str | None = None worker_type: str | None = None engine_output_type: str | None = None hf_config_name: str | None = None @@ -69,7 +69,9 @@ class OmniModelConfig(ModelConfig): @property def architectures(self) -> list[str]: - return [self.model_arch] + if self.model_arch is not None: + return [self.model_arch] + return super().architectures @property def embedding_size(self): diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py new file mode 100644 index 000000000..ad298aa9d --- /dev/null +++ b/vllm_omni/config/stage_config.py @@ -0,0 +1,487 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Stage Configuration System for vLLM-Omni. + +Pipeline structure (stages, types, data-flow) is defined in per-model YAML +files and is set by model developers at integration time. +Runtime parameters (gpu_memory_utilization, tp_size, etc.) come from CLI. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict + +# Pipeline YAMLs live alongside model code in model_executor/models// +_MODELS_DIR = Path(__file__).resolve().parent.parent / "model_executor" / "models" + + +def get_pipeline_path(model_dir: str, filename: str) -> Path: + """Return the full path to a pipeline YAML file. + + Args: + model_dir: Model subdirectory name (e.g., "qwen3_omni"). + filename: Name of the YAML file (e.g., "pipeline.yaml"). + + Returns: + Absolute path to the file. + """ + return _MODELS_DIR / model_dir / filename + + +logger = init_logger(__name__) + + +class StageType(str, Enum): + """Type of processing stage in the Omni pipeline.""" + + LLM = "llm" + DIFFUSION = "diffusion" + + +@dataclass +class StageConfig: + """Per-stage configuration β€” pipeline-structure fields only. + + Engine params (gpu_memory_utilization, tp_size, etc.) come from CLI, + NOT from this class. + """ + + # Identity + stage_id: int + model_stage: str + + # Stage type + stage_type: StageType = StageType.LLM + + input_sources: list[int] = field(default_factory=list) + custom_process_input_func: str | None = None + final_output: bool = False + final_output_type: str | None = None # "text", "audio", "image" + worker_type: str | None = None # "ar" or "generation" + scheduler_cls: str | None = None + hf_config_name: str | None = None + is_comprehension: bool = False + + # Runtime overrides (populated from CLI, not from pipeline YAML) + runtime_overrides: dict[str, Any] = field(default_factory=dict) + + def to_omegaconf(self) -> Any: + """Convert to OmegaConf for backward compatibility with OmniStage. + + Returns: + OmegaConf DictConfig with stage configuration in legacy format. + """ + # Build engine_args dict with required fields + engine_args: dict[str, Any] = { + "model_stage": self.model_stage, + } + + if self.worker_type: + engine_args["worker_type"] = self.worker_type + if self.scheduler_cls: + engine_args["scheduler_cls"] = self.scheduler_cls + if self.hf_config_name: + engine_args["hf_config_name"] = self.hf_config_name + + # Apply runtime overrides (CLI args) + for key, value in self.runtime_overrides.items(): + if key not in ("devices", "max_batch_size"): + engine_args[key] = value + + # Build runtime config + runtime: dict[str, Any] = { + "process": True, + "max_batch_size": self.runtime_overrides.get("max_batch_size", 1), + } + if "devices" in self.runtime_overrides: + runtime["devices"] = self.runtime_overrides["devices"] + + # Build full config dict + config_dict: dict[str, Any] = { + "stage_id": self.stage_id, + "stage_type": StageType(self.stage_type).value, + "engine_args": create_config(engine_args), + "runtime": create_config(runtime), + "engine_input_source": self.input_sources, # Legacy field name + "final_output": self.final_output, + "final_output_type": self.final_output_type, + "is_comprehension": self.is_comprehension, + } + + if self.custom_process_input_func: + config_dict["custom_process_input_func"] = self.custom_process_input_func + + return create_config(config_dict) + + +@dataclass +class ModelPipeline: + """Complete pipeline definition for a multi-stage model. + + Defined by model developers, bundled with the model, not user-editable. + """ + + model_type: str + stages: list[StageConfig] + + # Pipeline-wide behavior flags + async_chunk: bool = False + + # Optional distributed configuration + connectors: dict[str, Any] | None = None + edges: list[dict[str, Any]] | None = None + + def get_stage(self, stage_id: int) -> StageConfig | None: + """Look up a stage by its ID. + + Args: + stage_id: The stage ID to search for. + + Returns: + The matching StageConfig, or None if not found. + """ + for stage in self.stages: + if stage.stage_id == stage_id: + return stage + return None + + def validate_pipeline(self) -> list[str]: + """Validate pipeline topology at model integration time (not runtime). + + Checks: + - All stage IDs are unique + - All input_sources reference valid stage IDs + - At least one entry point (stage with empty input_sources) + + Returns: + List of validation error messages. Empty list if valid. + """ + errors: list[str] = [] + + if not self.stages: + errors.append("Topology has no stages defined") + return errors + + # Check for unique stage IDs + stage_ids = [s.stage_id for s in self.stages] + if len(stage_ids) != len(set(stage_ids)): + errors.append("Duplicate stage IDs found") + + stage_id_set = set(stage_ids) + + # Check input_sources reference valid stages + for stage in self.stages: + for source_id in stage.input_sources: + if source_id not in stage_id_set: + errors.append(f"Stage {stage.stage_id} references non-existent input source {source_id}") + if source_id == stage.stage_id: + errors.append(f"Stage {stage.stage_id} references itself as input source") + + # Check for at least one entry point + entry_points = [s for s in self.stages if not s.input_sources] + if not entry_points: + errors.append("No entry point found (stage with empty input_sources)") + + return errors + + +class StageConfigFactory: + """Factory that loads pipeline YAML and merges CLI overrides. + + Handles both single-stage and multi-stage models. + """ + + # Mapping of model types to directories under model_executor/models/. + PIPELINE_MODELS: dict[str, str] = { + "qwen3_omni_moe": "qwen3_omni", + "qwen2_5_omni": "qwen2_5_omni", + "bagel": "bagel", + "qwen3_tts": "qwen3_tts", + "mimo_audio": "mimo_audio", + "glm-image": "glm_image", + "cosyvoice3": "cosyvoice3", + "mammothmoda2": "mammoth_moda2", + } + + # Fallback: map HF architecture class names to pipeline dirs. + # Used when model_type collides with another model (e.g. MiMo Audio + # reports model_type="qwen2" which matches plain Qwen2, not our pipeline). + _ARCHITECTURE_MODELS: dict[str, str] = { + "MiMoAudioForConditionalGeneration": "mimo_audio", + "HunyuanImage3ForCausalMM": "hunyuan_image3", + } + + @classmethod + def create_from_model( + cls, + model: str, + cli_overrides: dict[str, Any] | None = None, + ) -> list[StageConfig] | None: + """Load pipeline YAML, merge with CLI overrides. + + Args: + model: Model name or path. + cli_overrides: CLI overrides from VllmConfig/OmniDiffusionConfig. + + Returns: + List of StageConfig objects with CLI overrides applied, + or None if no pipeline definition was found for this model. + """ + if cli_overrides is None: + cli_overrides = {} + + trust_remote_code = cli_overrides.get("trust_remote_code", True) + pipeline = cls._load_pipeline(model, trust_remote_code=trust_remote_code) + + if pipeline is None: + return None + + errors = pipeline.validate_pipeline() + if errors: + logger.warning(f"Pipeline validation warnings for {model}: {errors}") + + # Apply CLI overrides + result: list[StageConfig] = [] + for stage in pipeline.stages: + # Merge global CLI overrides + stage.runtime_overrides = cls._merge_cli_overrides(stage, cli_overrides) + result.append(stage) + + return result + + @classmethod + def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]: + """Single-stage diffusion - no YAML needed. + + Creates a default diffusion stage configuration for single-stage + diffusion models. Returns a legacy OmegaConf-compatible dict for + backward compatibility with OmniStage. + + Args: + kwargs: Engine arguments from CLI/API. + + Returns: + List containing a single config dict for the diffusion stage. + """ + # Calculate devices based on parallel config + devices = "0" + if "parallel_config" in kwargs: + num_devices = kwargs["parallel_config"].world_size + for i in range(1, num_devices): + devices += f",{i}" + + # Collect engine args – skip non-serializable objects + engine_args: dict[str, Any] = {} + for key, value in kwargs.items(): + if key in ("parallel_config",): + continue + engine_args[key] = value + + engine_args.setdefault("cache_backend", "none") + engine_args["model_stage"] = "diffusion" + + # Convert dtype to string for OmegaConf + if "dtype" in engine_args: + engine_args["dtype"] = str(engine_args["dtype"]) + + config_dict: dict[str, Any] = { + "stage_id": 0, + "stage_type": StageType.DIFFUSION.value, + "runtime": { + "process": True, + "devices": devices, + "max_batch_size": 1, + }, + "engine_args": create_config(engine_args), + "final_output": True, + "final_output_type": "image", + } + + return [config_dict] + + @classmethod + def _load_pipeline(cls, model: str, trust_remote_code: bool = True) -> ModelPipeline | None: + """Load pipeline YAML for the model. + + Args: + model: Model name or path. + trust_remote_code: Whether to trust remote code for HF config loading. + + Returns: + ModelPipeline if found, None otherwise. + """ + model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code) + if model_type is None: + return None + + pipeline_dir = cls.PIPELINE_MODELS.get(model_type) + + # Fallback: check HF architectures when model_type doesn't match + if pipeline_dir is None and hf_config is not None: + for arch in getattr(hf_config, "architectures", []) or []: + pipeline_dir = cls._ARCHITECTURE_MODELS.get(arch) + if pipeline_dir is not None: + model_type = pipeline_dir + break + + if pipeline_dir is None: + logger.debug(f"No pipeline mapping for model_type: {model_type}") + return None + + pipeline_path = get_pipeline_path(pipeline_dir, "pipeline.yaml") + + if not pipeline_path.exists(): + logger.debug(f"Pipeline file not found: {pipeline_path}") + return None + + return cls._parse_pipeline_yaml(pipeline_path, model_type) + + @classmethod + def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline: + """Parse a pipeline YAML file. + + Args: + path: Path to the YAML file. + model_type: Model type identifier. + + Returns: + ModelPipeline object. + """ + config_data = load_yaml_config(path) + + stages: list[StageConfig] = [] + for stage_data in config_data.stages: + # Use .get() for optional fields β€” idiomatic for OmegaConf DictConfig + stage_type_str = stage_data.get("stage_type", "llm") + stage_type = StageType(stage_type_str) if stage_type_str else StageType.LLM + + # Handle both 'input_sources' (new) and 'engine_input_source' (legacy) + input_sources = stage_data.get("input_sources", None) + if input_sources is None: + input_sources = stage_data.get("engine_input_source", []) + if input_sources is None: + input_sources = [] + input_sources = list(input_sources) + + stage = StageConfig( + stage_id=stage_data.stage_id, + model_stage=stage_data.model_stage, + stage_type=stage_type, + input_sources=input_sources, + custom_process_input_func=stage_data.get("custom_process_input_func", None), + final_output=stage_data.get("final_output", False), + final_output_type=stage_data.get("final_output_type", None), + worker_type=stage_data.get("worker_type", None), + scheduler_cls=stage_data.get("scheduler_cls", None), + hf_config_name=stage_data.get("hf_config_name", None), + is_comprehension=stage_data.get("is_comprehension", False), + ) + stages.append(stage) + + # Get pipeline-wide flags + async_chunk = config_data.get("async_chunk", False) + + # Get optional connector config + connectors = to_dict(config_data.connectors) if hasattr(config_data, "connectors") else None + edges = to_dict(config_data.edges) if hasattr(config_data, "edges") else None + + return ModelPipeline( + model_type=getattr(config_data, "model_type", model_type), + stages=stages, + async_chunk=async_chunk, + connectors=connectors, + edges=edges, + ) + + @classmethod + def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]: + """Auto-detect model_type from model directory. + + Args: + model: Model name or path. + trust_remote_code: Whether to trust remote code for HF config loading. + + Returns: + Tuple of (model_type, hf_config). Both may be None on failure. + """ + try: + from vllm.transformers_utils.config import get_config + + hf_config = get_config(model, trust_remote_code=trust_remote_code) + return hf_config.model_type, hf_config + except Exception as e: + logger.debug(f"Failed to auto-detect model type for {model}: {e}") + return None, None + + # Keys that should never be forwarded as engine overrides (internal / + # orchestrator-only knobs, complex objects, etc.). + _INTERNAL_KEYS: set[str] = { + "model", + "stage_configs_path", + "stage_id", + "stage_init_timeout", + "init_timeout", + "shm_threshold_bytes", + "worker_backend", + "ray_address", + "batch_timeout", + "log_stats", + "tokenizer", + "parallel_config", + } + + @classmethod + def _merge_cli_overrides( + cls, + stage: StageConfig, + cli_overrides: dict[str, Any], + ) -> dict[str, Any]: + """Merge CLI overrides into stage runtime config. + + All CLI arguments registered by engine config classes (e.g. + EngineArgs / OmniDiffusionConfig) are accepted as overrides + unless they appear in ``_INTERNAL_KEYS``. + + Handles: + - Global overrides (apply to all stages) + - Per-stage overrides (--stage-N-* format, take precedence) + + Args: + stage: The stage to merge overrides into. + cli_overrides: CLI arguments from VllmConfig/OmniDiffusionConfig. + + Returns: + Dict of runtime overrides for this stage. + """ + result: dict[str, Any] = {} + + # Apply global overrides – any key not in the internal blocklist + # is forwarded so that engine-registered params work out of the box. + for key, value in cli_overrides.items(): + if key in cls._INTERNAL_KEYS: + continue + if re.match(r"stage_\d+_", key): + # Per-stage keys handled below + continue + if value is not None: + result[key] = value + + # Apply per-stage overrides (--stage-N-* format, take precedence) + stage_prefix = f"stage_{stage.stage_id}_" + for key, value in cli_overrides.items(): + if key.startswith(stage_prefix) and value is not None: + param_name = key[len(stage_prefix) :] + if param_name in cls._INTERNAL_KEYS: + continue + result[param_name] = value + + return result diff --git a/vllm_omni/config/yaml_util.py b/vllm_omni/config/yaml_util.py new file mode 100644 index 000000000..09b27ca02 --- /dev/null +++ b/vllm_omni/config/yaml_util.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Centralized OmegaConf wrapper for vLLM-Omni. + +All OmegaConf usage in the project MUST go through this module. +Other modules should import these wrapper functions instead of +using OmegaConf directly. +""" + +from __future__ import annotations + +from typing import Any + +from omegaconf import DictConfig, OmegaConf + + +def load_yaml_config(path: str | Any) -> DictConfig: + """Load a YAML file and return it as a DictConfig. + + Args: + path: Path to the YAML file. + + Returns: + OmegaConf DictConfig with attribute-style access. + """ + return OmegaConf.load(path) # type: ignore[return-value] + + +def create_config(data: Any) -> DictConfig: + """Wrap a dict (or list) into a DictConfig. + + Args: + data: Dict, list, or other structure to wrap. + + Returns: + OmegaConf DictConfig / ListConfig. + """ + return OmegaConf.create(data) + + +def merge_configs(*cfgs: Any) -> dict: + """Deep-merge multiple configs and return a plain dict. + + Args: + *cfgs: DictConfig or dict objects to merge (left to right). + + Returns: + Plain dict with merged, resolved values. + """ + merged = OmegaConf.merge(*cfgs) + return OmegaConf.to_container(merged, resolve=True) # type: ignore[return-value] + + +def to_dict(obj: Any, *, resolve: bool = True) -> Any: + """Convert a DictConfig (or similar) to a plain dict. + + Args: + obj: OmegaConf container to convert. + resolve: Whether to resolve interpolations (default True). + + Returns: + Plain dict. + """ + return OmegaConf.to_container(obj, resolve=resolve) # type: ignore[return-value] diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index 1c98d0724..d7165eda2 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -984,6 +984,80 @@ def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int], return refresh_cache_context +def enable_cache_for_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Enable cache-dit for Flux.2-dev pipeline. + + Args: + pipeline: The Flux2 pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + Returns: + A refresh function that can be called with a new ``num_inference_steps`` + to update the cache context for the pipeline. + """ + # Build DBCacheConfig for transformer + db_cache_config = _build_db_cache_config(cache_config) + + calibrator = None + if cache_config.enable_taylorseer: + taylorseer_order = cache_config.taylorseer_order + calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) + logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + + # Build ParamsModifier for transformer + modifier = ParamsModifier( + cache_config=db_cache_config, + calibrator_config=calibrator, + ) + + logger.info( + f"Enabling cache-dit on Flux transformer with BlockAdapter: " + f"Fn={db_cache_config.Fn_compute_blocks}, " + f"Bn={db_cache_config.Bn_compute_blocks}, " + f"W={db_cache_config.max_warmup_steps}, " + ) + + # Enable cache-dit using BlockAdapter for transformer + cache_dit.enable_cache( + ( + BlockAdapter( + transformer=pipeline.transformer, + blocks=[ + pipeline.transformer.transformer_blocks, + pipeline.transformer.single_transformer_blocks, + ], + forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2], + params_modifiers=[modifier], + ) + ), + cache_config=db_cache_config, + ) + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh cache context for the transformer with new num_inference_steps. + + Args: + pipeline: The Flux2 pipeline instance. + num_inference_steps: New number of inference steps. + """ + if cache_config.scm_steps_mask_policy is None: + cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) + else: + cache_dit.refresh_context( + pipeline.transformer, + cache_config=DBCacheConfig().reset( + num_inference_steps=num_inference_steps, + steps_computation_mask=cache_dit.steps_mask( + mask_policy=cache_config.scm_steps_mask_policy, + total_steps=num_inference_steps, + ), + steps_computation_policy=cache_config.scm_steps_policy, + ), + verbose=verbose, + ) + + return refresh_cache_context + + # Register custom cache-dit enablers after function definitions CUSTOM_DIT_ENABLERS.update( { @@ -998,6 +1072,7 @@ CUSTOM_DIT_ENABLERS.update( "LTX2Pipeline": enable_cache_for_ltx2, "LTX2ImageToVideoPipeline": enable_cache_for_ltx2, "BagelPipeline": enable_cache_for_bagel, + "Flux2Pipeline": enable_cache_for_flux2, } ) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 38366469e..634cb3294 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -463,11 +463,13 @@ class OmniDiffusionConfig: @property def is_moe(self) -> bool: num_experts = self.tf_model_config.get("num_experts", None) + if not isinstance(num_experts, (list, tuple, int)): + return False if isinstance(num_experts, int): - return num_experts > 1 + return num_experts > 0 if isinstance(num_experts, (list, tuple)): - return any(isinstance(n, int) and n > 1 for n in num_experts) + return any(isinstance(n, int) and n > 0 for n in num_experts) return False diff --git a/vllm_omni/diffusion/models/flux2/__init__.py b/vllm_omni/diffusion/models/flux2/__init__.py new file mode 100644 index 000000000..bce893e69 --- /dev/null +++ b/vllm_omni/diffusion/models/flux2/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Flux2 diffusion model components.""" + +from vllm_omni.diffusion.models.flux2.flux2_transformer import ( + Flux2Transformer2DModel, +) +from vllm_omni.diffusion.models.flux2.pipeline_flux2 import ( + Flux2Pipeline, + get_flux2_post_process_func, +) + +__all__ = [ + "Flux2Pipeline", + "Flux2Transformer2DModel", + "get_flux2_post_process_func", +] diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py new file mode 100644 index 000000000..040f2779a --- /dev/null +++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py @@ -0,0 +1,770 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable +from types import SimpleNamespace +from typing import Any + +import torch +from diffusers.models.embeddings import ( + TimestepEmbedding, + Timesteps, + get_1d_rotary_pos_embed, +) +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.normalization import AdaLayerNormContinuous +from diffusers.utils import is_torch_npu_available +from torch import nn +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.layers.rope import RotaryEmbedding + + +class Flux2SwiGLU(nn.Module): + """SwiGLU activation used by Flux2.""" + + def __init__(self): + super().__init__() + self.gate_fn = nn.SiLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return self.gate_fn(x1) * x2 + + +class Flux2FeedForward(nn.Module): + def __init__( + self, + dim: int, + dim_out: int | None = None, + mult: float = 3.0, + inner_dim: int | None = None, + bias: bool = False, + ): + super().__init__() + if inner_dim is None: + inner_dim = int(dim * mult) + dim_out = dim_out or dim + + self.linear_in = MergedColumnParallelLinear( + dim, + [inner_dim, inner_dim], + bias=bias, + return_bias=False, + ) + self.act_fn = Flux2SwiGLU() + self.linear_out = RowParallelLinear( + inner_dim, + dim_out, + bias=bias, + input_is_parallel=True, + return_bias=False, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.linear_in(x) + x = self.act_fn(x) + return self.linear_out(x) + + +class Flux2Attention(nn.Module): + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + added_kv_proj_dim: int | None = None, + added_proj_bias: bool | None = True, + out_bias: bool = True, + eps: float = 1e-5, + out_dim: int = None, + elementwise_affine: bool = True, + ): + super().__init__() + self.head_dim = dim_head + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.query_dim = query_dim + self.out_dim = out_dim if out_dim is not None else query_dim + self.heads = out_dim // dim_head if out_dim is not None else heads + self.dropout = dropout + self.added_kv_proj_dim = added_kv_proj_dim + + self.to_qkv = QKVParallelLinear( + hidden_size=query_dim, + head_size=self.head_dim, + total_num_heads=self.heads, + bias=bias, + ) + self.query_num_heads = self.to_qkv.num_heads + self.kv_num_heads = self.to_qkv.num_kv_heads + + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + + self.to_out = nn.ModuleList( + [ + RowParallelLinear( + self.inner_dim, + self.out_dim, + bias=out_bias, + input_is_parallel=True, + return_bias=False, + ), + nn.Dropout(dropout), + ] + ) + + if added_kv_proj_dim is not None: + self.norm_added_q = RMSNorm(dim_head, eps=eps) + self.norm_added_k = RMSNorm(dim_head, eps=eps) + self.add_kv_proj = QKVParallelLinear( + hidden_size=added_kv_proj_dim, + head_size=self.head_dim, + total_num_heads=self.heads, + bias=added_proj_bias, + ) + self.add_query_num_heads = self.add_kv_proj.num_heads + self.add_kv_num_heads = self.add_kv_proj.num_kv_heads + self.to_add_out = RowParallelLinear( + self.inner_dim, + query_dim, + bias=out_bias, + input_is_parallel=True, + return_bias=False, + ) + + self.rope = RotaryEmbedding(is_neox_style=False) + self.attn = Attention( + num_heads=self.query_num_heads, + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + num_kv_heads=self.kv_num_heads, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + qkv, _ = self.to_qkv(hidden_states) + query, key, value = qkv.chunk(3, dim=-1) + + encoder_query = encoder_key = encoder_value = None + if encoder_hidden_states is not None and self.added_kv_proj_dim is not None: + encoder_qkv, _ = self.add_kv_proj(encoder_hidden_states) + encoder_query, encoder_key, encoder_value = encoder_qkv.chunk(3, dim=-1) + + query = query.unflatten(-1, (self.query_num_heads, -1)) + key = key.unflatten(-1, (self.kv_num_heads, -1)) + value = value.unflatten(-1, (self.kv_num_heads, -1)) + + query = self.norm_q(query) + key = self.norm_k(key) + + if encoder_hidden_states is not None and self.added_kv_proj_dim is not None: + encoder_query = encoder_query.unflatten(-1, (self.add_query_num_heads, -1)) + encoder_key = encoder_key.unflatten(-1, (self.add_kv_num_heads, -1)) + encoder_value = encoder_value.unflatten(-1, (self.add_kv_num_heads, -1)) + + encoder_query = self.norm_added_q(encoder_query) + encoder_key = self.norm_added_k(encoder_key) + + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + if image_rotary_emb is not None: + cos, sin = image_rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + query = self.rope(query, cos, sin) + key = self.rope(key, cos, sin) + + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + hidden_states = self.attn(query, key, value, attn_metadata) + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + + if encoder_hidden_states is not None: + context_len = encoder_hidden_states.shape[1] + encoder_hidden_states, hidden_states = hidden_states.split_with_sizes( + [context_len, hidden_states.shape[1] - context_len], + dim=1, + ) + encoder_hidden_states = self.to_add_out(encoder_hidden_states) + + hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[1](hidden_states) + + if encoder_hidden_states is not None: + return hidden_states, encoder_hidden_states + return hidden_states + + +class Flux2ParallelSelfAttention(nn.Module): + """ + Parallel attention block that fuses QKV projections with MLP input projections. + """ + + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + bias: bool = False, + out_bias: bool = True, + eps: float = 1e-5, + out_dim: int = None, + elementwise_affine: bool = True, + mlp_ratio: float = 4.0, + mlp_mult_factor: int = 2, + ): + super().__init__() + self.head_dim = dim_head + self.inner_dim = out_dim if out_dim is not None else dim_head * heads + self.query_dim = query_dim + self.out_dim = out_dim if out_dim is not None else query_dim + self.heads = out_dim // dim_head if out_dim is not None else heads + self.dropout = dropout + + self.mlp_ratio = mlp_ratio + self.mlp_hidden_dim = int(query_dim * self.mlp_ratio) + self.mlp_mult_factor = mlp_mult_factor + + self.to_qkv_mlp_proj = ColumnParallelLinear( + self.query_dim, + self.inner_dim * 3 + self.mlp_hidden_dim * self.mlp_mult_factor, + bias=bias, + gather_output=True, + ) + self.mlp_act_fn = Flux2SwiGLU() + + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + + self.to_out = ColumnParallelLinear( + self.inner_dim + self.mlp_hidden_dim, + self.out_dim, + bias=out_bias, + gather_output=True, + ) + self.rope = RotaryEmbedding(is_neox_style=False) + self.attn = Attention( + num_heads=self.heads, + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + hidden_states, _ = self.to_qkv_mlp_proj(hidden_states) + qkv, mlp_hidden_states = torch.split( + hidden_states, + [3 * self.inner_dim, self.mlp_hidden_dim * self.mlp_mult_factor], + dim=-1, + ) + + query, key, value = qkv.chunk(3, dim=-1) + query = query.unflatten(-1, (self.heads, -1)) + key = key.unflatten(-1, (self.heads, -1)) + value = value.unflatten(-1, (self.heads, -1)) + + query = self.norm_q(query) + key = self.norm_k(key) + + if image_rotary_emb is not None: + cos, sin = image_rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + query = self.rope(query, cos, sin) + key = self.rope(key, cos, sin) + + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + attn_output = self.attn(query, key, value, attn_metadata) + attn_output = attn_output.flatten(2, 3).to(query.dtype) + + mlp_hidden_states = self.mlp_act_fn(mlp_hidden_states) + hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=-1) + hidden_states, _ = self.to_out(hidden_states) + return hidden_states + + +class Flux2SingleTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 3.0, + eps: float = 1e-6, + bias: bool = False, + ): + super().__init__() + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.attn = Flux2ParallelSelfAttention( + query_dim=dim, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=dim, + bias=bias, + out_bias=bias, + eps=eps, + mlp_ratio=mlp_ratio, + mlp_mult_factor=2, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, + temb_mod_params: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + split_hidden_states: bool = False, + text_seq_len: int | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if encoder_hidden_states is not None: + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + mod_shift, mod_scale, mod_gate = temb_mod_params + + norm_hidden_states = self.norm(hidden_states) + norm_hidden_states = (1 + mod_scale) * norm_hidden_states + mod_shift + + joint_attention_kwargs = joint_attention_kwargs or {} + attn_output = self.attn( + hidden_states=norm_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + hidden_states = hidden_states + mod_gate * attn_output + if hidden_states.dtype == torch.float16: + hidden_states = hidden_states.clip(-65504, 65504) + + if split_hidden_states: + encoder_hidden_states, hidden_states = hidden_states[:, :text_seq_len], hidden_states[:, text_seq_len:] + return encoder_hidden_states, hidden_states + return hidden_states + + +class Flux2TransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 3.0, + eps: float = 1e-6, + bias: bool = False, + ): + super().__init__() + self.mlp_hidden_dim = int(dim * mlp_ratio) + + self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + + self.attn = Flux2Attention( + query_dim=dim, + added_kv_proj_dim=dim, + dim_head=attention_head_dim, + heads=num_attention_heads, + out_dim=dim, + bias=bias, + added_proj_bias=bias, + out_bias=bias, + eps=eps, + ) + + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ff = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias) + + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ff_context = Flux2FeedForward(dim=dim, dim_out=dim, mult=mlp_ratio, bias=bias) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb_mod_params_img: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...], + temb_mod_params_txt: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...], + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + joint_attention_kwargs = joint_attention_kwargs or {} + + (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = temb_mod_params_img + (c_shift_msa, c_scale_msa, c_gate_msa), (c_shift_mlp, c_scale_mlp, c_gate_mlp) = temb_mod_params_txt + + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = (1 + scale_msa) * norm_hidden_states + shift_msa + + norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states) + norm_encoder_hidden_states = (1 + c_scale_msa) * norm_encoder_hidden_states + c_shift_msa + + attn_output, context_attn_output = self.attn( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + **joint_attention_kwargs, + ) + + attn_output = gate_msa * attn_output + hidden_states = hidden_states + attn_output + + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + ff_output = self.ff(norm_hidden_states) + hidden_states = hidden_states + gate_mlp * ff_output + + context_attn_output = c_gate_msa * context_attn_output + encoder_hidden_states = encoder_hidden_states + context_attn_output + + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) + norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp + context_ff_output = self.ff_context(norm_encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states + c_gate_mlp * context_ff_output + if encoder_hidden_states.dtype == torch.float16: + encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504) + + return encoder_hidden_states, hidden_states + + +class Flux2PosEmbed(nn.Module): + def __init__(self, theta: int, axes_dim: list[int]): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # Expected ids shape: [S, len(self.axes_dim)] + cos_out = [] + sin_out = [] + pos = ids.float() + is_mps = ids.device.type == "mps" + is_npu = ids.device.type == "npu" + freqs_dtype = torch.float32 if (is_mps or is_npu) else torch.float64 + # Unlike Flux 1, loop over len(self.axes_dim) rather than ids.shape[-1] + for i in range(len(self.axes_dim)): + freqs_cis = get_1d_rotary_pos_embed( + self.axes_dim[i], + pos[..., i], + theta=self.theta, + use_real=False, + freqs_dtype=freqs_dtype, + ) + cos_out.append(freqs_cis.real) + sin_out.append(freqs_cis.imag) + freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) + freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) + return freqs_cos, freqs_sin + + +class Flux2TimestepGuidanceEmbeddings(nn.Module): + def __init__( + self, + in_channels: int = 256, + embedding_dim: int = 6144, + bias: bool = False, + guidance_embeds: bool = True, + ): + super().__init__() + self.time_proj = Timesteps(num_channels=in_channels, flip_sin_to_cos=True, downscale_freq_shift=0) + self.timestep_embedder = TimestepEmbedding( + in_channels=in_channels, + time_embed_dim=embedding_dim, + sample_proj_bias=bias, + ) + + if guidance_embeds: + self.guidance_embedder = TimestepEmbedding( + in_channels=in_channels, + time_embed_dim=embedding_dim, + sample_proj_bias=bias, + ) + else: + self.guidance_embedder = None + + def forward(self, timestep: torch.Tensor, guidance: torch.Tensor | None) -> torch.Tensor: + timesteps_proj = self.time_proj(timestep) + timesteps_emb = self.timestep_embedder(timesteps_proj.to(timestep.dtype)) + + if guidance is not None and self.guidance_embedder is not None: + guidance_proj = self.time_proj(guidance) + guidance_emb = self.guidance_embedder(guidance_proj.to(guidance.dtype)) + return timesteps_emb + guidance_emb + return timesteps_emb + + +class Flux2Modulation(nn.Module): + def __init__(self, dim: int, mod_param_sets: int = 2, bias: bool = False): + super().__init__() + self.mod_param_sets = mod_param_sets + + self.linear = nn.Linear(dim, dim * 3 * self.mod_param_sets, bias=bias) + self.act_fn = nn.SiLU() + + def forward(self, temb: torch.Tensor) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]: + mod = self.act_fn(temb) + mod = self.linear(mod) + + if mod.ndim == 2: + mod = mod.unsqueeze(1) + mod_params = torch.chunk(mod, 3 * self.mod_param_sets, dim=-1) + # Return tuple of 3-tuples of modulation params shift/scale/gate + return tuple(mod_params[3 * i : 3 * (i + 1)] for i in range(self.mod_param_sets)) + + +class Flux2Transformer2DModel(nn.Module): + """ + The Transformer model introduced in Flux 2. + """ + + _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"] + + def __init__( + self, + patch_size: int = 1, + in_channels: int = 128, + out_channels: int | None = None, + num_layers: int = 8, + num_single_layers: int = 48, + attention_head_dim: int = 128, + num_attention_heads: int = 48, + joint_attention_dim: int = 15360, + timestep_guidance_channels: int = 256, + mlp_ratio: float = 3.0, + axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32), + rope_theta: int = 2000, + eps: float = 1e-6, + guidance_embeds: bool = True, + ): + super().__init__() + self.stacked_params_mapping = None + self.out_channels = out_channels or in_channels + self.inner_dim = num_attention_heads * attention_head_dim + self.config = SimpleNamespace( + patch_size=patch_size, + in_channels=in_channels, + out_channels=self.out_channels, + num_layers=num_layers, + num_single_layers=num_single_layers, + attention_head_dim=attention_head_dim, + num_attention_heads=num_attention_heads, + joint_attention_dim=joint_attention_dim, + timestep_guidance_channels=timestep_guidance_channels, + mlp_ratio=mlp_ratio, + axes_dims_rope=axes_dims_rope, + rope_theta=rope_theta, + eps=eps, + guidance_embeds=guidance_embeds, + ) + + self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope) + self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings( + in_channels=timestep_guidance_channels, + embedding_dim=self.inner_dim, + bias=False, + guidance_embeds=guidance_embeds, + ) + + self.double_stream_modulation_img = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False) + self.double_stream_modulation_txt = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False) + self.single_stream_modulation = Flux2Modulation(self.inner_dim, mod_param_sets=1, bias=False) + + self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False) + self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False) + + self.transformer_blocks = nn.ModuleList( + [ + Flux2TransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + bias=False, + ) + for _ in range(num_layers) + ] + ) + + self.single_transformer_blocks = nn.ModuleList( + [ + Flux2SingleTransformerBlock( + dim=self.inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + mlp_ratio=mlp_ratio, + eps=eps, + bias=False, + ) + for _ in range(num_single_layers) + ] + ) + + self.norm_out = AdaLayerNormContinuous( + self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False + ) + self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=False) + + @property + def dtype(self) -> torch.dtype: + return next(self.parameters()).dtype + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + timestep: torch.LongTensor = None, + img_ids: torch.Tensor = None, + txt_ids: torch.Tensor = None, + guidance: torch.Tensor | None = None, + joint_attention_kwargs: dict[str, Any] | None = None, + return_dict: bool = True, + ) -> torch.Tensor | Transformer2DModelOutput: + joint_attention_kwargs = joint_attention_kwargs or {} + + num_txt_tokens = encoder_hidden_states.shape[1] + + timestep = timestep.to(hidden_states.dtype) * 1000 + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) * 1000 + + temb = self.time_guidance_embed(timestep, guidance) + + double_stream_mod_img = self.double_stream_modulation_img(temb) + double_stream_mod_txt = self.double_stream_modulation_txt(temb) + single_stream_mod = self.single_stream_modulation(temb)[0] + + hidden_states = self.x_embedder(hidden_states) + encoder_hidden_states = self.context_embedder(encoder_hidden_states) + + if img_ids.ndim == 3: + img_ids = img_ids[0] + if txt_ids.ndim == 3: + txt_ids = txt_ids[0] + + if is_torch_npu_available(): + freqs_cos_image, freqs_sin_image = self.pos_embed(img_ids.cpu()) + image_rotary_emb = (freqs_cos_image.npu(), freqs_sin_image.npu()) + freqs_cos_text, freqs_sin_text = self.pos_embed(txt_ids.cpu()) + text_rotary_emb = (freqs_cos_text.npu(), freqs_sin_text.npu()) + else: + image_rotary_emb = self.pos_embed(img_ids) + text_rotary_emb = self.pos_embed(txt_ids) + concat_rotary_emb = ( + torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0), + torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0), + ) + + for index_block, block in enumerate(self.transformer_blocks): + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb_mod_params_img=double_stream_mod_img, + temb_mod_params_txt=double_stream_mod_txt, + image_rotary_emb=concat_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + for index_block, block in enumerate(self.single_transformer_blocks): + hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=None, + temb_mod_params=single_stream_mod, + image_rotary_emb=concat_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + hidden_states = hidden_states[:, num_txt_tokens:, ...] + hidden_states = self.norm_out(hidden_states, temb) + output = self.proj_out(hidden_states) + + if not return_dict: + return (output,) + + return Transformer2DModelOutput(sample=output) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + (".to_qkv.", ".to_q.", "q"), + (".to_qkv.", ".to_k.", "k"), + (".to_qkv.", ".to_v.", "v"), + (".add_kv_proj", ".add_q_proj", "q"), + (".add_kv_proj", ".add_k_proj", "k"), + (".add_kv_proj", ".add_v_proj", "v"), + ] + # Expose packed shard mappings for LoRA handling of fused projections. + self.stacked_params_mapping = stacked_params_mapping + + params_dict = dict(self.named_parameters()) + + for name, buffer in self.named_buffers(): + if name.endswith(".beta") or name.endswith(".eps"): + params_dict[name] = buffer + + loaded_params: set[str] = set() + for name, loaded_weight in weights: + original_name = name + mapped = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in original_name: + continue + name = original_name.replace(weight_name, param_name) + param = params_dict.get(name) + if param is None: + break + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + mapped = True + break + if mapped: + continue + + name = original_name + if name not in params_dict and ".to_out.0." in name: + name = name.replace(".to_out.0.", ".to_out.") + # Some GGUF checkpoints include quantized tensors for modules that + # are intentionally left unquantized in this model. + param = params_dict.get(name) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py new file mode 100644 index 000000000..7f07adf6f --- /dev/null +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -0,0 +1,1088 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import inspect +import json +import logging +import math +import os +from collections.abc import Callable, Iterable +from typing import Any, cast + +import numpy as np +import PIL.Image +import torch +from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.flux2.pipeline_flux2 import UPSAMPLING_MAX_IMAGE_SIZE +from diffusers.pipelines.flux2.system_messages import ( + SYSTEM_MESSAGE, + SYSTEM_MESSAGE_UPSAMPLING_I2I, + SYSTEM_MESSAGE_UPSAMPLING_T2I, +) +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import AutoProcessor, Mistral3ForConditionalGeneration, PixtralProcessor +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.flux2 import Flux2Transformer2DModel +from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific + +logger = logging.getLogger(__name__) + + +class Flux2ImageProcessor(VaeImageProcessor): + """Image processor to preprocess the reference image for Flux2.""" + + def __init__( + self, + do_resize: bool = True, + vae_scale_factor: int = 16, + vae_latent_channels: int = 32, + do_normalize: bool = True, + do_convert_rgb: bool = True, + ): + super().__init__( + do_resize=do_resize, + vae_scale_factor=vae_scale_factor, + vae_latent_channels=vae_latent_channels, + do_normalize=do_normalize, + do_convert_rgb=do_convert_rgb, + ) + + @staticmethod + def check_image_input( + image: PIL.Image.Image, + max_aspect_ratio: int = 8, + min_side_length: int = 64, + max_area: int = 1024 * 1024, + ) -> PIL.Image.Image: + if not isinstance(image, PIL.Image.Image): + raise ValueError(f"Image must be a PIL.Image.Image, got {type(image)}") + + width, height = image.size + if width < min_side_length or height < min_side_length: + raise ValueError(f"Image too small: {width}x{height}. Both dimensions must be at least {min_side_length}px") + + aspect_ratio = max(width / height, height / width) + if aspect_ratio > max_aspect_ratio: + raise ValueError( + f"Aspect ratio too extreme: {width}x{height} (ratio: {aspect_ratio:.1f}:1). " + f"Maximum allowed ratio is {max_aspect_ratio}:1" + ) + + if width * height > max_area: + logger.warning("Image area exceeds recommended maximum; resizing will be applied.") + + return image + + @staticmethod + def _resize_to_target_area(image: PIL.Image.Image, target_area: int = 1024 * 1024) -> PIL.Image.Image: + image_width, image_height = image.size + scale = math.sqrt(target_area / (image_width * image_height)) + width = int(image_width * scale) + height = int(image_height * scale) + return image.resize((width, height), PIL.Image.Resampling.LANCZOS) + + @staticmethod + def _resize_if_exceeds_area(image: PIL.Image.Image, target_area: int = 1024 * 1024) -> PIL.Image.Image: + image_width, image_height = image.size + if image_width * image_height <= target_area: + return image + return Flux2ImageProcessor._resize_to_target_area(image, target_area) + + def _resize_and_crop(self, image: PIL.Image.Image, width: int, height: int) -> PIL.Image.Image: + image_width, image_height = image.size + left = (image_width - width) // 2 + top = (image_height - height) // 2 + right = left + width + bottom = top + height + return image.crop((left, top, right, bottom)) + + @staticmethod + def concatenate_images(images: list[PIL.Image.Image]) -> PIL.Image.Image: + if len(images) == 1: + return images[0].copy() + + images = [img.convert("RGB") if img.mode != "RGB" else img for img in images] + total_width = sum(img.width for img in images) + max_height = max(img.height for img in images) + background_color = (255, 255, 255) + new_img = PIL.Image.new("RGB", (total_width, max_height), background_color) + + x_offset = 0 + for img in images: + y_offset = (max_height - img.height) // 2 + new_img.paste(img, (x_offset, y_offset)) + x_offset += img.width + + return new_img + + +def get_flux2_post_process_func( + od_config: OmniDiffusionConfig, +): + if od_config.output_type == "latent": + return lambda x: x + model_name = od_config.model + if os.path.exists(model_name): + model_path = model_name + else: + model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + + vae_config_path = os.path.join(model_path, "vae/config.json") + with open(vae_config_path) as f: + vae_config = json.load(f) + vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8 + + image_processor = Flux2ImageProcessor(vae_scale_factor=vae_scale_factor * 2) + + def post_process_func(images: torch.Tensor): + return image_processor.postprocess(images) + + return post_process_func + + +# Copied from diffusers.pipelines.flux2.pipeline_flux2.format_input +def format_input( + prompts: list[str], + system_message: str = SYSTEM_MESSAGE, + images: list[PIL.Image.Image] | list[list[PIL.Image.Image]] = None, +) -> list[list[dict[str, Any]]]: + """ + Format a batch of text prompts into the conversation format expected by apply_chat_template. Optionally, add images + to the input. + + Args: + prompts: List of text prompts + system_message: System message to use (default: CREATIVE_SYSTEM_MESSAGE) + images (optional): List of images to add to the input. + + Returns: + `list[list[dict[str, Any]]]`: List of conversations, where each conversation is a list of message dicts + """ + # Remove [IMG] tokens from prompts to avoid Pixtral validation issues + # when truncation is enabled. The processor counts [IMG] tokens and fails + # if the count changes after truncation. + cleaned_txt = [prompt.replace("[IMG]", "") for prompt in prompts] + + if images is None or len(images) == 0: + return [ + [ + { + "role": "system", + "content": [{"type": "text", "text": system_message}], + }, + {"role": "user", "content": [{"type": "text", "text": prompt}]}, + ] + for prompt in cleaned_txt + ] + else: + assert len(images) == len(prompts), "Number of images must match number of prompts" + messages = [ + [ + { + "role": "system", + "content": [{"type": "text", "text": system_message}], + }, + ] + for _ in cleaned_txt + ] + + for i, (el, batch_images) in enumerate(zip(messages, images)): + # optionally add the images per batch element. + if batch_images is not None: + el.append( + { + "role": "user", + "content": [{"type": "image", "image": image_obj} for image_obj in batch_images], + } + ) + # add the text. + el.append( + { + "role": "user", + "content": [{"type": "text", "text": cleaned_txt[i]}], + } + ) + + return messages + + +# Copied from diffusers.pipelines.flux2.pipeline_flux2._validate_and_process_images +def _validate_and_process_images( + images: list[list[PIL.Image.Image]] | list[PIL.Image.Image], + image_processor: Flux2ImageProcessor, + upsampling_max_image_size: int, +) -> list[list[PIL.Image.Image]]: + # Simple validation: ensure it's a list of PIL images or list of lists of PIL images + if not images: + return [] + + # Check if it's a list of lists or a list of images + if isinstance(images[0], PIL.Image.Image): + # It's a list of images, convert to list of lists + images = [[im] for im in images] + + # potentially concatenate multiple images to reduce the size + images = [[image_processor.concatenate_images(img_i)] if len(img_i) > 1 else img_i for img_i in images] + + # cap the pixels + images = [ + [image_processor._resize_if_exceeds_area(img_i, upsampling_max_image_size) for img_i in img_i] + for img_i in images + ] + return images + + +# Copied from diffusers.pipelines.flux2.pipeline_flux2.compute_empirical_mu +def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: + a1, b1 = 8.73809524e-05, 1.89833333 + a2, b2 = 0.00016927, 0.45666666 + + if image_seq_len > 4300: + mu = a2 * image_seq_len + b2 + return float(mu) + + m_200 = a2 * image_seq_len + b2 + m_10 = a1 * image_seq_len + b1 + + a = (m_200 - m_10) / 190.0 + b = m_200 - 200.0 * a + mu = a * num_steps + b + + return float(mu) + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +) -> tuple[torch.Tensor, int]: + r""" + Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles + custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. + + Args: + scheduler (`SchedulerMixin`): + The scheduler to get timesteps from. + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` + must be `None`. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + timesteps (`List[int]`, *optional*): + Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, + `num_inference_steps` and `sigmas` must be `None`. + sigmas (`List[float]`, *optional*): + Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, + `num_inference_steps` and `timesteps` must be `None`. + + Returns: + `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the + second element is the number of inference steps. + """ + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents(encoder_output: torch.Tensor, generator: torch.Generator = None, sample_mode: str = "sample"): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +class Flux2Pipeline(nn.Module, SupportImageInput): + """Flux2 pipeline for text-to-image generation.""" + + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + support_image_input = True + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + self._execution_device = get_local_device() + model = od_config.model + # Check if model is a local path + local_files_only = os.path.exists(model) + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + self.text_encoder = Mistral3ForConditionalGeneration.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ) + self.tokenizer = PixtralProcessor.from_pretrained( + model, subfolder="tokenizer", local_files_only=local_files_only + ) + self.vae = AutoencoderKLFlux2.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( + self._execution_device + ) + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel) + self.transformer = Flux2Transformer2DModel(**transformer_kwargs) + + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.tokenizer_max_length = 512 + self.default_sample_size = 128 + + self.system_message = SYSTEM_MESSAGE + self.system_message_upsampling_t2i = SYSTEM_MESSAGE_UPSAMPLING_T2I + self.system_message_upsampling_i2i = SYSTEM_MESSAGE_UPSAMPLING_I2I + self.upsampling_max_image_size = UPSAMPLING_MAX_IMAGE_SIZE + + self._guidance_scale = None + self._attention_kwargs = None + self._num_timesteps = None + self._current_timestep = None + self._interrupt = False + + @staticmethod + def _get_mistral_3_small_prompt_embeds( + text_encoder: Mistral3ForConditionalGeneration, + tokenizer: AutoProcessor, + prompt: str | list[str], + dtype: torch.dtype | None = None, + device: torch.device | None = None, + max_sequence_length: int = 512, + system_message: str = SYSTEM_MESSAGE, + hidden_states_layers: list[int] = (10, 20, 30), + ): + dtype = text_encoder.dtype if dtype is None else dtype + device = text_encoder.device if device is None else device + + prompt = [prompt] if isinstance(prompt, str) else prompt + + # Format input messages + messages_batch = format_input(prompts=prompt, system_message=system_message) + + # Process all messages at once + inputs = tokenizer.apply_chat_template( + messages_batch, + add_generation_prompt=False, + tokenize=True, + return_dict=True, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=max_sequence_length, + ) + + # Move to device + input_ids = inputs["input_ids"].to(device) + attention_mask = inputs["attention_mask"].to(device) + + # Forward pass through the model + output = text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + use_cache=False, + ) + + # Only use outputs from intermediate layers and stack them + out = torch.stack([output.hidden_states[k] for k in hidden_states_layers], dim=1) + out = out.to(dtype=dtype, device=device) + + batch_size, num_channels, seq_len, hidden_dim = out.shape + prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim) + + return prompt_embeds + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_text_ids + def _prepare_text_ids( + x: torch.Tensor, # (B, L, D) or (L, D) + t_coord: torch.Tensor | None = None, + ): + B, L, _ = x.shape + out_ids = [] + + for i in range(B): + t = torch.arange(1) if t_coord is None else t_coord[i] + h = torch.arange(1) + w = torch.arange(1) + seq_positions = torch.arange(L) + + coords = torch.cartesian_prod(t, h, w, seq_positions) + out_ids.append(coords) + + return torch.stack(out_ids) + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_latent_ids + def _prepare_latent_ids( + latents: torch.Tensor, # (B, C, H, W) + ): + r""" + Generates 4D position coordinates (T, H, W, L) for latent tensors. + + Args: + latents (torch.Tensor): + Latent tensor of shape (B, C, H, W) + + Returns: + torch.Tensor: + Position IDs tensor of shape (B, H*W, 4) All batches share the same coordinate structure: T=0, + H=[0..H-1], W=[0..W-1], L=0 + """ + + batch_size, _, height, width = latents.shape + + t = torch.arange(1) # [0] - time dimension + h = torch.arange(height) + w = torch.arange(width) + layer_ids = torch.arange(1) # [0] - layer dimension + + # Create position IDs: (H*W, 4) + latent_ids = torch.cartesian_prod(t, h, w, layer_ids) + + # Expand to batch: (B, H*W, 4) + latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1) + + return latent_ids + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_image_ids + def _prepare_image_ids( + image_latents: list[torch.Tensor], # [(1, C, H, W), (1, C, H, W), ...] + scale: int = 10, + ): + r""" + Generates 4D time-space coordinates (T, H, W, L) for a sequence of image latents. + + This function creates a unique coordinate for every pixel/patch across all input latent with different + dimensions. + + Args: + image_latents (List[torch.Tensor]): + A list of image latent feature tensors, typically of shape (C, H, W). + scale (int, optional): + A factor used to define the time separation (T-coordinate) between latents. T-coordinate for the i-th + latent is: 'scale + scale * i'. Defaults to 10. + + Returns: + torch.Tensor: + The combined coordinate tensor. Shape: (1, N_total, 4) Where N_total is the sum of (H * W) for all + input latents. + + Coordinate Components (Dimension 4): + - T (Time): The unique index indicating which latent image the coordinate belongs to. + - H (Height): The row index within that latent image. + - W (Width): The column index within that latent image. + - L (Seq. Length): A sequence length dimension, which is always fixed at 0 (size 1) + """ + + if not isinstance(image_latents, list): + raise ValueError(f"Expected `image_latents` to be a list, got {type(image_latents)}.") + + # create time offset for each reference image + t_coords = [scale + scale * t for t in torch.arange(0, len(image_latents))] + t_coords = [t.view(-1) for t in t_coords] + + image_latent_ids = [] + for x, t in zip(image_latents, t_coords): + x = x.squeeze(0) + _, height, width = x.shape + + x_ids = torch.cartesian_prod(t, torch.arange(height), torch.arange(width), torch.arange(1)) + image_latent_ids.append(x_ids) + + image_latent_ids = torch.cat(image_latent_ids, dim=0) + image_latent_ids = image_latent_ids.unsqueeze(0) + + return image_latent_ids + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._patchify_latents + def _patchify_latents(latents): + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + latents = latents.reshape(batch_size, num_channels_latents * 4, height // 2, width // 2) + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpatchify_latents + def _unpatchify_latents(latents): + batch_size, num_channels_latents, height, width = latents.shape + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width) + latents = latents.permute(0, 1, 4, 2, 5, 3) + latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), height * 2, width * 2) + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._pack_latents + def _pack_latents(latents): + """ + pack latents: (batch_size, num_channels, height, width) -> (batch_size, height * width, num_channels) + """ + + batch_size, num_channels, height, width = latents.shape + latents = latents.reshape(batch_size, num_channels, height * width).permute(0, 2, 1) + + return latents + + @staticmethod + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpack_latents_with_ids + def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> list[torch.Tensor]: + """ + using position ids to scatter tokens into place + """ + x_list = [] + for data, pos in zip(x, x_ids): + _, ch = data.shape # noqa: F841 + h_ids = pos[:, 1].to(torch.int64) + w_ids = pos[:, 2].to(torch.int64) + + h = torch.max(h_ids) + 1 + w = torch.max(w_ids) + 1 + + flat_ids = h_ids * w + w_ids + + out = torch.zeros((h * w, ch), device=data.device, dtype=data.dtype) + out.scatter_(0, flat_ids.unsqueeze(1).expand(-1, ch), data) + + # reshape from (H * W, C) to (H, W, C) and permute to (C, H, W) + + out = out.view(h, w, ch).permute(2, 0, 1) + x_list.append(out) + + return torch.stack(x_list, dim=0) + + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline.upsample_prompt + def upsample_prompt( + self, + prompt: str | list[str], + images: list[PIL.Image.Image] | list[list[PIL.Image.Image]] = None, + temperature: float = 0.15, + device: torch.device = None, + ) -> list[str]: + prompt = [prompt] if isinstance(prompt, str) else prompt + device = self.text_encoder.device if device is None else device + + # Set system message based on whether images are provided + if images is None or len(images) == 0 or images[0] is None: + system_message = SYSTEM_MESSAGE_UPSAMPLING_T2I + else: + system_message = SYSTEM_MESSAGE_UPSAMPLING_I2I + + # Validate and process the input images + if images: + images = _validate_and_process_images(images, self.image_processor, self.upsampling_max_image_size) + + # Format input messages + messages_batch = format_input(prompts=prompt, system_message=system_message, images=images) + + # Process all messages at once + # with image processing a too short max length can throw an error in here. + inputs = self.tokenizer.apply_chat_template( + messages_batch, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=2048, + ) + + # Move to device + inputs["input_ids"] = inputs["input_ids"].to(device) + inputs["attention_mask"] = inputs["attention_mask"].to(device) + + if "pixel_values" in inputs: + inputs["pixel_values"] = inputs["pixel_values"].to(device, self.text_encoder.dtype) + + # Generate text using the model's generate method + generated_ids = self.text_encoder.generate( + **inputs, + max_new_tokens=512, + do_sample=True, + temperature=temperature, + use_cache=True, + ) + + # Decode only the newly generated tokens (skip input tokens) + # Extract only the generated portion + input_length = inputs["input_ids"].shape[1] + generated_tokens = generated_ids[:, input_length:] + + upsampled_prompt = self.tokenizer.tokenizer.batch_decode( + generated_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=True + ) + return upsampled_prompt + + def encode_prompt( + self, + prompt: str | list[str], + device: torch.device | None = None, + num_images_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + max_sequence_length: int = 512, + text_encoder_out_layers: tuple[int, ...] = (10, 20, 30), + ): + device = device or self._execution_device + + if prompt is None: + prompt = "" + + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt_embeds is None: + prompt_embeds = self._get_mistral_3_small_prompt_embeds( + text_encoder=self.text_encoder, + tokenizer=self.tokenizer, + prompt=prompt, + device=device, + max_sequence_length=max_sequence_length, + system_message=self.system_message, + hidden_states_layers=text_encoder_out_layers, + ) + + batch_size, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + text_ids = self._prepare_text_ids(prompt_embeds) + text_ids = text_ids.to(device) + return prompt_embeds, text_ids + + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._encode_vae_image + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if image.ndim != 4: + raise ValueError(f"Expected image dims 4, got {image.ndim}.") + + image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax") + image_latents = self._patchify_latents(image_latents) + + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(image_latents.device, image_latents.dtype) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps) + image_latents = (image_latents - latents_bn_mean) / latents_bn_std + + return image_latents + + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline.prepare_latents + def prepare_latents( + self, + batch_size, + num_latents_channels, + height, + width, + dtype, + device, + generator: torch.Generator, + latents: torch.Tensor | None = None, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, num_latents_channels * 4, height // 2, width // 2) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device=device, dtype=dtype) + + latent_ids = self._prepare_latent_ids(latents) + latent_ids = latent_ids.to(device) + + latents = self._pack_latents(latents) # [B, C, H, W] -> [B, H*W, C] + return latents, latent_ids + + # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline.prepare_image_latents + def prepare_image_latents( + self, + images: list[torch.Tensor], + batch_size, + generator: torch.Generator, + device, + dtype, + ): + image_latents = [] + for image in images: + image = image.to(device=device, dtype=dtype) + image_latent = self._encode_vae_image(image=image, generator=generator) + image_latents.append(image_latent) # (1, 128, 32, 32) + + image_latent_ids = self._prepare_image_ids(image_latents) + + # Pack each latent and concatenate + packed_latents = [] + for latent in image_latents: + # latent: (1, 128, 32, 32) + packed = self._pack_latents(latent) # (1, 1024, 128) + packed = packed.squeeze(0) # (1024, 128) - remove batch dim + packed_latents.append(packed) + + # Concatenate all reference tokens along sequence dimension + image_latents = torch.cat(packed_latents, dim=0) # (N*1024, 128) + image_latents = image_latents.unsqueeze(0) # (1, N*1024, 128) + + image_latents = image_latents.repeat(batch_size, 1, 1) + image_latent_ids = image_latent_ids.repeat(batch_size, 1, 1) + image_latent_ids = image_latent_ids.to(device) + + return image_latents, image_latent_ids + + def check_inputs( + self, + prompt, + height, + width, + prompt_embeds=None, + callback_on_step_end_tensor_inputs=None, + ): + if ( + height is not None + and height % (self.vae_scale_factor * 2) != 0 + or width is not None + and width % (self.vae_scale_factor * 2) != 0 + ): + logger.warning( + "`height` and `width` have to be divisible by %s but are %s and %s. " + "Dimensions will be resized accordingly", + self.vae_scale_factor * 2, + height, + width, + ) + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, " + f"but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + def forward( + self, + req: OmniDiffusionRequest, + image: PIL.Image.Image | list[PIL.Image.Image] | None = None, + prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + sigmas: list[float] | None = None, + guidance_scale: float | None = 4.0, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int, dict], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 512, + text_encoder_out_layers: tuple[int, ...] = (10, 20, 30), + caption_upsample_temperature: float = None, + ) -> DiffusionOutput: + if len(req.prompts) > 1: + logger.warning( + """This model only supports a single prompt, not a batched request.""", + """Taking only the first image for now.""", + ) + first_prompt = req.prompts[0] + prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") + + if ( + raw_image := None + if isinstance(first_prompt, str) + else first_prompt.get("multi_modal_data", {}).get("image") + ) is None: + pass # use image from param list + elif isinstance(raw_image, list): + image = [PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image, im) for im in raw_image] + else: + image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else cast(PIL.Image.Image, raw_image) + + height = req.sampling_params.height or height + width = req.sampling_params.width or width + num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + sigmas = req.sampling_params.sigmas or sigmas + guidance_scale = ( + req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale + ) + generator = req.sampling_params.generator or generator + num_images_per_prompt = ( + req.sampling_params.num_outputs_per_prompt + if req.sampling_params.num_outputs_per_prompt > 0 + else num_images_per_prompt + ) + max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length + text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers) + + req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts] + if any(p is not None for p in req_prompt_embeds): + # If at list one prompt is provided as an embedding, + # Then assume that the user wants to provide embeddings for all prompts, and enter this if block + # If the user in fact provides mixed input format, req_prompt_embeds will have some None's + # And `torch.stack` automatically raises an exception for us + prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore # intentionally expect TypeError + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt=prompt, + height=height, + width=width, + prompt_embeds=prompt_embeds, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + + # 3. prepare text embeddings + if caption_upsample_temperature: + prompt = self.upsample_prompt(prompt, images=image, temperature=caption_upsample_temperature, device=device) + prompt_embeds, text_ids = self.encode_prompt( + prompt=prompt, + prompt_embeds=prompt_embeds, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + text_encoder_out_layers=text_encoder_out_layers, + ) + + # 4. process images + if image is not None and not isinstance(image, list): + image = [image] + + condition_images = None + if image is not None: + for img in image: + self.image_processor.check_image_input(img) + + condition_images = [] + for img in image: + image_width, image_height = img.size + if image_width * image_height > 1024 * 1024: + img = self.image_processor._resize_to_target_area(img, 1024 * 1024) + image_width, image_height = img.size + + multiple_of = self.vae_scale_factor * 2 + image_width = (image_width // multiple_of) * multiple_of + image_height = (image_height // multiple_of) * multiple_of + img = self.image_processor.preprocess(img, height=image_height, width=image_width, resize_mode="crop") + condition_images.append(img) + height = height or image_height + width = width or image_width + + height = height or self.default_sample_size * self.vae_scale_factor + width = width or self.default_sample_size * self.vae_scale_factor + + # 5. prepare latent variables + num_channels_latents = self.transformer.config.in_channels // 4 + latents, latent_ids = self.prepare_latents( + batch_size=batch_size * num_images_per_prompt, + num_latents_channels=num_channels_latents, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=device, + generator=generator, + latents=latents, + ) + + image_latents = None + image_latent_ids = None + if condition_images is not None: + image_latents, image_latent_ids = self.prepare_image_latents( + images=condition_images, + batch_size=batch_size * num_images_per_prompt, + generator=generator, + device=device, + dtype=self.vae.dtype, + ) + + # 6. Prepare timesteps + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + if hasattr(self.scheduler.config, "use_flow_sigmas") and self.scheduler.config.use_flow_sigmas: + sigmas = None + image_seq_len = latents.shape[1] + mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=num_inference_steps) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + mu=mu, + ) + self._num_timesteps = len(timesteps) + + # 7. Denoising loop + # We set the index here to remove DtoH sync, helpful especially during compilation. + # Check out more details here: https://github.com/huggingface/diffusers/pull/11696 + self.scheduler.set_begin_index(0) + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + timestep = t.expand(latents.shape[0]).to(latents.dtype) + + latent_model_input = latents.to(self.transformer.dtype) + latent_image_ids = latent_ids + + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1).to(self.transformer.dtype) + latent_image_ids = torch.cat([latent_ids, image_latent_ids], dim=1) + + noise_pred = self.transformer( + hidden_states=latent_model_input, # (B, image_seq_len, C) + timestep=timestep / 1000, + guidance=None, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, # B, text_seq_len, 4 + img_ids=latent_image_ids, # B, image_seq_len, 4 + joint_attention_kwargs=self.attention_kwargs, + return_dict=False, + )[0] + + noise_pred = noise_pred[:, : latents.size(1) :] + + # compute the previous noisy sample x_t -> x_t-1 + latents_dtype = latents.dtype + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + if latents.dtype != latents_dtype and torch.backends.mps.is_available(): + latents = latents.to(latents_dtype) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + self._current_timestep = None + + latents = self._unpack_latents_with_ids(latents, latent_ids) + + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to( + latents.device, latents.dtype + ) + latents = latents * latents_bn_std + latents_bn_mean + latents = self._unpatchify_latents(latents) + if output_type == "latent": + image = latents + else: + if latents.dtype != self.vae.dtype: + latents = latents.to(self.vae.dtype) + image = self.vae.decode(latents, return_dict=False)[0] + + return DiffusionOutput(output=image) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index ba15a62b0..9fc6dcf7f 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -17,7 +17,7 @@ import torch.nn.functional as F from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from torch import nn -from transformers import AutoTokenizer, UMT5EncoderModel +from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig @@ -174,8 +174,17 @@ class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): ] self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + # Helios checkpoints store embed_tokens under ``shared.weight`` only, + # but the published config sets ``tie_word_embeddings=False``. When + # transformers sees ``tie=False`` it creates a separate + # ``encoder.embed_tokens.weight`` that is never loaded from the + # checkpoint, leaving it as all-zeros. This silently destroys prompt + # encoding and produces grey/meaningless video output. Force tying so + # that ``embed_tokens`` shares ``shared.weight`` as intended. + text_enc_cfg = AutoConfig.from_pretrained(model, subfolder="text_encoder", local_files_only=local_files_only) + text_enc_cfg.tie_word_embeddings = True self.text_encoder = UMT5EncoderModel.from_pretrained( - model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only + model, subfolder="text_encoder", config=text_enc_cfg, torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) self.vae = AutoencoderKLWan.from_pretrained( model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py b/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py index 98a3ac07b..cbc6a8ad1 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Hunyuan Image 3 diffusion model components.""" +from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import HunyuanFusedMoE from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_image_3_transformer import ( HunyuanImage3Model, HunyuanImage3Text2ImagePipeline, @@ -10,8 +11,4 @@ from vllm_omni.diffusion.models.hunyuan_image_3.pipeline_hunyuan_image_3 import HunyuanImage3Pipeline, ) -__all__ = [ - "HunyuanImage3Pipeline", - "HunyuanImage3Model", - "HunyuanImage3Text2ImagePipeline", -] +__all__ = ["HunyuanImage3Pipeline", "HunyuanImage3Model", "HunyuanImage3Text2ImagePipeline", "HunyuanFusedMoE"] diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py new file mode 100644 index 000000000..5ada5ceb8 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +from vllm.model_executor.layers.fused_moe import SharedFusedMoE +from vllm.utils.import_utils import resolve_obj_by_qualname + +from vllm_omni.platforms import current_omni_platform + + +class HunyuanFusedMoEDefault(SharedFusedMoE): + def __init__(self, *, prefix: str = "", **kwargs: Any) -> None: + super().__init__(prefix=prefix, **kwargs) + self._prefix = prefix + self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True) + + def _initialize_kernel_hook(self, module: Any, args: Any, kwargs: Any) -> None: + if self.quant_method: + self.quant_method.process_weights_after_loading(self) + self._init_hook_handle.remove() + + def forward(self, hidden_states: Any, router_logits: Any) -> Any: + return super().forward(hidden_states, router_logits) + + +class HunyuanFusedMoE: + def __new__(cls, *, prefix: str = "", **kwargs: Any) -> Any: + op_name = "hunyuan_fused_moe" + current_omni_platform.prepare_diffusion_op_runtime(op_name) + impl = resolve_obj_by_qualname( + current_omni_platform.get_diffusion_model_impl_qualname(op_name), + ) + return impl(prefix=prefix, **kwargs) + + @classmethod + def make_expert_params_mapping( + cls, + model: Any, + ckpt_gate_proj_name: str, + ckpt_down_proj_name: str, + ckpt_up_proj_name: str, + num_experts: int, + num_redundant_experts: int = 0, + ) -> list[tuple[str, str, int, str]]: + impl = resolve_obj_by_qualname( + current_omni_platform.get_diffusion_model_impl_qualname("hunyuan_fused_moe"), + ) + return impl.make_expert_params_mapping( + model, + ckpt_gate_proj_name=ckpt_gate_proj_name, + ckpt_down_proj_name=ckpt_down_proj_name, + ckpt_up_proj_name=ckpt_up_proj_name, + num_experts=num_experts, + num_redundant_experts=num_redundant_experts, + ) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py index 0c2f9e290..a89931550 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py @@ -32,10 +32,8 @@ from transformers.modeling_utils import PreTrainedModel from vllm.config import CacheConfig from vllm.distributed import ( get_tensor_model_parallel_world_size, - tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.fused_moe import SharedFusedMoE from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -63,6 +61,7 @@ from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.distributed.parallel_state import get_pp_group from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import HunyuanFusedMoE logger = logging.getLogger(__name__) @@ -1417,7 +1416,7 @@ class HunYuanSparseMoeBlock(nn.Module): final_hidden_states = final_hidden_states[0] + final_hidden_states[1] if self.tp_size > 1: - final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(final_hidden_states) return final_hidden_states.view(orig_shape) @@ -1565,22 +1564,6 @@ class HunYuanAttention(nn.Module): return output, None, past_key_value -class HunyuanFusedMoE(SharedFusedMoE): - def __init__(self, *, prefix: str = "", **kwargs): - super().__init__(prefix=prefix, **kwargs) - self._prefix = prefix - - self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True) - - def _initialize_kernel_hook(self, module, args, kwargs): - if self.quant_method: - self.quant_method.process_weights_after_loading(self) - self._init_hook_handle.remove() - - def forward(self, hidden_states, router_logits): - return super().forward(hidden_states, router_logits) - - class HunyuanImage3DecoderLayer(nn.Module): def __init__(self, config: HunyuanImage3Config, layer_idx: int, prefix: str = ""): super().__init__() @@ -2454,7 +2437,6 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) latents = callback_outputs.pop("latents", latents) - # call the callback, if provided if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index ccf35ea4e..8d8e523d6 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py @@ -18,11 +18,10 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.data import OmniDiffusionConfig -from vllm_omni.diffusion.distributed.parallel_state import ( - get_sequence_parallel_rank, - get_sequence_parallel_world_size, - get_sp_group, +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, ) from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.platforms import current_omni_platform @@ -50,6 +49,7 @@ class FeedForward(nn.Module): class LongCatImageAttention(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, query_dim: int, heads: int = 8, dim_head: int = 64, @@ -64,7 +64,7 @@ class LongCatImageAttention(nn.Module): pre_only: bool = False, ): super().__init__() - + self.parallel_config = parallel_config self.head_dim = dim_head self.inner_dim = out_dim if out_dim is not None else dim_head * heads self.query_dim = query_dim @@ -209,12 +209,8 @@ class LongCatImageAttention(nn.Module): encoder_query = self.norm_added_q(encoder_query) encoder_key = self.norm_added_k(encoder_key) - # Check if SP is enabled from forward context (set by LongCatImageTransformer2DModel) - forward_ctx = get_forward_context() - sp_size = forward_ctx.sequence_parallel_size - use_sp_joint_attention = sp_size > 1 and not forward_ctx.split_text_embed_in_sp - - if use_sp_joint_attention: + sp_size = self.parallel_config.sequence_parallel_size + if sp_size is not None and sp_size > 1: # SP Mode: Use common helper for RoPE + joint attention hidden_states = self._sp_attention_with_rope( img_query=query, @@ -248,12 +244,22 @@ class LongCatImageAttention(nn.Module): # In SP mode, image part is chunked: (B, txt_len + img_len/SP, D) # Check if SP is enabled and we have text_seq_len info - forward_ctx = get_forward_context() - sp_size = forward_ctx.sequence_parallel_size + sp_size = self.parallel_config.sequence_parallel_size text_seq_len = kwargs.get("text_seq_len", None) - use_sp_single_stream = sp_size > 1 and not forward_ctx.split_text_embed_in_sp and text_seq_len is not None + forward_ctx = get_forward_context() + + if ( + sp_size is not None + and sp_size > 1 + and not forward_ctx.split_text_embed_in_sp + and text_seq_len is not None + ): + # Ensure that the SP split won't cause out of bounds issues. + if text_seq_len < 0 or text_seq_len > query.shape[1]: + raise ValueError( + f"text_seq_len={text_seq_len} is out of bounds for sequence length {query.shape[1]}" + ) - if use_sp_single_stream: # SP Mode for single-stream block: # Split QKV into text and image parts, then use common helper hidden_states = self._sp_attention_with_rope( @@ -301,6 +307,7 @@ class LongCatImageAttention(nn.Module): class LongCatImageTransformerBlock(nn.Module): def __init__( self, + parallel_config: DiffusionParallelConfig, dim: int, num_attention_heads: int, attention_head_dim: int, @@ -309,10 +316,12 @@ class LongCatImageTransformerBlock(nn.Module): ): super().__init__() + self.parallel_config = parallel_config self.norm1 = AdaLayerNormZero(dim) self.norm1_context = AdaLayerNormZero(dim) self.attn = LongCatImageAttention( + parallel_config=parallel_config, query_dim=dim, added_kv_proj_dim=dim, dim_head=attention_head_dim, @@ -429,6 +438,61 @@ class LongCatImageTimestepEmbeddings(nn.Module): return timesteps_emb +class RoPEPreparer(nn.Module): + """ + This module encapsulates RoPE computation to enable _sp_plan sharding + for text / image components. + + For LongCat, which uses dual-stream attention, this means that text + components are replicated across SP ranks, while image components are + sharded. + """ + + def __init__(self, pos_embed: LongCatImagePosEmbed): + super().__init__() + self.pos_embed = pos_embed + + def forward( + self, + txt_ids: torch.Tensor, + img_ids: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Compute RoPE embeddings for text and image sequences. + + Args: + txt_ids: Text position IDs (txt_seq_len, n_axes) + img_ids: Image position IDs (img_seq_len, n_axes) + + Returns: + Tuple of cosine / sine components for text & image + in the order: (txt_cos, txt_sin, img_cos, img_sin) + + NOTE: careful about output orders if this is refactored in the + future; we need to match the _sp_plan indices, since text + components (0 & 1) need to be replicated across SP ranks, + while image components (2 & 3) must be sharded. + """ + # Concatenate and compute RoPE for full sequence + ids = torch.cat((txt_ids, img_ids), dim=0) + + if current_omni_platform.is_npu(): + freqs_cos, freqs_sin = self.pos_embed(ids.cpu()) + freqs_cos = freqs_cos.npu() + freqs_sin = freqs_sin.npu() + else: + freqs_cos, freqs_sin = self.pos_embed(ids) + + # Split into text and image portions + txt_len = txt_ids.shape[0] + txt_cos = freqs_cos[:txt_len] + txt_sin = freqs_sin[:txt_len] + img_cos = freqs_cos[txt_len:] + img_sin = freqs_sin[txt_len:] + + return txt_cos, txt_sin, img_cos, img_sin + + class LongCatImageSingleTransformerBlock(nn.Module): """ Single-stream Transformer block for LongCat with SP (Sequence Parallelism) support. @@ -437,7 +501,14 @@ class LongCatImageSingleTransformerBlock(nn.Module): This keeps the block logic clean and centralizes SP logic in the attention layer. """ - def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float = 4.0): + def __init__( + self, + parallel_config: DiffusionParallelConfig, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + mlp_ratio: float = 4.0, + ): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) @@ -448,6 +519,7 @@ class LongCatImageSingleTransformerBlock(nn.Module): # SP handling is delegated to LongCatImageAttention via text_seq_len kwarg self.attn = LongCatImageAttention( + parallel_config=parallel_config, query_dim=dim, dim_head=attention_head_dim, heads=num_attention_heads, @@ -511,6 +583,23 @@ class LongCatImageTransformer2DModel(nn.Module): _repeated_blocks = ["LongCatImageTransformerBlock", "LongCatImageSingleTransformerBlock"] + # Sequence Parallelism for LongCat (following diffusers' _cp_plan pattern) + _sp_plan = { + "": { + # Chunk the hidden states prior to the forward() + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Shard RoPE image embeddings after rope_preparer computes them + # Outputs 0, 1 are text components, so they aren't sharded + # Outputs 2, 3 are image components and are sharded + "rope_preparer": { + 2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), + 3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), + }, + # Gather at the last linear projection + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + def __init__( self, od_config: OmniDiffusionConfig, @@ -534,6 +623,7 @@ class LongCatImageTransformer2DModel(nn.Module): self.parallel_config = od_config.parallel_config self.pos_embed = LongCatImagePosEmbed(theta=10000, axes_dim=axes_dims_rope) + self.rope_preparer = RoPEPreparer(self.pos_embed) self.time_embed = LongCatImageTimestepEmbeddings(embedding_dim=self.inner_dim) @@ -543,6 +633,7 @@ class LongCatImageTransformer2DModel(nn.Module): self.transformer_blocks = nn.ModuleList( [ LongCatImageTransformerBlock( + parallel_config=self.parallel_config, dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, @@ -554,6 +645,7 @@ class LongCatImageTransformer2DModel(nn.Module): self.single_transformer_blocks = nn.ModuleList( [ LongCatImageSingleTransformerBlock( + parallel_config=self.parallel_config, dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, @@ -581,40 +673,11 @@ class LongCatImageTransformer2DModel(nn.Module): return_dict: bool = True, ) -> torch.FloatTensor | Transformer2DModelOutput: fwd_context = get_forward_context() - # Before: hidden_states shape = (B, img_seq_len, in_channels) - # After: hidden_states shape = (B, img_seq_len // SP, in_channels) sp_size = self.parallel_config.sequence_parallel_size - # Store SP size in forward context for sub-modules to access if sp_size is not None and sp_size > 1: - sp_world_size = get_sequence_parallel_world_size() - sp_rank = get_sequence_parallel_rank() - original_shape = hidden_states.shape - hidden_states = torch.chunk(hidden_states, sp_world_size, dim=1)[sp_rank] - # LongCat uses dual-stream (text + image) with joint attention - # Text embeddings should be replicated across SP ranks for correctness - fwd_context.sequence_parallel_size = sp_size fwd_context.split_text_embed_in_sp = False - # Mark SP as active so attention layers; we need this to ensure we use - # Ulysses instead of NoParallelAttention since we don't set an sp plan - # for this model. - # TODO: would be nice to refactor this to use sp_plan if possible to - # tracking this directly, even though we only have one level. - fwd_context._sp_shard_depth = 1 - - # Debug log (only first forward) - if not hasattr(self, "_sp_forward_logged"): - self._sp_forward_logged = True - logger.info( - f"[LongCat Transformer] SP enabled: sp_size={sp_size}, world_size={sp_world_size}, " - f"rank={sp_rank}, original_shape={original_shape}, chunked_shape={hidden_states.shape}" - ) - else: - fwd_context._sp_shard_depth = 0 - if not hasattr(self, "_sp_forward_logged"): - self._sp_forward_logged = True - logger.info(f"[LongCat Transformer] SP disabled: sp_size={sp_size}") - + # Hidden states are sharded prior to forward() when sp is active hidden_states = self.x_embedder(hidden_states) timestep = timestep.to(hidden_states.dtype) * 1000 @@ -622,46 +685,17 @@ class LongCatImageTransformer2DModel(nn.Module): temb = self.time_embed(timestep, hidden_states.dtype) encoder_hidden_states = self.context_embedder(encoder_hidden_states) - ids = torch.cat((txt_ids, img_ids), dim=0) + # Compute RoPE embeddings via rope_preparer module + # _sp_plan will automatically shard img_cos/img_sin (outputs 2, 3) + # txt_cos/txt_sin (outputs 0, 1) remain replicated for dual-stream attention + txt_cos, txt_sin, img_cos, img_sin = self.rope_preparer(txt_ids, img_ids) - if current_omni_platform.is_npu(): - freqs_cos, freqs_sin = self.pos_embed(ids.cpu()) - image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu()) - else: - image_rotary_emb = self.pos_embed(ids) - - # SP: Chunk RoPE embeddings along sequence dimension - if sp_size is not None and sp_size > 1: - sp_world_size = get_sequence_parallel_world_size() - sp_rank = get_sequence_parallel_rank() - freqs_cos, freqs_sin = image_rotary_emb - txt_len = txt_ids.shape[0] - - # Split RoPE into text and image portions - # txt_freqs: (txt_seq_len, head_dim) - keep full for all ranks - # img_freqs: (img_seq_len, head_dim) -> (img_seq_len // SP, head_dim) - txt_freqs_cos = freqs_cos[:txt_len] - txt_freqs_sin = freqs_sin[:txt_len] - img_freqs_cos = freqs_cos[txt_len:] - img_freqs_sin = freqs_sin[txt_len:] - - # Chunk image RoPE for each SP rank - # img_freqs_cos: (img_seq_len // SP, head_dim) - # img_freqs_sin: (img_seq_len // SP, head_dim) - img_freqs_cos = torch.chunk(img_freqs_cos, sp_world_size, dim=0)[sp_rank] - img_freqs_sin = torch.chunk(img_freqs_sin, sp_world_size, dim=0)[sp_rank] - - # Optionally chunk text RoPE if split_text_embed_in_sp is True - if get_forward_context().split_text_embed_in_sp: - txt_freqs_cos = torch.chunk(txt_freqs_cos, sp_world_size, dim=0)[sp_rank] - txt_freqs_sin = torch.chunk(txt_freqs_sin, sp_world_size, dim=0)[sp_rank] - - # Reconstruct image_rotary_emb with chunked values - # Final shape: (txt_seq_len + img_seq_len // SP, head_dim) - image_rotary_emb = ( - torch.cat([txt_freqs_cos, img_freqs_cos], dim=0), - torch.cat([txt_freqs_sin, img_freqs_sin], dim=0), - ) + # Reconstruct image_rotary_emb with chunked values + # Final shape: (txt_seq_len + img_seq_len // SP, head_dim) + image_rotary_emb = ( + torch.cat([txt_cos, img_cos], dim=0), + torch.cat([txt_sin, img_sin], dim=0), + ) for block in self.transformer_blocks: encoder_hidden_states, hidden_states = block( @@ -680,13 +714,9 @@ class LongCatImageTransformer2DModel(nn.Module): ) hidden_states = self.norm_out(hidden_states, temb) - output = self.proj_out(hidden_states) - # SP: All-gather output to reconstruct full sequence - if sp_size is not None and sp_size > 1: - output = get_sp_group().all_gather(output, dim=1) - # Mark SP as inactive after gathering - get_forward_context()._sp_shard_depth = 0 + # proj_out gathers for sequence parallel + output = self.proj_out(hidden_states) if not return_dict: return (output,) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 1b48c9c7d..d5af2de93 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -132,6 +132,11 @@ _DIFFUSION_MODELS = { "pipeline_helios", "HeliosPipeline", ), + "Flux2Pipeline": ( + "flux2", + "pipeline_flux2", + "Flux2Pipeline", + ), } @@ -305,6 +310,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "OmniGen2Pipeline": "get_omnigen2_post_process_func", "HeliosPipeline": "get_helios_post_process_func", "HeliosPyramidPipeline": "get_helios_post_process_func", + "Flux2Pipeline": "get_flux2_post_process_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { diff --git a/vllm_omni/distributed/omni_coordinator/__init__.py b/vllm_omni/distributed/omni_coordinator/__init__.py new file mode 100644 index 000000000..cbef920d4 --- /dev/null +++ b/vllm_omni/distributed/omni_coordinator/__init__.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .load_balancer import ( + LoadBalancer, + LoadBalancingPolicy, + RandomBalancer, + Task, +) +from .messages import InstanceEvent, InstanceInfo, InstanceList, StageStatus +from .omni_coord_client_for_hub import OmniCoordClientForHub +from .omni_coord_client_for_stage import OmniCoordClientForStage +from .omni_coordinator import OmniCoordinator + +__all__ = [ + "OmniCoordinator", + "StageStatus", + "InstanceEvent", + "InstanceInfo", + "InstanceList", + "OmniCoordClientForStage", + "OmniCoordClientForHub", + "Task", + "LoadBalancer", + "LoadBalancingPolicy", + "RandomBalancer", +] diff --git a/vllm_omni/distributed/omni_coordinator/load_balancer.py b/vllm_omni/distributed/omni_coordinator/load_balancer.py new file mode 100644 index 000000000..15a079b0a --- /dev/null +++ b/vllm_omni/distributed/omni_coordinator/load_balancer.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import random +from abc import ABC, abstractmethod +from enum import Enum +from typing import Any, TypedDict + +from .messages import InstanceInfo + + +class Task(TypedDict, total=False): + """Task structure passed from async_omni (stage.submit(task)). + + Mirrors the dict built in AsyncOmni with request_id, engine_inputs, + sampling_params. Future load-balancing policies may use these fields. + """ + + request_id: str + engine_inputs: Any + sampling_params: Any + + +class LoadBalancingPolicy(str, Enum): + """Enumeration for load balancing policies. + + Only ``RANDOM`` is implemented. Additional policies (e.g. round-robin, + least-connections) can be added in the future. + """ + + RANDOM = "random" + + +class LoadBalancer(ABC): + """Abstract base class for load balancers. + + Subclasses implement :meth:`select` to choose an instance for a given task. + """ + + @abstractmethod + def select(self, task: Task, instances: list[InstanceInfo]) -> int: + """Route a task to one of the available instances. + + Args: + task: The task to route. Not used by the random policy but reserved + for future strategies that may inspect task metadata. + instances: List of available instances to choose from. + + Returns: + Index of the selected instance in ``instances``. + + Raises: + ValueError: If ``instances`` is empty. + """ + + raise NotImplementedError + + +class RandomBalancer(LoadBalancer): + """Load balancer that selects an instance uniformly at random. + + This is the initial and only policy supported. It intentionally ignores + the task payload and chooses a random index from the provided instance + list. More sophisticated policies (e.g. round-robin, least-connections) + can be implemented as additional subclasses of :class:`LoadBalancer`. + """ + + def select(self, task: Task, instances: list[InstanceInfo]) -> int: # noqa: ARG002 + if not instances: + raise ValueError("instances must not be empty") + + return random.randrange(len(instances)) + + +__all__ = [ + "Task", + "LoadBalancingPolicy", + "LoadBalancer", + "RandomBalancer", +] diff --git a/vllm_omni/distributed/omni_coordinator/messages.py b/vllm_omni/distributed/omni_coordinator/messages.py new file mode 100644 index 000000000..2bb590139 --- /dev/null +++ b/vllm_omni/distributed/omni_coordinator/messages.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + + +class StageStatus(str, Enum): + """Enumeration for stage instance status.""" + + UP = "up" # Instance is ready and available + DOWN = "down" # Instance is shutdown gracefully + ERROR = "error" # Instance encountered an error or timeout + + +@dataclass +class InstanceEvent: + """Wire payload from OmniCoordClientForStage to OmniCoordinator. + + Schema for Stage β†’ Coordinator events over ZMQ: + input_addr, output_addr, stage_id, status, queue_length, event_type. + """ + + input_addr: str # Stage instance input ZMQ address (e.g., "tcp://host:port") + output_addr: str # Stage instance output ZMQ address (e.g., "tcp://host:port") + stage_id: int # Stage ID + event_type: str # "update" | "heartbeat" + status: StageStatus # Current status + queue_length: int # Current queue length + + +@dataclass +class InstanceInfo: + """Metadata for a single stage instance. + + This type is stored in OmniCoordinator's internal registry and is also + published to hubs via :class:`InstanceList`. + """ + + input_addr: str # Stage instance input ZMQ address (e.g., "tcp://host:port") + output_addr: str # Stage instance output ZMQ address (e.g., "tcp://host:port") + stage_id: int # Stage ID of this instance + status: StageStatus # Current status of the instance + queue_length: int # Current queue length of this instance + last_heartbeat: float # Timestamp of the last heartbeat received (seconds) + registered_at: float # Timestamp when the instance was registered (seconds) + + +@dataclass +class InstanceList: + """Container for instance list updates. + + OmniCoordinator publishes an :class:`InstanceList` whenever its view of + active instances changes. OmniCoordClientForHub caches the latest value + and exposes it to AsyncOmni and the load balancer. + """ + + instances: list[InstanceInfo] + timestamp: float # Time when the list was last updated (seconds) diff --git a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_hub.py b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_hub.py new file mode 100644 index 000000000..9081e4591 --- /dev/null +++ b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_hub.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import logging +import threading +from time import sleep, time +from typing import Any + +import zmq + +from .messages import InstanceInfo, InstanceList, StageStatus + +logger = logging.getLogger(__name__) + + +class OmniCoordClientForHub: + """Client for AsyncOmni side to receive instance list updates. + + This client maintains a SUB socket connected to OmniCoordinator's PUB + endpoint and caches the latest :class:`InstanceList` in memory for use by + the load balancer and routing logic. + """ + + def __init__(self, coord_zmq_addr: str) -> None: + """Initialize client and start receive thread (socket created in thread).""" + self._coord_zmq_addr = coord_zmq_addr + + self._ctx = zmq.Context() + self._lock = threading.Lock() + self._instance_list: InstanceList | None = None + self._closed = False + self._stop_event = threading.Event() + self._init_done = threading.Event() + self._init_error: list[BaseException] = [] + + self._thread = threading.Thread(target=self._recv_loop, daemon=True) + self._thread.start() + + self._init_done.wait(timeout=5.0) + if self._init_error: + raise RuntimeError(f"Failed to connect to coordinator at {self._coord_zmq_addr}") from self._init_error[0] + + def _decode_instance_list(self, payload: dict[str, Any]) -> InstanceList: + """Convert a JSON-decoded dict into an :class:`InstanceList`.""" + instances_payload = payload.get("instances", []) + instances: list[InstanceInfo] = [] + for inst in instances_payload: + instances.append( + InstanceInfo( + input_addr=inst["input_addr"], + output_addr=inst["output_addr"], + stage_id=int(inst["stage_id"]), + status=StageStatus(inst["status"]), + queue_length=int(inst["queue_length"]), + last_heartbeat=float(inst["last_heartbeat"]), + registered_at=float(inst["registered_at"]), + ) + ) + + timestamp = float(payload.get("timestamp", time())) + return InstanceList(instances=instances, timestamp=timestamp) + + def _recv_loop(self) -> None: + """Background loop that receives and caches instance lists.""" + sub = None + try: + sub = self._ctx.socket(zmq.SUB) + sub.setsockopt(zmq.SUBSCRIBE, b"") + sub.setsockopt(zmq.RCVTIMEO, 100) # 100ms timeout, avoids busy-wait + sub.connect(self._coord_zmq_addr) + except (zmq.ZMQError, OSError) as e: + self._init_error.append(e) + sub = None + finally: + self._init_done.set() + + try: + while not self._stop_event.is_set(): + # (Re)create and connect SUB socket if needed. + if sub is None: + try: + sub = self._ctx.socket(zmq.SUB) + sub.setsockopt(zmq.SUBSCRIBE, b"") + sub.setsockopt(zmq.RCVTIMEO, 100) # 100ms timeout, avoids busy-wait + sub.connect(self._coord_zmq_addr) + except (zmq.ZMQError, OSError) as e: + logger.error( + "Hub client failed to connect to coordinator at %s, will retry", + self._coord_zmq_addr, + exc_info=e, + ) + if sub is not None: + try: + sub.close() + except zmq.ZMQError: + pass + sub = None + sleep(1.0) + continue + + try: + data = sub.recv() + except zmq.Again: + continue + except zmq.ZMQError as e: + logger.error("Hub client recv failed, will reconnect", exc_info=e) + try: + sub.close() + except zmq.ZMQError: + pass + sub = None + sleep(1.0) + continue + + try: + payload = json.loads(data.decode("utf-8")) + inst_list = self._decode_instance_list(payload) + with self._lock: + self._instance_list = inst_list + except ( + json.JSONDecodeError, + KeyError, + ValueError, + TypeError, + AttributeError, + ) as e: + logger.warning("Invalid instance list message, skipping: %s", e) + finally: + try: + if sub is not None: + sub.close() + except zmq.ZMQError: + pass + try: + self._ctx.term() + except zmq.ZMQError: + pass + + def get_instance_list(self) -> InstanceList: + """Return the latest cached :class:`InstanceList`. + + If no update has been received yet, returns an empty list with + ``timestamp=0.0``. + """ + with self._lock: + if self._instance_list is None: + return InstanceList(instances=[], timestamp=0.0) + return self._instance_list + + def get_instances_for_stage(self, stage_id: int) -> InstanceList: + """Return instances filtered by ``stage_id``.""" + base = self.get_instance_list() + filtered = [inst for inst in base.instances if inst.stage_id == stage_id] + return InstanceList(instances=filtered, timestamp=base.timestamp) + + def close(self) -> None: + """Close the SUB socket and stop the background thread.""" + if self._closed: + raise RuntimeError("Client already closed") + + self._closed = True + self._stop_event.set() + self._thread.join(timeout=1.0) diff --git a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py new file mode 100644 index 000000000..cd5c357bb --- /dev/null +++ b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import logging +import threading +import time +from dataclasses import asdict + +import zmq + +from .messages import InstanceEvent, StageStatus + +logger = logging.getLogger(__name__) + + +class OmniCoordClientForStage: + """Client used by stage instances to send events to OmniCoordinator. + + This client maintains a DEALER socket connected to OmniCoordinator's + ROUTER endpoint and sends JSON-encoded events describing instance status. + """ + + def __init__( + self, + coord_zmq_addr: str, + input_addr: str, + output_addr: str, + stage_id: int, + ) -> None: + """Initialize client and send initial registration / status-up event.""" + self._coord_zmq_addr = coord_zmq_addr + self._input_addr = input_addr + self._output_addr = output_addr + self._stage_id = stage_id + + self._ctx = zmq.Context() + self._socket = self._ctx.socket(zmq.DEALER) + try: + self._socket.connect(self._coord_zmq_addr) + except zmq.ZMQError as e: + self._socket.close() + raise RuntimeError(f"Failed to connect to coordinator at {self._coord_zmq_addr}: {e}") from e + + self._status = StageStatus.UP + self._queue_length = 0 + self._closed = False + self._heartbeat_interval = 5.0 + self._stop_event = threading.Event() + self._send_lock = threading.Lock() + + self._send_event("update") + + self._heartbeat_thread = threading.Thread( + target=self._heartbeat_loop, + daemon=True, + ) + self._heartbeat_thread.start() + + def _reconnect(self) -> bool: + """Best-effort reconnect with up to ``max_retries`` attempts. + + Each attempt closes the current socket/context, sleeps 5 seconds, + then creates a new DEALER socket and reconnects to the coordinator. + Caller must hold ``_send_lock``. + Returns True on success, False if all attempts fail. + """ + while not self._stop_event.is_set() and not self._closed: + try: + self._socket.close(0) + except zmq.ZMQError: + pass + try: + self._ctx.term() + except zmq.ZMQError: + pass + + time.sleep(5.0) + + try: + self._ctx = zmq.Context() + self._socket = self._ctx.socket(zmq.DEALER) + self._socket.connect(self._coord_zmq_addr) + return True + except zmq.ZMQError as e: + logger.error( + "Stage client reconnect failed, will retry in 5s (coord=%s)", + self._coord_zmq_addr, + exc_info=e, + ) + continue + return False + + def _send_event(self, event_type: str) -> None: + """Send an InstanceEvent to OmniCoordinator. + + Wire format: input_addr, output_addr, stage_id, status, queue_length, event_type. + For "update": includes status and queue_length from instance state. + For "heartbeat": status and queue_length are null. + + On send failure (ZMQError / RuntimeError), attempts to reconnect up + to 3 times (5s sleep each) and retries the send once after a + successful reconnect. Raises if reconnect or the retry send fails. + """ + if self._closed: + raise RuntimeError("Client already closed") + + event = InstanceEvent( + input_addr=self._input_addr, + output_addr=self._output_addr, + stage_id=self._stage_id, + event_type=event_type, + status=self._status, + queue_length=self._queue_length, + ) + data = json.dumps(asdict(event)).encode("utf-8") + + with self._send_lock: + try: + self._socket.send(data, flags=zmq.NOBLOCK) + return + except zmq.Again: + logger.debug("Send buffer full, dropping event") + return + except (RuntimeError, zmq.ZMQError) as e: + # First send failed; try reconnecting a few times. + if not self._reconnect: + logger.error("Failed to send event and reconnect to coordinator", exc_info=e) + raise + + # Reconnected successfully; try sending once more. + try: + self._socket.send(data, flags=zmq.NOBLOCK) + except zmq.Again: + logger.debug("Send buffer full after reconnect, dropping event") + except (RuntimeError, zmq.ZMQError) as e2: + logger.error("Failed to send event after successful reconnect", exc_info=e2) + raise + + def update_info( + self, + status: StageStatus | None = None, + queue_length: int | None = None, + ) -> None: + """Update instance information and notify OmniCoordinator. + + At least one of ``status`` or ``queue_length`` must be provided. + """ + if status is None and queue_length is None: + raise ValueError("At least one of status or queue_length must be provided") + + if status is not None: + self._status = status + if queue_length is not None: + self._queue_length = queue_length + + self._send_event("update") + + def _heartbeat_loop(self) -> None: + """Periodically send heartbeat events while the client is alive.""" + while not self._stop_event.wait(timeout=self._heartbeat_interval): + if self._closed: + break + + try: + self._send_event("heartbeat") + except (RuntimeError, zmq.ZMQError): + break + + def close(self) -> None: + """Send a final down event and close the underlying socket.""" + if self._closed: + raise RuntimeError("Client already closed") + + # Stop heartbeat thread first to avoid concurrent sends during shutdown. + self._stop_event.set() + if hasattr(self, "_heartbeat_thread"): + self._heartbeat_thread.join(timeout=1.0) + + # Mark status as DOWN and send one last update. + self._status = StageStatus.DOWN + try: + self._send_event("update") + except zmq.ZMQError: + pass # Socket may already be broken, proceed with close + + # Close DEALER socket and terminate this client's context. + self._socket.close(0) + try: + self._ctx.term() + except zmq.ZMQError: + pass + self._closed = True diff --git a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py new file mode 100644 index 000000000..7ff608f2f --- /dev/null +++ b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py @@ -0,0 +1,330 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import json +import logging +import threading +from dataclasses import asdict +from time import time +from typing import Any + +import zmq + +from .messages import InstanceEvent, InstanceInfo, InstanceList, StageStatus + +logger = logging.getLogger(__name__) + + +class OmniCoordinator: + """Coordinator for stage instances and hub clients. + + This service receives instance events from :class:`OmniCoordClientForStage` + via a ZMQ ROUTER socket and publishes active instance lists to + :class:`OmniCoordClientForHub` via a PUB socket. + + The coordinator maintains an in-memory registry of all known instances, + including their status, queue length, and heartbeat timestamps. A + background thread periodically checks for heartbeat timeouts and marks + unhealthy instances as ``StageStatus.ERROR``. + """ + + def __init__( + self, + router_zmq_addr: str, + pub_zmq_addr: str, + heartbeat_timeout: float = 30.0, + ) -> None: + """Initialize coordinator and start background service loops. + + Args: + router_zmq_addr: ZMQ address to bind the ROUTER socket. + pub_zmq_addr: ZMQ address to bind the PUB socket. + heartbeat_timeout: Seconds before an instance is considered + unhealthy if no heartbeat / update is received. + """ + self._router_zmq_addr = router_zmq_addr + self._pub_zmq_addr = pub_zmq_addr + self._heartbeat_timeout = heartbeat_timeout + + # Dedicated ZMQ context for this coordinator instance. + self._ctx = zmq.Context() + self._router = self._ctx.socket(zmq.ROUTER) + self._router.bind(self._router_zmq_addr) + + self._pub = self._ctx.socket(zmq.PUB) + self._pub.bind(self._pub_zmq_addr) + + self._instances: dict[str, InstanceInfo] = {} + self._lock = threading.Lock() + self._pub_lock = threading.Lock() + + self._publish_min_interval: float = 0.1 # seconds + self._pending_broadcast: bool = False + + self._running = True + self._closed = False + self._stop_event = threading.Event() + + self._router.setsockopt(zmq.RCVTIMEO, 100) + + self._recv_thread = threading.Thread(target=self._recv_loop, daemon=True) + self._recv_thread.start() + + self._periodic_thread = threading.Thread(target=self._periodic_loop, daemon=True) + self._periodic_thread.start() + + def get_active_instances(self) -> InstanceList: + """Return an :class:`InstanceList` of active (UP) instances only.""" + with self._lock: + active = [inst for inst in self._instances.values() if inst.status == StageStatus.UP] + return InstanceList(instances=active, timestamp=time()) + + def add_new_instance(self, event: InstanceEvent) -> None: + """Add a new instance based on an incoming event.""" + with self._lock: + self._add_new_instance_locked(event) + self.publish_instance_list_update() + + def update_instance_info(self, event: InstanceEvent) -> None: + """Update an existing instance based on an incoming event.""" + with self._lock: + self._update_instance_info_locked(event) + self.publish_instance_list_update() + + def remove_instance(self, event: InstanceEvent) -> None: + """Mark an instance as removed / down based on an incoming event. + + This marks the instance's status as DOWN or ERROR (depending on the + event) but keeps it in the internal registry. It is removed from the + *active* instance list published to hubs. + """ + with self._lock: + self._remove_instance_locked(event) + self.publish_instance_list_update() + + def publish_instance_list_update(self) -> None: + """Publish the current active instance list to all subscribers.""" + active_list = self.get_active_instances() + payload = asdict(active_list) + data = json.dumps(payload).encode("utf-8") + + with self._pub_lock: + try: + # PUB socket is best-effort; drop update if not ready. + self._pub.send(data, flags=zmq.NOBLOCK) + except (zmq.Again, zmq.ZMQError): + # Silently ignore send failures; next update will catch up. + return + + def _schedule_broadcast(self, force: bool) -> None: + """Schedule a broadcast, optionally bypassing throttling. + + When ``force`` is True, publish immediately. Otherwise, mark a pending + broadcast that will be flushed by the periodic loop at most once per + ``_publish_min_interval``. + """ + if force: + self.publish_instance_list_update() + else: + self._pending_broadcast = True + + def _mark_instance_error_locked(self, info: InstanceInfo) -> None: + """Mark instance as ERROR (e.g. after heartbeat timeout).""" + info.status = StageStatus.ERROR + + def _check_heartbeat_timeouts(self) -> None: + """Mark instances as ERROR if their heartbeat has timed out.""" + now = time() + timed_out = False + gc_ttl = 600.0 # 10 minutes + + with self._lock: + to_delete: list[str] = [] + + for input_addr, info in self._instances.items(): + if info.status == StageStatus.UP and now - info.last_heartbeat > self._heartbeat_timeout: + self._mark_instance_error_locked(info) + timed_out = True + elif info.status in (StageStatus.DOWN, StageStatus.ERROR) and now - info.last_heartbeat > gc_ttl: + to_delete.append(input_addr) + + for input_addr in to_delete: + del self._instances[input_addr] + if timed_out: + # Instance liveness changed; force immediate broadcast. + self._schedule_broadcast(force=True) + + def close(self) -> None: + """Shut down background threads and close all ZMQ sockets.""" + if self._closed: + raise RuntimeError("Coordinator already closed") + + self._closed = True + self._running = False + self._stop_event.set() + + # Wait for threads to exit before closing sockets. + for thread in (self._recv_thread, self._periodic_thread): + thread.join(timeout=1.0) + + try: + self._router.close(0) + except zmq.ZMQError: + pass + + try: + self._pub.close(0) + except zmq.ZMQError: + pass + + try: + self._ctx.term() + except zmq.ZMQError: + pass + + def _parse_instance_event(self, data: dict[str, Any]) -> InstanceEvent | None: + """Parse wire payload dict into InstanceEvent. Returns None if invalid.""" + try: + return InstanceEvent( + input_addr=str(data["input_addr"]), + output_addr=str(data["output_addr"]), + stage_id=int(data["stage_id"]), + event_type=str(data["event_type"]), + status=StageStatus(data.get("status")), + queue_length=data.get("queue_length"), + ) + except (KeyError, ValueError, TypeError): + return None + + def _recv_loop(self) -> None: + """Background loop that receives and processes instance events.""" + while self._running: + try: + frames = self._router.recv_multipart() + except zmq.Again: + # RCVTIMEO expired, loop to recheck _running. + continue + except zmq.ZMQError: + # Socket likely closed or context terminated. + break + + if not frames: + continue + + payload = frames[-1] + try: + data = json.loads(payload.decode("utf-8")) + event = self._parse_instance_event(data) + except json.JSONDecodeError as e: + logger.warning("Invalid JSON in instance event, dropping: %s", e) + continue + if event is None: + logger.warning("Malformed instance event, dropping") + continue + + self._handle_event(event) + + def _periodic_loop(self) -> None: + """Periodic loop to check heartbeat timeouts and flush broadcasts. + + Heartbeat timeouts are checked on their original cadence, while + queue_length / non-liveness updates are coalesced and flushed at + most once per ``_publish_min_interval``. + """ + heartbeat_interval = max(1.0, min(self._heartbeat_timeout / 2.0, 5.0)) + loop_interval = self._publish_min_interval + + last_heartbeat_check = 0.0 + while self._running: + now = time() + + if now - last_heartbeat_check >= heartbeat_interval: + self._check_heartbeat_timeouts() + last_heartbeat_check = now + + if self._pending_broadcast: + self.publish_instance_list_update() + self._pending_broadcast = False + + if self._stop_event.wait(timeout=loop_interval): + break + + def _handle_event(self, event: InstanceEvent) -> None: + """Dispatch an incoming event to the appropriate handler.""" + try: + input_addr = event.input_addr + + # Heartbeat: only update last_heartbeat; if previously ERROR, + # promote back to UP and broadcast once. + if event.event_type == "heartbeat": + promote = False + with self._lock: + info = self._instances.get(input_addr) + if info is not None: + info.last_heartbeat = time() + if info.status == StageStatus.ERROR: + info.status = StageStatus.UP + promote = True + if promote: + self._schedule_broadcast(force=True) + return + + # Check-and-act under single lock to avoid TOCTOU race (duplicate + # registration when concurrent events arrive for the same instance). + with self._lock: + force_broadcast = False + if input_addr not in self._instances: + self._add_new_instance_locked(event) + force_broadcast = True + else: + if event.status == StageStatus.DOWN: + self._remove_instance_locked(event) + force_broadcast = True + else: + self._update_instance_info_locked(event) + + # New instances / DOWN events are broadcast immediately; other + # updates (e.g. queue_length changes) are throttled via the + # periodic loop. + self._schedule_broadcast(force=force_broadcast) + except (KeyError, ValueError, TypeError) as e: + logger.warning("Dropping malformed event: %s", e) + + def _add_new_instance_locked(self, event: InstanceEvent) -> None: + input_addr = event.input_addr + if not input_addr: + raise KeyError("input_addr required") + stage_id = event.stage_id + if stage_id < 0: + raise KeyError("stage_id required and must be non-negative") + + now = time() + info = InstanceInfo( + input_addr=input_addr, + output_addr=event.output_addr, + stage_id=stage_id, + status=event.status, + queue_length=event.queue_length, + last_heartbeat=now, + registered_at=now, + ) + self._instances[input_addr] = info + + def _update_instance_info_locked(self, event: InstanceEvent) -> None: + input_addr = event.input_addr + info = self._instances[input_addr] + + if event.status is not None: + info.status = event.status + + if event.queue_length is not None: + info.queue_length = event.queue_length + + def _remove_instance_locked(self, event: InstanceEvent) -> None: + input_addr = event.input_addr + info = self._instances.get(input_addr) + if info is None: + return + + info.status = StageStatus.DOWN diff --git a/vllm_omni/distributed/ray_utils/utils.py b/vllm_omni/distributed/ray_utils/utils.py index 07513b660..519460256 100644 --- a/vllm_omni/distributed/ray_utils/utils.py +++ b/vllm_omni/distributed/ray_utils/utils.py @@ -176,6 +176,26 @@ def start_ray_actor( runtime_env={"env_vars": {"PYTHONPATH": os.environ.get("PYTHONPATH", "")}, "CUDA_LAUNCH_BLOCKING": "1"}, ).remote() - worker_actor.run.remote(worker_entry_fn, *args, **kwargs) + task_ref = worker_actor.run.remote(worker_entry_fn, *args, **kwargs) - return worker_actor + return worker_actor, task_ref + + +def is_ray_task_alive(task_ref: Any, **kwargs): + """Checks ray task status. Returns FALSE if ray task has exited for any reason.""" + if not RAY_AVAILABLE: + raise ImportError("ray is required to query ray tasks") + + ready, _ = ray.wait([task_ref], **kwargs) + return not bool(ready) + + +def get_ray_task_error(task_ref: Any, **kwargs) -> Exception | None: + """Gets ray task. Returns RayTaskError if ray instance exited with any error, else None.""" + if not RAY_AVAILABLE: + raise ImportError("ray is required to query ray tasks") + + try: + ray.get(task_ref, **kwargs) + except Exception as e: + return e diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index c4f719703..d733e7bc4 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -71,7 +71,7 @@ class OmniEngineArgs(EngineArgs): stage_id: int = 0 model_stage: str = "thinker" - model_arch: str = "Qwen2_5OmniForConditionalGeneration" + model_arch: str | None = None engine_output_type: str | None = None hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None @@ -233,7 +233,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): stage_id: int = 0 model_stage: str = "thinker" - model_arch: str = "Qwen2_5OmniForConditionalGeneration" + model_arch: str | None = None engine_output_type: str | None = None hf_config_name: str | None = None custom_process_next_stage_input_func: str | None = None diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 3c3cb4cd6..1ba95124a 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -140,6 +140,15 @@ class AsyncOmni(OmniBase): getattr(self, "_inline_engine", None), ) + async def get_supported_tasks(self) -> set[str]: + """Return supported tasks based on stage output modalities and capabilities.""" + tasks: set[str] = set() + if "text" in self.output_modalities or any(stage.is_comprehension for stage in self.stage_list): + tasks.add("generate") + if "audio" in self.output_modalities: + tasks.add("speech") + return tasks + def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]: """Create default diffusion stage configuration.""" # TODO: here is different from the Omni class. We should merge the two in the future. @@ -831,16 +840,6 @@ class AsyncOmni(OmniBase): stage.submit(abort_task) return None - async def get_supported_tasks(self) -> set[str]: - tasks: set[str] = set() - has_comprehension = any(stage.is_comprehension for stage in self.stage_list) - if has_comprehension: - tasks.add("generate") - for stage in self.stage_list: - if getattr(stage, "final_output_type", None) == "audio": - tasks.add("speech") - return tasks if tasks else {"generate"} - async def get_vllm_config(self) -> VllmConfig: for stage in self.stage_list: if stage.is_comprehension: diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index c414b972b..488b986d8 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -15,13 +15,14 @@ import huggingface_hub import msgspec.msgpack import torch import zmq -from omegaconf import OmegaConf from tqdm.auto import tqdm from vllm import SamplingParams from vllm.logger import init_logger from vllm.utils.network_utils import make_zmq_socket from vllm.v1.utils import get_engine_client_zmq_addr +from vllm_omni.config.stage_config import StageConfigFactory +from vllm_omni.config.yaml_util import create_config from vllm_omni.distributed.omni_connectors import ( get_stage_connector_config, initialize_orchestrator_connectors, @@ -224,46 +225,35 @@ class OmniBase: cache_config = self._get_default_cache_config(cache_backend) return cache_config - def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]: - """Create default diffusion stage configuration.""" - # We temporally create a default config for diffusion stage. - # In the future, we should merge the default config with the user-provided config. - # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. + def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> list[dict[str, Any]]: + """Create default diffusion stage configuration. + + Uses StageConfigFactory for typed configuration creation while + maintaining backward compatibility with the legacy format. + + Args: + kwargs: Engine arguments from CLI/API. + + Returns: + List containing a single OmegaConf config for the diffusion stage. + """ + # Normalize dtype if "dtype" in kwargs and not isinstance(kwargs["dtype"], str): if not isinstance(kwargs["dtype"], torch.dtype): raise TypeError(f"Provided dtype must be a string or torch.dtype, got {type(kwargs['dtype']).__name__}") kwargs["dtype"] = str(kwargs["dtype"]).removeprefix("torch.") + # Normalize cache config before passing to factory cache_backend = kwargs.get("cache_backend", "none") cache_config = self._normalize_cache_config(cache_backend, kwargs.get("cache_config", None)) - # TODO: hack, calculate devices based on parallel config. - devices = "0" - if "parallel_config" in kwargs: - num_devices = kwargs["parallel_config"].world_size - for i in range(1, num_devices): - devices += f",{i}" - default_stage_cfg = [ - { - "stage_id": 0, - "stage_type": "diffusion", - "runtime": { - "process": True, - "devices": devices, - "max_batch_size": 1, - }, - "engine_args": OmegaConf.create( - { - **kwargs, - "cache_backend": cache_backend, - "cache_config": cache_config, - } - ), - "final_output": True, - "final_output_type": "image", - } - ] - default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" - return default_stage_cfg + + # Update kwargs with normalized values + kwargs_copy = dict(kwargs) + kwargs_copy["cache_backend"] = cache_backend + kwargs_copy["cache_config"] = cache_config + + # Use the factory to create default diffusion config + return StageConfigFactory.create_default_diffusion(kwargs_copy) def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]: """Resolve stage configs and inject defaults shared by orchestrator/headless.""" @@ -290,7 +280,7 @@ class OmniBase: if getattr(cfg, "stage_type", None) != "diffusion": continue if not hasattr(cfg, "engine_args") or cfg.engine_args is None: - cfg.engine_args = OmegaConf.create({}) + cfg.engine_args = create_config({}) if kwargs.get("lora_path") is not None: if not hasattr(cfg.engine_args, "lora_path") or cfg.engine_args.lora_path is None: cfg.engine_args.lora_path = kwargs["lora_path"] diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 27f470b4b..e06e50c83 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -25,8 +25,6 @@ from vllm_omni.engine.input_processor import OmniInputProcessor, reinject_omni_f from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.entrypoints.utils import ( filter_dataclass_kwargs, - load_stage_configs_from_model, - load_stage_configs_from_yaml, resolve_model_config_path, ) @@ -43,8 +41,6 @@ class OmniLLM(LLM): Args: model: Model name or path to load - stage_configs_path: Optional path to YAML file containing stage - configurations. If None, configurations are loaded from the model. log_stats: Whether to enable statistics logging compilation_config: Optional compilation configuration. Can be an integer (compilation level), dict, or CompilationConfig instance. @@ -68,7 +64,6 @@ class OmniLLM(LLM): def __init__( self, model: str, - stage_configs_path: str | None = None, log_stats: bool = False, compilation_config: int | dict[str, Any] | CompilationConfig | None = None, hf_overrides: dict[str, Any] | None = None, @@ -86,13 +81,8 @@ class OmniLLM(LLM): self.batch_timeout = batch_timeout self.log_stats: bool = bool(log_stats) - # Load stage configurations - if stage_configs_path is None: - self.config_path = resolve_model_config_path(model) - self.stage_configs = load_stage_configs_from_model(model) - else: - self.config_path = stage_configs_path - self.stage_configs = load_stage_configs_from_yaml(stage_configs_path) + # Resolve model config path for connectors + self.config_path = resolve_model_config_path(model) # Initialize connectors self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors( @@ -154,6 +144,7 @@ class OmniLLM(LLM): compilation_config=compilation_config_instance, structured_outputs_config=structured_outputs_instance, omni_kv_config=omni_kv_config, + hf_overrides=hf_overrides or {}, **filter_dataclass_kwargs(OmniEngineArgs, kwargs), ) diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 899d8ddff..b40311d19 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -36,7 +36,12 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.distributed.omni_connectors import build_stage_connectors from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase -from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_actor +from vllm_omni.distributed.ray_utils.utils import ( + get_ray_task_error, + is_ray_task_alive, + kill_ray_actor, + start_ray_actor, +) from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs, OmniEngineArgs from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM @@ -264,7 +269,10 @@ class OmniStage: self.engine_args = stage_config.engine_args self.model_stage = stage_config.engine_args.model_stage self.requires_multimodal_data = getattr(stage_config.runtime, "requires_multimodal_data", False) - self.engine_input_source = getattr(stage_config, "engine_input_source", []) + # Support both 'input_sources' (new format) and 'engine_input_source' (legacy) + self.engine_input_source = ( + getattr(stage_config, "input_sources", None) or getattr(stage_config, "engine_input_source", []) or [] + ) self.engine_output_type = getattr(stage_config.engine_args, "engine_output_type", None) self.engine_outputs = None self.is_comprehension = getattr(stage_config, "is_comprehension", False) @@ -303,6 +311,8 @@ class OmniStage: self._in_q: mp.queues.Queue | ZmqQueue | str | None = None self._out_q: mp.queues.Queue | ZmqQueue | str | None = None self._proc: mp.Process | None = None + self._ray_actor: Any | None = None + self._ray_task_ref: Any | None = None self._shm_threshold_bytes: int = 65536 self._stage_init_timeout: int = stage_init_timeout # Callback used by the orchestrator's output handler to stash @@ -544,7 +554,7 @@ class OmniStage: os.environ["VLLM_LOGGING_PREFIX"] = new_env if worker_backend == "ray": if is_async: - self._ray_actor = start_ray_actor( + self._ray_actor, self._ray_task_ref = start_ray_actor( _stage_worker_async_entry, ray_placement_group, self.stage_id, @@ -556,7 +566,7 @@ class OmniStage: stage_init_timeout=self._stage_init_timeout, ) else: - self._ray_actor = start_ray_actor( + self._ray_actor, self._ray_task_ref = start_ray_actor( _stage_worker, ray_placement_group, self.stage_id, @@ -629,9 +639,10 @@ class OmniStage: pass self._out_q = None - if hasattr(self, "_ray_actor") and self._ray_actor: + if self._ray_actor is not None: kill_ray_actor(self._ray_actor) self._ray_actor = None + self._ray_task_ref = None elif self._proc is not None: try: self._proc.join(timeout=5) @@ -695,8 +706,17 @@ class OmniStage: return None try: return self._out_q.get_nowait() - except Exception: - return None + except queue.Empty: + pass + except Exception as e: + logger.error("Unexpected error when collecting OmniStage output queue:", exc_info=e) + self.stop_stage_worker() + raise + if self._proc is not None and not self._proc.is_alive(): + raise RuntimeError(f"OmniStage Worker process died unexpectedly with exit code {self._proc.exitcode}") + if self._ray_task_ref is not None and not is_ray_task_alive(self._ray_task_ref, timeout=0): + e = get_ray_task_error(self._ray_task_ref, timeout=0) + raise RuntimeError("OmniStage Ray actor died unexpectedly") from e def process_engine_inputs( self, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 6404e56d9..82770e02d 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -19,10 +19,9 @@ from typing import Annotated, Any, cast import httpx import vllm.envs as envs -from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile, WebSocket from fastapi.responses import JSONResponse, StreamingResponse from PIL import Image -from pydantic import BaseModel, Field from starlette.datastructures import State from starlette.routing import Route from vllm import SamplingParams @@ -99,6 +98,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( ) from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams, OmniTextPrompt from vllm_omni.lora.request import LoRARequest @@ -106,30 +106,6 @@ from vllm_omni.lora.utils import stable_lora_int_id logger = init_logger(__name__) router = APIRouter() -profiler_router = APIRouter() - - -def _should_enable_profiler_endpoints(args: Namespace) -> bool: - # Check upstream vLLM's profiler_config - profiler_config = getattr(args, "profiler_config", None) - if profiler_config is not None: - # profiler_config exists, check if profiler is set - profiler = getattr(profiler_config, "profiler", None) - if profiler is not None: - return True - - # TODO: remove this env after refactoring torch profiler to CLI args - env_value = os.environ.get("VLLM_TORCH_PROFILER_DIR") - return env_value is not None - - -class ProfileRequest(BaseModel): - """Request model for profiling endpoints.""" - - stages: list[int] | None = Field( - default=None, - description="List of stage IDs to profile. If None, profiles all stages.", - ) def _remove_route_from_router( @@ -178,6 +154,10 @@ class _DiffusionServingModels: def __init__(self, base_model_paths: list[BaseModelPath]) -> None: self._base_model_paths = base_model_paths + @property + def base_model_paths(self) -> list[BaseModelPath]: + return self._base_model_paths + async def show_available_models(self) -> ModelList: return ModelList( data=[ @@ -241,7 +221,10 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, else: supported_tasks = ("generate",) if not supported_tasks: - supported_tasks = ("generate",) + # Only default to "generate" when get_supported_tasks is not implemented; + # TTS-only models intentionally return an empty set. + if not hasattr(engine_client, "get_supported_tasks"): + supported_tasks = ("generate",) # OMNI: Pass supported_tasks to build_app (required by upstream vLLM) app = build_openai_app(args, supported_tasks) @@ -252,11 +235,6 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, await omni_init_app_state(engine_client, app.state, args) - # Conditionally register profiler endpoints based on config or env var - if _should_enable_profiler_endpoints(args): - logger.warning("Profiler endpoints are enabled. This should ONLY be used for local development!") - app.include_router(profiler_router) - vllm_config = await engine_client.get_vllm_config() # Check if pure diffusion mode (vllm_config will be None) @@ -287,7 +265,6 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, ssl_certfile=args.ssl_certfile, ssl_ca_certs=args.ssl_ca_certs, ssl_cert_reqs=args.ssl_cert_reqs, - ssl_ciphers=args.ssl_ciphers, h11_max_incomplete_event_size=args.h11_max_incomplete_event_size, h11_max_header_count=args.h11_max_header_count, **uvicorn_kwargs, @@ -731,6 +708,10 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, request_logger=request_logger ) + state.openai_streaming_speech = OmniStreamingSpeechHandler( + speech_service=state.openai_serving_speech, + ) + state.openai_serving_video = OmniOpenAIServingVideo( engine_client, model_name=served_model_names[0] if served_model_names else None, @@ -878,113 +859,29 @@ async def list_voices(raw_request: Request): if handler is None: return base(raw_request).create_error_response(message="The model does not support Speech API") - # Get all speakers (both model built-in and uploaded) speakers = sorted(handler.supported_speakers) if handler.supported_speakers else [] - - # Get uploaded speakers details - uploaded_speakers = [] - if hasattr(handler, "uploaded_speakers"): - for voice_name, info in handler.uploaded_speakers.items(): - uploaded_speakers.append( - { - "name": info.get("name", voice_name), - "consent": info.get("consent", ""), - "created_at": info.get("created_at", 0), - "file_size": info.get("file_size", 0), - "mime_type": info.get("mime_type", ""), - } - ) - - return JSONResponse(content={"voices": speakers, "uploaded_voices": uploaded_speakers}) + return JSONResponse(content={"voices": speakers}) -@router.post( - "/v1/audio/voices", - responses={ - HTTPStatus.OK.value: {"model": dict}, - HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, - }, -) -async def upload_voice( - raw_request: Request, - audio_sample: UploadFile = File(...), - consent: str = Form(...), - name: str = Form(...), -): - """Upload a new voice sample for voice cloning. +@router.websocket("/v1/audio/speech/stream") +async def streaming_speech(websocket: WebSocket): + """WebSocket endpoint for streaming text input TTS. - Uploads an audio file that can be used as a reference for voice cloning - in Base task TTS requests. The voice can then be referenced by name - in subsequent TTS requests. - - Args: - audio_sample: Audio file (max 10MB) - consent: Consent recording ID - name: Name for the new voice - raw_request: Raw FastAPI request - - Returns: - JSON response with voice information + Accepts text incrementally, splits at sentence boundaries, and + returns audio per sentence. See serving_speech_stream.py for protocol. """ - handler = Omnispeech(raw_request) + handler = getattr(websocket.app.state, "openai_streaming_speech", None) if handler is None: - return base(raw_request).create_error_response(message="The model does not support Speech API") - - try: - # Upload the voice - result = await handler.upload_voice(audio_sample, consent, name) - - return JSONResponse(content={"success": True, "voice": result}) - - except ValueError as e: - return base(raw_request).create_error_response(message=str(e)) - except Exception as e: - logger.exception(f"Failed to upload voice: {e}") - return base(raw_request).create_error_response(message=f"Failed to upload voice: {str(e)}") - - -@router.delete( - "/v1/audio/voices/{name}", - responses={ - HTTPStatus.OK.value: {"model": dict}, - HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, - HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, - }, -) -async def delete_voice(name: str, raw_request: Request): - """Delete an uploaded voice. - - Deletes the voice sample and associated metadata. Also removes any - cached voice clone prompts for this voice. - - Args: - name: Name of the voice to delete - raw_request: Raw FastAPI request - - Returns: - JSON response indicating success or failure - """ - handler = Omnispeech(raw_request) - if handler is None: - return base(raw_request).create_error_response(message="The model does not support Speech API") - - try: - # Delete the voice - success = await handler.delete_voice(name) - if not success: - return JSONResponse( - content={"success": False, "error": f"Voice '{name}' not found"}, - status_code=HTTPStatus.NOT_FOUND.value, - ) - - return JSONResponse(content={"success": True, "message": f"Voice '{name}' deleted successfully"}) - - except ValueError as e: - return base(raw_request).create_error_response(message=str(e)) - except Exception as e: - logger.exception(f"Failed to delete voice '{name}': {e}") - return base(raw_request).create_error_response(message=f"Failed to delete voice: {str(e)}") + await websocket.accept() + await websocket.send_json( + { + "type": "error", + "message": "Streaming speech is not available", + } + ) + await websocket.close() + return + await handler.handle_session(websocket) # Health and Model endpoints for diffusion mode @@ -1626,58 +1523,6 @@ def apply_stage_default_sampling_params( setattr(sampling_params, param_name, param_value) -@profiler_router.post("/start_profile") -async def start_profile(raw_request: Request, request: ProfileRequest | None = None): - """Start profiling for the engine. - - Args: - request: Optional request body with stages to profile. - - stages: List of stage IDs to profile. If None, profiles all stages. - - Example: - POST /start_profile - {"stages": [0, 1]} # Profile only stages 0 and 1 - """ - try: - stages = request.stages if request else None - logger.info("Starting profiler for stages: %s", stages if stages else "all") - engine_client = raw_request.app.state.engine_client - result = await engine_client.start_profile(stages=stages) - logger.info("Profiler started.") - return JSONResponse(content=result) - except Exception as e: - logger.exception("Failed to start profiler: %s", e) - raise HTTPException( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to start profiler: {str(e)}" - ) - - -@profiler_router.post("/stop_profile") -async def stop_profile(raw_request: Request, request: ProfileRequest | None = None): - """Stop profiling for the engine. - - Args: - request: Optional request body with stages to stop profiling. - - stages: List of stage IDs to stop profiling. If None, stops all stages. - - Example: - POST /stop_profile - {"stages": [0, 1]} # Stop profiling only stages 0 and 1 - """ - try: - stages = request.stages if request else None - logger.info("Stopping profiler for stages: %s", stages if stages else "all") - engine_client = raw_request.app.state.engine_client - result = await engine_client.stop_profile(stages=stages) - logger.info("Profiler stopped.") - return JSONResponse(content=result) - except Exception as e: - logger.exception("Failed to stop profiler: %s", e) - raise HTTPException( - status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=f"Failed to stop profiler: {str(e)}" - ) - - async def _run_video_generation( request: VideoGenerationRequest, raw_request: Request, diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 871787173..e57743226 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -71,9 +71,9 @@ class OpenAICreateSpeechRequest(BaseModel): @model_validator(mode="after") def validate_streaming_constraints(self) -> "OpenAICreateSpeechRequest": if self.stream: - if self.response_format != "pcm": + if self.response_format not in ("pcm", "wav"): raise ValueError( - "Streaming (stream=true) requires response_format='pcm'. " + "Streaming (stream=true) requires response_format not in ('pcm', 'wav'). " f"Got response_format='{self.response_format}'." ) if self.speed is None: @@ -100,3 +100,52 @@ class CreateAudio(BaseModel): class AudioResponse(BaseModel): audio_data: bytes | str media_type: str + + +class StreamingSpeechSessionConfig(BaseModel): + """Configuration sent as the first WebSocket message for streaming TTS.""" + + model: str | None = None + voice: str | None = None + task_type: Literal["CustomVoice", "VoiceDesign", "Base"] | None = None + language: str | None = None + instructions: str | None = None + response_format: Literal["wav", "pcm", "flac", "mp3", "aac", "opus"] = "wav" + speed: float | None = Field(default=1.0, ge=0.25, le=4.0) + max_new_tokens: int | None = Field(default=None, ge=1) + initial_codec_chunk_frames: int | None = Field( + default=None, + ge=0, + description="Initial chunk size for reduced TTFA. Overrides stage config for this session.", + ) + ref_audio: str | None = None + ref_text: str | None = None + x_vector_only_mode: bool | None = None + stream_audio: bool = Field( + default=False, + description=( + "If true, send raw PCM audio chunks progressively over WebSocket. " + "Requires response_format='pcm'. Speed adjustment is not supported when streaming." + ), + ) + split_granularity: Literal["sentence", "clause"] = Field( + default="sentence", + description=( + "Text splitting granularity: 'sentence' splits on .!?γ€‚οΌοΌŸ, " + "'clause' also splits on CJK commas ， and semicolons οΌ›." + ), + ) + + @model_validator(mode="after") + def validate_streaming_constraints(self) -> "StreamingSpeechSessionConfig": + if self.stream_audio: + if self.response_format != "pcm": + raise ValueError( + "WebSocket streaming audio (stream_audio=true) requires response_format='pcm'. " + f"Got response_format='{self.response_format}'." + ) + if self.speed is None: + self.speed = 1.0 + elif self.speed != 1.0: + raise ValueError("Speed adjustment is not supported when stream_audio=true. Set speed=1.0 or omit it.") + return self diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 7bcf75ace..3b5d67542 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -4,6 +4,7 @@ import json import math import os import re +import struct import time from pathlib import Path from typing import Any @@ -20,6 +21,7 @@ from vllm.utils import random_uuid from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin from vllm_omni.entrypoints.openai.metadata_manager import MetadataManager from vllm_omni.entrypoints.openai.protocol.audio import ( + AudioResponse, CreateAudio, OpenAICreateSpeechRequest, ) @@ -27,8 +29,8 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) -# TTS Configuration (currently supports Qwen3-TTS) -_TTS_MODEL_STAGES: set[str] = {"qwen3_tts"} +# TTS Configuration +_TTS_MODEL_STAGES: set[str] = {"qwen3_tts", "fish_speech_slow_ar"} _TTS_LANGUAGES: set[str] = { "Auto", "Chinese", @@ -47,6 +49,47 @@ _TTS_MAX_NEW_TOKENS_MIN = 1 _TTS_MAX_NEW_TOKENS_MAX = 4096 +def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample: int = 16) -> bytes: + """Create a WAV header with placeholder size values for streaming. + + Uses 0xFFFFFFFF as placeholder for data size fields, which is accepted + by most audio clients and matches OpenAI's streaming WAV implementation. + + Args: + sample_rate: Audio sample rate in Hz + num_channels: Number of audio channels (1 for mono, 2 for stereo) + bits_per_sample: Bits per sample (typically 16) + + Returns: + 44-byte WAV header as bytes + """ + byte_rate = sample_rate * num_channels * bits_per_sample // 8 + block_align = num_channels * bits_per_sample // 8 + + # Use 0xFFFFFFFF as placeholder for unknown size (streaming) + placeholder_size = 0xFFFFFFFF + + # ref https://docs.fileformat.com/audio/wav/ + header = struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", # ChunkID + placeholder_size, # ChunkSize (placeholder) + b"WAVE", # Format + b"fmt ", # Subchunk1ID + 16, # Subchunk1Size (16 for PCM) + 1, # AudioFormat (1 for PCM) + num_channels, # NumChannels + sample_rate, # SampleRate + byte_rate, # ByteRate + block_align, # BlockAlign + bits_per_sample, # BitsPerSample + b"data", # Subchunk2ID + placeholder_size, # Subchunk2Size (placeholder) + ) + + return header + + def _sanitize_filename(filename: str) -> str: """Sanitize filename to prevent path traversal attacks. @@ -97,6 +140,10 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # Find and cache the TTS stage (if any) during initialization self._tts_stage = self._find_tts_stage() self._is_tts = self._tts_stage is not None + self._is_fish_speech = ( + self._tts_stage is not None and getattr(self._tts_stage, "model_stage", None) == "fish_speech_slow_ar" + ) + self._fish_speech_tokenizer = None # Cache TTS configuration values (computed once, reused per request) self._max_instructions_length = self._compute_max_instructions_length() @@ -574,8 +621,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): wav_np = np.mean(wav_np, axis=-1) return wav_np.tolist(), int(sr) - async def _generate_pcm_chunks(self, generator, request_id: str): - """Generate PCM audio chunks for streaming response. + async def _generate_audio_chunks(self, generator, request_id: str, response_format: str = "pcm"): + """Generate audio chunks for streaming response. Handles two audio output modes from the engine: - Cumulative mode (list): Engine returns growing list of chunks; @@ -586,12 +633,15 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): Args: generator: Async generator from the engine request_id: Request identifier for logging + response_format: Audio format (pcm or wav) Yields: - Raw PCM bytes for each audio chunk + Raw audio bytes for each chunk (with WAV header for first chunk if wav format) """ prev_count = 0 sample_rate_val = 24000 + first_chunk = True + try: async for res in generator: audio_output, audio_key = self._extract_audio_output(res) @@ -622,6 +672,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ) if chunk_np.ndim > 1: chunk_np = chunk_np.squeeze() + # For WAV format, emit header before first audio chunk + if response_format == "wav" and first_chunk: + # Assert that sample rate has been set from chunk metadata (not just default) + # This ensures the WAV header contains the correct sample rate + assert sr_raw is not None, ( + "First audio chunk must include sample rate metadata for WAV streaming" + ) + wav_header = _create_wav_header(sample_rate=sample_rate_val, num_channels=1, bits_per_sample=16) + yield wav_header + first_chunk = False + + # Convert audio to PCM bytes audio_obj = CreateAudio( audio_tensor=chunk_np, sample_rate=sample_rate_val, @@ -636,6 +698,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): raise except Exception as e: logger.exception("Streaming speech generation failed for %s: %s", request_id, e) + raise @staticmethod def _extract_audio_output(res) -> tuple[dict | None, str | None]: @@ -650,7 +713,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): mm = getattr(ro, "multimodal_output", None) if ro else None if not mm: return None, None - key = "audio" if "audio" in mm else None + key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None) return mm, key def _build_tts_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: @@ -721,6 +784,194 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return params + def _build_fish_speech_prompt( + self, + request: OpenAICreateSpeechRequest, + ref_audio_data: tuple[list[float], int] | None = None, + ) -> dict[str, Any]: + """Build prompt for Fish Speech S2 Pro. + + Without voice cloning: + <|im_start|>user\\n<|speaker:0|>{text}<|im_end|>\\n<|im_start|>assistant\\n<|voice|> + + With voice cloning (ref_audio + ref_text): + <|im_start|>system\\n<|speaker:0|>{ref_text}<|audio_start|>{semantic_tokens}<|audio_end|><|im_end|> + <|im_start|>user\\n<|speaker:0|>{text}<|im_end|>\\n<|im_start|>assistant\\n<|voice|> + """ + from transformers import AutoTokenizer + + if self._fish_speech_tokenizer is None: + model_name = self.engine_client.model_config.model + self._fish_speech_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + + tokenizer = self._fish_speech_tokenizer + model_name = self.engine_client.model_config.model + + if ref_audio_data is not None and request.ref_text: + # Voice cloning: encode reference audio and build system message. + from vllm_omni.model_executor.models.fish_speech.dac_encoder import ( + encode_reference_audio, + ) + + wav_samples, sr = ref_audio_data + semantic_token_ids = encode_reference_audio(model_name, wav_samples, sr) + + # Build system message with ref text + audio tokens. + audio_start_id = tokenizer.encode("<|audio_start|>", add_special_tokens=False) + audio_end_id = tokenizer.encode("<|audio_end|>", add_special_tokens=False) + + # System content: <|speaker:0|>{ref_text}<|audio_start|>{codes}<|audio_end|> + prefix_text = f"<|speaker:0|>{request.ref_text}" + prefix_ids = tokenizer.encode(prefix_text, add_special_tokens=False) + system_content_ids = prefix_ids + audio_start_id + semantic_token_ids + audio_end_id + + # Manually build system turn: <|im_start|>system\n{content}<|im_end|>\n + im_start = tokenizer.encode("<|im_start|>", add_special_tokens=False) + im_end = tokenizer.encode("<|im_end|>", add_special_tokens=False) + system_tag = tokenizer.encode("system\n", add_special_tokens=False) + newline = tokenizer.encode("\n", add_special_tokens=False) + system_ids = im_start + system_tag + system_content_ids + im_end + newline + + # User turn via chat template. + user_text = f"<|speaker:0|>{request.input}" + user_messages = [{"role": "user", "content": user_text}] + user_ids = tokenizer.apply_chat_template(user_messages, tokenize=True, add_generation_prompt=True) + prompt_ids = system_ids + user_ids + else: + # No voice cloning: simple user message. + user_text = f"<|speaker:0|>{request.input}" + messages = [{"role": "user", "content": user_text}] + prompt_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True) + + # Append <|voice|> token to signal voice generation. + voice_token_id = tokenizer.encode("<|voice|>", add_special_tokens=False) + prompt_ids = prompt_ids + voice_token_id + + additional_information: dict[str, Any] = { + "text": [request.input], + "max_new_tokens": [request.max_new_tokens or 4096], + } + + return { + "prompt_token_ids": prompt_ids, + "additional_information": additional_information, + } + + async def _prepare_speech_generation( + self, + request: OpenAICreateSpeechRequest, + ) -> tuple[str, Any, dict[str, Any]]: + if self.engine_client.errored: + raise self.engine_client.dead_error + + if self._is_fish_speech: + if not request.input or not request.input.strip(): + raise ValueError("Input text cannot be empty") + ref_audio_data = None + if request.ref_audio is not None: + if not request.ref_text or not request.ref_text.strip(): + raise ValueError("Voice cloning requires 'ref_text' (transcript of the reference audio)") + wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + ref_audio_data = (wav_list, sr) + prompt = self._build_fish_speech_prompt(request, ref_audio_data=ref_audio_data) + tts_params = {} + elif self._is_tts: + validation_error = self._validate_tts_request(request) + if validation_error: + raise ValueError(validation_error) + + tts_params = self._build_tts_params(request) + if request.ref_audio is not None: + wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + tts_params["ref_audio"] = [[wav_list, sr]] + + ph_len = self._estimate_prompt_len(tts_params) + prompt = {"prompt_token_ids": [1] * ph_len, "additional_information": tts_params} + else: + tts_params = {} + prompt = {"prompt": request.input} + + request_id = f"speech-{random_uuid()}" + model_type = ( + "fish_speech" + if self._is_fish_speech + else tts_params.get("task_type", ["unknown"])[0] + if self._is_tts + else "generic" + ) + logger.info( + "TTS speech request %s: text=%r, model=%s", + request_id, + request.input[:50] + "..." if len(request.input) > 50 else request.input, + model_type, + ) + + sampling_params_list = self.engine_client.default_sampling_params_list + + # Override Stage-0 max_tokens if caller specified max_new_tokens (Fish Speech). + if self._is_fish_speech and request.max_new_tokens is not None and sampling_params_list: + import copy + + sampling_params_list = copy.deepcopy(sampling_params_list) + sampling_params_list[0].max_tokens = request.max_new_tokens + + generator = self.engine_client.generate( + prompt=prompt, + request_id=request_id, + sampling_params_list=sampling_params_list, + output_modalities=["audio"], + ) + return request_id, generator, tts_params + + async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest): + """Yield raw PCM bytes for a speech request as soon as chunks are decoded.""" + request_id, generator, _ = await self._prepare_speech_generation(request) + async for chunk in self._generate_pcm_chunks(generator, request_id): + yield chunk + + async def _generate_audio_bytes( + self, + request: OpenAICreateSpeechRequest, + ) -> tuple[bytes, str]: + request_id, generator, _ = await self._prepare_speech_generation(request) + + final_output: OmniRequestOutput | None = None + async for res in generator: + final_output = res + + if final_output is None: + raise ValueError("No output generated from the model.") + + audio_output, audio_key = self._extract_audio_output(final_output) + if audio_key is None: + raise ValueError("TTS model did not produce audio output.") + + audio_tensor = audio_output[audio_key] + sr_raw = audio_output.get("sr", 24000) + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) + + if isinstance(audio_tensor, list): + import torch + + audio_tensor = torch.cat(audio_tensor, dim=-1) + if hasattr(audio_tensor, "float"): + audio_tensor = audio_tensor.float().detach().cpu().numpy() + + if audio_tensor.ndim > 1: + audio_tensor = audio_tensor.squeeze() + + audio_obj = CreateAudio( + audio_tensor=audio_tensor, + sample_rate=sample_rate, + response_format=request.response_format or "wav", + speed=request.speed or 1.0, + stream_format=request.stream_format, + base64_encode=False, + ) + audio_response: AudioResponse = self.create_audio(audio_obj) + return audio_response.audio_data, audio_response.media_type + async def create_speech( self, request: OpenAICreateSpeechRequest, @@ -742,97 +993,43 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): - ref_text: Transcript of reference audio (Base task) - x_vector_only_mode: Use speaker embedding only (Base task) - Streaming is supported via stream=True with response_format='pcm'. - Each Code2Wav chunk is yielded as raw PCM bytes as soon as it is decoded. + Streaming is supported via stream=True with response_format='pcm' or 'wav'. + Each Code2Wav chunk is yielded as raw audio bytes as soon as it is decoded. + For WAV format, a header with placeholder size values is emitted first. """ error_check_ret = await self._check_model(request) if error_check_ret is not None: logger.error("Error with model %s", error_check_ret) return error_check_ret - if self.engine_client.errored: - raise self.engine_client.dead_error - - request_id = f"speech-{random_uuid()}" - try: - if self._is_tts: - # Validate TTS parameters - validation_error = self._validate_tts_request(request) - if validation_error: - return self.create_error_response(validation_error) - - tts_params = self._build_tts_params(request) - if request.ref_audio is not None: - wav_list, sr = await self._resolve_ref_audio(request.ref_audio) - tts_params["ref_audio"] = [[wav_list, sr]] - - # Prompt length must match model-side embeddings; values are placeholders. - ph_len = self._estimate_prompt_len(tts_params) - prompt = {"prompt_token_ids": [1] * ph_len, "additional_information": tts_params} - else: - tts_params = {} - prompt = {"prompt": request.input} - - logger.info( - "TTS speech request %s: text=%r, task_type=%s", - request_id, - request.input[:50] + "..." if len(request.input) > 50 else request.input, - tts_params.get("task_type", ["unknown"])[0], - ) - - sampling_params_list = self.engine_client.default_sampling_params_list - - generator = self.engine_client.generate( - prompt=prompt, - request_id=request_id, - sampling_params_list=sampling_params_list, - output_modalities=["audio"], - ) - if request.stream: + # Determine response format and media type for streaming + response_format = (request.response_format or "wav").lower() + + # Only pcm and wav support streaming without post-processing + if response_format not in ["pcm", "wav"]: + return self.create_error_response( + f"Streaming is only supported for 'pcm' and 'wav' formats. " + f"Got '{response_format}'. For other formats, use stream=False." + ) + + # Check if speed adjustment is requested (not compatible with streaming) + if request.speed is not None and request.speed != 1.0: + return self.create_error_response( + "Streaming is not supported with speed adjustment. " + "Use stream=False or remove the speed parameter." + ) + + media_type = "audio/wav" if response_format == "wav" else "audio/pcm" + request_id, generator, _ = await self._prepare_speech_generation(request) return StreamingResponse( - self._generate_pcm_chunks(generator, request_id), - media_type="audio/pcm", + self._generate_audio_chunks(generator, request_id, response_format), + media_type=media_type, ) - # Non-streaming: collect final output - final_output: OmniRequestOutput | None = None - async for res in generator: - final_output = res - - if final_output is None: - return self.create_error_response("No output generated from the model.") - - audio_output, audio_key = self._extract_audio_output(final_output) - if audio_key is None: - return self.create_error_response("TTS model did not produce audio output.") - - audio_tensor = audio_output[audio_key] - sr_raw = audio_output.get("sr", 24000) - sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw - sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) - - # async_chunk mode accumulates chunks as a list; concat first. - if isinstance(audio_tensor, list): - import torch - - audio_tensor = torch.cat(audio_tensor, dim=-1) - if hasattr(audio_tensor, "float"): - audio_tensor = audio_tensor.float().detach().cpu().numpy() - if audio_tensor.ndim > 1: - audio_tensor = audio_tensor.squeeze() - - audio_obj = CreateAudio( - audio_tensor=audio_tensor, - sample_rate=sample_rate, - response_format=request.response_format or "wav", - speed=request.speed or 1.0, - stream_format=request.stream_format, - base64_encode=False, - ) - audio_response = self.create_audio(audio_obj) - return Response(content=audio_response.audio_data, media_type=audio_response.media_type) + audio_bytes, media_type = await self._generate_audio_bytes(request) + return Response(content=audio_bytes, media_type=media_type) except asyncio.CancelledError: return self.create_error_response("Client disconnected") diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py new file mode 100644 index 000000000..d88c88b65 --- /dev/null +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -0,0 +1,289 @@ +"""WebSocket handler for streaming text input TTS. + +Accepts text incrementally via WebSocket, buffers and splits at sentence +boundaries, and generates audio per sentence using the existing TTS pipeline. + +Protocol: + Client -> Server: + {"type": "session.config", ...} # Session config (sent once first) + {"type": "input.text", "text": "..."} # Text chunks + {"type": "input.done"} # End of input + + Server -> Client: + {"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "wav"} + + {"type": "audio.done", "sentence_index": 0} + {"type": "session.done", "total_sentences": N} + {"type": "error", "message": "..."} +""" + +import asyncio +import json +from contextlib import aclosing + +from fastapi import WebSocket, WebSocketDisconnect +from pydantic import ValidationError +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai.protocol.audio import ( + OpenAICreateSpeechRequest, + StreamingSpeechSessionConfig, +) +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech +from vllm_omni.entrypoints.openai.text_splitter import ( + SPLIT_CLAUSE, + SPLIT_SENTENCE, + SentenceSplitter, +) + +logger = init_logger(__name__) + +_DEFAULT_IDLE_TIMEOUT = 30.0 # seconds +_DEFAULT_CONFIG_TIMEOUT = 10.0 # seconds +_PCM_SAMPLE_RATE = 24000 +_MAX_CONFIG_MESSAGE_SIZE = 4 * 1024 * 1024 # allow large ref_audio payloads +_MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024 + + +class OmniStreamingSpeechHandler: + """Handles WebSocket sessions for streaming text-input TTS. + + Each WebSocket connection is an independent session. Text arrives + incrementally, is split at sentence boundaries, and audio is generated + per sentence using the existing OmniOpenAIServingSpeech pipeline. + + Args: + speech_service: The existing TTS serving instance (reused for + validation and audio generation). + idle_timeout: Max seconds to wait for a message before closing. + config_timeout: Max seconds to wait for the initial session.config. + """ + + def __init__( + self, + speech_service: OmniOpenAIServingSpeech, + idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, + ) -> None: + self._speech_service = speech_service + self._idle_timeout = idle_timeout + self._config_timeout = config_timeout + + async def handle_session(self, websocket: WebSocket) -> None: + """Main session loop for a single WebSocket connection.""" + await websocket.accept() + + try: + # 1. Wait for session.config + config = await self._receive_config(websocket) + if config is None: + return # Error already sent, connection closing + + # Validate model if specified + if config.model and hasattr(self._speech_service, "_check_model"): + error = await self._speech_service._check_model( + OpenAICreateSpeechRequest(input="ping", model=config.model) + ) + if error is not None: + await self._send_error(websocket, str(error)) + return + + boundary_re = SPLIT_CLAUSE if config.split_granularity == "clause" else SPLIT_SENTENCE + splitter = SentenceSplitter(boundary_re=boundary_re) + sentence_index = 0 + + # 2. Receive text chunks until input.done + while True: + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._idle_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Idle timeout: no message received") + return + + if len(raw) > _MAX_INPUT_TEXT_MESSAGE_SIZE: + await self._send_error(websocket, "input.text message too large") + continue + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON message") + continue + + if not isinstance(msg, dict): + await self._send_error(websocket, "WebSocket messages must be JSON objects") + continue + + msg_type = msg.get("type") + + if msg_type == "input.text": + text = msg.get("text", "") + if not isinstance(text, str): + await self._send_error(websocket, "input.text requires a string value") + continue + sentences = splitter.add_text(text) + for sentence in sentences: + await self._generate_and_send(websocket, config, sentence, sentence_index) + sentence_index += 1 + + elif msg_type == "input.done": + # Flush remaining buffer + remaining = splitter.flush() + if remaining: + await self._generate_and_send(websocket, config, remaining, sentence_index) + sentence_index += 1 + + # Send session.done + await websocket.send_json( + { + "type": "session.done", + "total_sentences": sentence_index, + } + ) + return + + else: + await self._send_error( + websocket, + f"Unknown message type: {msg_type}", + ) + + except WebSocketDisconnect: + logger.info("Streaming speech: client disconnected") + except Exception as e: + logger.exception("Streaming speech session error: %s", e) + try: + await self._send_error(websocket, f"Internal error: {e}") + except Exception: + logger.debug("Failed to send error to streaming speech client", exc_info=True) + + async def _receive_config(self, websocket: WebSocket) -> StreamingSpeechSessionConfig | None: + """Wait for and validate the session.config message.""" + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._config_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Timeout waiting for session.config") + return None + + if len(raw) > _MAX_CONFIG_MESSAGE_SIZE: + await self._send_error(websocket, "session.config message too large") + return None + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON in session.config") + return None + + if not isinstance(msg, dict): + await self._send_error(websocket, "session.config must be a JSON object") + return None + + if msg.get("type") != "session.config": + await self._send_error( + websocket, + f"Expected session.config, got: {msg.get('type')}", + ) + return None + + try: + config = StreamingSpeechSessionConfig(**{k: v for k, v in msg.items() if k != "type"}) + except ValidationError as e: + await self._send_error(websocket, f"Invalid session config: {e}") + return None + + return config + + async def _generate_and_send( + self, + websocket: WebSocket, + config: StreamingSpeechSessionConfig, + sentence_text: str, + sentence_index: int, + ) -> None: + """Generate audio for a single sentence and send it over WebSocket.""" + response_format = config.response_format or "wav" + + request = OpenAICreateSpeechRequest( + input=sentence_text, + model=config.model, + voice=config.voice, + task_type=config.task_type, + language=config.language, + instructions=config.instructions, + response_format=response_format, + speed=config.speed, + max_new_tokens=config.max_new_tokens, + initial_codec_chunk_frames=config.initial_codec_chunk_frames, + ref_audio=config.ref_audio, + ref_text=config.ref_text, + x_vector_only_mode=config.x_vector_only_mode, + stream=config.stream_audio, + ) + + start_payload = { + "type": "audio.start", + "sentence_index": sentence_index, + "sentence_text": sentence_text, + "format": response_format, + } + if config.stream_audio and response_format == "pcm": + start_payload["sample_rate"] = _PCM_SAMPLE_RATE + await websocket.send_json(start_payload) + + total_bytes = 0 + generation_failed = False + request_id = None + try: + if config.stream_audio: + request_id, generator, _ = await self._speech_service._prepare_speech_generation(request) + async with aclosing(self._speech_service._generate_pcm_chunks(generator, request_id)) as stream: + async for chunk in stream: + total_bytes += len(chunk) + await websocket.send_bytes(chunk) + else: + audio_bytes, _ = await self._speech_service._generate_audio_bytes(request) + total_bytes = len(audio_bytes) + await websocket.send_bytes(audio_bytes) + except WebSocketDisconnect: + if request_id is not None: + try: + await self._speech_service.engine_client.abort(request_id) + except Exception: + logger.debug("Failed to abort streaming speech request %s", request_id, exc_info=True) + raise + except Exception as e: + generation_failed = True + logger.error("Generation failed for sentence %d: %s", sentence_index, e) + await self._send_error(websocket, f"Generation failed for sentence {sentence_index}: {e}") + finally: + try: + await websocket.send_json( + { + "type": "audio.done", + "sentence_index": sentence_index, + "total_bytes": total_bytes, + "error": generation_failed, + } + ) + except Exception: + logger.debug("Failed to send audio.done for sentence %d", sentence_index, exc_info=True) + + @staticmethod + async def _send_error(websocket: WebSocket, message: str) -> None: + """Send an error message to the client.""" + try: + await websocket.send_json( + { + "type": "error", + "message": message, + } + ) + except Exception: + pass # Connection may already be closed; safe to ignore diff --git a/vllm_omni/entrypoints/openai/text_splitter.py b/vllm_omni/entrypoints/openai/text_splitter.py new file mode 100644 index 000000000..9c02de139 --- /dev/null +++ b/vllm_omni/entrypoints/openai/text_splitter.py @@ -0,0 +1,120 @@ +"""Multi-language sentence boundary detector for streaming TTS input. + +Buffers incoming text and splits at sentence boundaries (English and CJK), +yielding complete sentences for audio generation. +""" + +import re +from re import Pattern + +# Maximum buffer size (in characters) to prevent unbounded memory growth. +_MAX_BUFFER_SIZE = 100_000 # ~100 KB of text + +# Sentence-level: .!? + CJK sentence-ending γ€‚οΌοΌŸ +# NOTE: English requires trailing whitespace to confirm a boundary β€” +# end-of-string is NOT treated as a boundary (that is what flush() is for). +SPLIT_SENTENCE = re.compile( + r"(?<=[.!?])\s+" + r"|(?<=[γ€‚οΌοΌŸ])" +) + +# Clause-level: adds CJK commas ， and semicolons οΌ› +SPLIT_CLAUSE = re.compile( + r"(?<=[.!?])\s+" + r"|(?<=[γ€‚οΌοΌŸοΌŒοΌ›])" +) + +# Default alias +_SENTENCE_BOUNDARY_RE = SPLIT_SENTENCE + + +class SentenceSplitter: + """Incremental sentence splitter for streaming text input. + + Buffers text and yields complete sentences when boundaries are detected. + Designed for TTS pipelines where text arrives incrementally (e.g., from STT). + + Args: + min_sentence_length: Minimum character length for a sentence. + Sentences shorter than this are kept in the buffer to avoid + splitting on abbreviations like "Dr." or "U.S.". + boundary_re: Custom compiled regex for sentence boundaries. + Use ``SPLIT_SENTENCE`` (default) for sentence-level splitting, + ``SPLIT_CLAUSE`` for finer-grained clause-level splitting, + or pass your own ``re.Pattern``. + """ + + def __init__( + self, + min_sentence_length: int = 2, + boundary_re: Pattern[str] | None = None, + ) -> None: + self._buffer: str = "" + self._min_sentence_length = min_sentence_length + self._boundary_re = boundary_re or _SENTENCE_BOUNDARY_RE + + @property + def buffer(self) -> str: + """Current buffered text.""" + return self._buffer + + def add_text(self, text: str) -> list[str]: + """Add text to the buffer and return any complete sentences. + + Args: + text: Incoming text chunk. + + Returns: + List of complete sentences extracted from the buffer. + May be empty if no sentence boundary was found. + + Raises: + ValueError: If the buffer exceeds the maximum size. + """ + if not text: + return [] + + self._buffer += text + if len(self._buffer) > _MAX_BUFFER_SIZE: + raise ValueError( + f"Text buffer exceeded maximum size ({_MAX_BUFFER_SIZE} chars). " + "Consider adding sentence-ending punctuation to your input." + ) + return self._extract_sentences() + + def flush(self) -> str | None: + """Flush remaining buffered text as a final sentence. + + Returns: + The remaining buffered text (stripped), or None if buffer is empty. + """ + remaining = self._buffer.strip() + self._buffer = "" + return remaining if remaining else None + + def _extract_sentences(self) -> list[str]: + """Split buffer at sentence boundaries, keeping incomplete text buffered.""" + parts = self._boundary_re.split(self._buffer) + + if len(parts) <= 1: + # No boundary found β€” keep everything in buffer + return [] + + sentences: list[str] = [] + carry = "" + # All parts except the last are complete sentences + for i in range(len(parts) - 1): + text = carry + parts[i] + carry = "" + stripped = text.strip() + if len(stripped) >= self._min_sentence_length: + sentences.append(stripped) + elif stripped: + # Too short (e.g. "Dr.") β€” carry forward to next part + carry = text + # else: empty, skip + + # Last part stays in buffer (may be incomplete) + self._buffer = carry + parts[-1] + + return sentences diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index de4a71719..99f2042e2 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -9,7 +9,7 @@ from collections.abc import Callable from multiprocessing import shared_memory as _shm from typing import Any -from omegaconf import OmegaConf +from vllm_omni.config.yaml_util import to_dict as _omega_to_dict logger = logging.getLogger(__name__) @@ -315,7 +315,7 @@ def _to_dict(x: Any) -> dict[str, Any]: try: if isinstance(x, dict): return dict(x) - return OmegaConf.to_container(x, resolve=True) # type: ignore[arg-type] + return _omega_to_dict(x) except Exception: try: return dict(x) diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 69ce73fc4..0e31bfa7c 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -5,11 +5,11 @@ from dataclasses import asdict, fields, is_dataclass from pathlib import Path from typing import Any, get_args, get_origin -from omegaconf import OmegaConf from vllm.logger import init_logger from vllm.transformers_utils.config import get_config, get_hf_file_to_dict from vllm.transformers_utils.repo_utils import file_or_path_exists +from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs from vllm_omni.entrypoints.stage_utils import _to_dict from vllm_omni.platforms import current_omni_platform @@ -234,6 +234,11 @@ def resolve_model_config_path(model: str) -> str: def load_stage_configs_from_model(model: str, base_engine_args: dict | None = None) -> list: """Load stage configurations from model's default config file. + .. deprecated:: + This is the legacy OmegaConf-based loading path. New code should use + ``StageConfigFactory.create_from_model()`` instead. This function will + be removed once all callers are migrated (see PR series [2/N]). + Loads stage configurations based on the model type and device type. First tries to load a device-specific YAML file from stage_configs/{device_type}/ directory. If not found, falls back to the default config file. @@ -259,6 +264,9 @@ def load_stage_configs_from_model(model: str, base_engine_args: dict | None = No def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None = None) -> list: """Load stage configurations from a YAML file. + .. deprecated:: + Legacy OmegaConf-based loader. Will be removed in PR series [2/N]. + Args: config_path: Path to the YAML configuration file @@ -267,17 +275,17 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None """ if base_engine_args is None: base_engine_args = {} - config_data = OmegaConf.load(config_path) + config_data = load_yaml_config(config_path) stage_args = config_data.stage_args global_async_chunk = config_data.get("async_chunk", False) - # Convert any nested dataclass objects to dicts before creating OmegaConf + # Convert any nested dataclass objects to dicts before creating DictConfig base_engine_args = _convert_dataclasses_to_dict(base_engine_args) - base_engine_args = OmegaConf.create(base_engine_args) + base_engine_args = create_config(base_engine_args) for stage_arg in stage_args: base_engine_args_tmp = base_engine_args.copy() # Update base_engine_args with stage-specific engine_args if they exist if hasattr(stage_arg, "engine_args") and stage_arg.engine_args is not None: - base_engine_args_tmp = OmegaConf.merge(base_engine_args_tmp, stage_arg.engine_args) + base_engine_args_tmp = create_config(merge_configs(base_engine_args_tmp, stage_arg.engine_args)) stage_type = getattr(stage_arg, "stage_type", "llm") if hasattr(stage_arg, "runtime") and stage_arg.runtime is not None and stage_type != "diffusion": runtime_cfg = stage_arg.runtime @@ -312,7 +320,7 @@ def load_and_resolve_stage_configs( if not stage_configs: if default_stage_cfg_factory is not None: default_stage_cfg = default_stage_cfg_factory() - stage_configs = OmegaConf.create(default_stage_cfg) + stage_configs = create_config(default_stage_cfg) else: stage_configs = [] else: @@ -347,6 +355,7 @@ def get_final_stage_id_for_e2e( output_modalities = default_modalities try: + final_stage_id_for_e2e = last_stage_id for _sid in range(last_stage_id, -1, -1): if ( getattr(stage_list[_sid], "final_output", False) @@ -354,8 +363,6 @@ def get_final_stage_id_for_e2e( ): final_stage_id_for_e2e = _sid break - if final_stage_id_for_e2e < 0: - final_stage_id_for_e2e = last_stage_id except Exception as e: logger.debug( "[Orchestrator] Failed to determine final stage for E2E; \ diff --git a/vllm_omni/model_executor/models/fish_speech/__init__.py b/vllm_omni/model_executor/models/fish_speech/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/models/fish_speech/configuration_fish_speech.py b/vllm_omni/model_executor/models/fish_speech/configuration_fish_speech.py new file mode 100644 index 000000000..9ed4ed013 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/configuration_fish_speech.py @@ -0,0 +1,154 @@ +"""Configuration classes for Fish Speech S2 Pro (fish_qwen3_omni). + +The HuggingFace config uses field names from the original Fish Speech codebase +(dim, n_head, n_layer, etc.). This module re-exports them under the standard +Transformers/Qwen3 attribute names so that vLLM's ``Qwen3Model`` can consume +the Slow AR config transparently. +""" + +from __future__ import annotations + +from transformers import PretrainedConfig + + +class FishSpeechSlowARConfig(PretrainedConfig): + """Slow AR (text_model) config -- Qwen3-based transformer. + + Maps Fish Speech field names to Qwen3-compatible attribute names so + ``vllm.model_executor.models.qwen3.Qwen3Model`` works out of the box. + """ + + model_type = "fish_qwen3" + + def __init__( + self, + vocab_size: int = 155776, + dim: int = 2560, + n_head: int = 32, + n_local_heads: int = 8, + head_dim: int = 128, + n_layer: int = 36, + intermediate_size: int = 9728, + attention_qk_norm: bool = True, + rope_base: float = 1_000_000.0, + max_seq_len: int = 32768, + tie_word_embeddings: bool = True, + codebook_size: int = 4096, + num_codebooks: int = 10, + semantic_begin_id: int = 0, + semantic_end_id: int = 0, + scale_codebook_embeddings: bool = False, + rms_norm_eps: float = 1e-6, + **kwargs, + ): + # Fish Speech field names β†’ standard Transformers/Qwen3 names. + self.hidden_size = dim + self.num_attention_heads = n_head + self.num_key_value_heads = n_local_heads + self.head_dim = head_dim + self.num_hidden_layers = n_layer + self.intermediate_size = intermediate_size + self.max_position_embeddings = max_seq_len + self.rope_theta = rope_base + self.rms_norm_eps = rms_norm_eps + self.hidden_act = "silu" + self.attention_bias = False + self.attention_qk_norm = attention_qk_norm + + # Fish Speech codec / codebook fields. + self.codebook_size = codebook_size + self.num_codebooks = num_codebooks + self.semantic_begin_id = semantic_begin_id + self.semantic_end_id = semantic_end_id + self.scale_codebook_embeddings = scale_codebook_embeddings + + super().__init__( + vocab_size=vocab_size, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class FishSpeechFastARConfig(PretrainedConfig): + """Fast AR (audio_decoder) config -- 4-layer residual codebook predictor.""" + + model_type = "fish_qwen3_audio_decoder" + + def __init__( + self, + vocab_size: int = 4096, + num_codebooks: int = 10, + dim: int = 2560, + n_head: int = 32, + n_local_heads: int = 8, + head_dim: int = 128, + n_layer: int = 4, + intermediate_size: int = 9728, + max_seq_len: int = 11, + text_dim: int = 2560, + audio_hidden_dim: int = 5120, + attention_qk_norm: bool = False, + rope_base: float = 1_000_000.0, + rms_norm_eps: float = 1e-6, + **kwargs, + ): + self.hidden_size = dim + self.num_attention_heads = n_head + self.num_key_value_heads = n_local_heads + self.head_dim = head_dim + self.num_hidden_layers = n_layer + self.intermediate_size = intermediate_size + self.max_position_embeddings = max_seq_len + self.rope_theta = rope_base + self.rms_norm_eps = rms_norm_eps + self.hidden_act = "silu" + self.attention_bias = False + self.attention_qk_norm = attention_qk_norm + self.num_codebooks = num_codebooks + self.text_dim = text_dim + self.audio_hidden_dim = audio_hidden_dim + + super().__init__(vocab_size=vocab_size, **kwargs) + + +class FishSpeechConfig(PretrainedConfig): + """Top-level config for Fish Speech S2 Pro (fish_qwen3_omni). + + Wraps ``text_config`` (Slow AR) and ``audio_decoder_config`` (Fast AR). + """ + + model_type = "fish_qwen3_omni" + sub_configs = { + "text_config": FishSpeechSlowARConfig, + "audio_decoder_config": FishSpeechFastARConfig, + } + + def __init__( + self, + text_config: dict | FishSpeechSlowARConfig | None = None, + audio_decoder_config: dict | FishSpeechFastARConfig | None = None, + semantic_start_token_id: int = 151678, + semantic_end_token_id: int = 155773, + audio_pad_token_id: int = 151677, + **kwargs, + ): + if isinstance(text_config, dict): + text_config = FishSpeechSlowARConfig(**text_config) + if isinstance(audio_decoder_config, dict): + audio_decoder_config = FishSpeechFastARConfig(**audio_decoder_config) + + self.text_config = text_config or FishSpeechSlowARConfig() + self.audio_decoder_config = audio_decoder_config or FishSpeechFastARConfig() + + self.semantic_start_token_id = semantic_start_token_id + self.semantic_end_token_id = semantic_end_token_id + self.audio_pad_token_id = audio_pad_token_id + + # Propagate semantic IDs into text_config for convenience. + self.text_config.semantic_begin_id = semantic_start_token_id + self.text_config.semantic_end_id = semantic_end_token_id + + super().__init__(**kwargs) + + def get_text_config(self, **kwargs) -> FishSpeechSlowARConfig: + return self.text_config diff --git a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py new file mode 100644 index 000000000..084a863ad --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py @@ -0,0 +1,111 @@ +"""DAC codec encoder for Fish Speech S2 Pro voice cloning. + +Encodes reference audio into VQ codes for use as prompt conditioning. +Runs on CPU in the API server process -- loaded lazily on first use. +""" + +from __future__ import annotations + +import os + +import numpy as np +import torch +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.fish_speech.dac_utils import ( + DAC_SAMPLE_RATE, + build_dac_codec, +) + +logger = init_logger(__name__) + +_codec_cache: dict[str, nn.Module] = {} + + +def _load_dac_codec(model_path: str) -> nn.Module: + """Load the DAC codec model from codec.pth (cached per model_path).""" + if model_path in _codec_cache: + return _codec_cache[model_path] + + codec_path = os.path.join(model_path, "codec.pth") + if not os.path.exists(codec_path): + from transformers.utils.hub import cached_file + + cached = cached_file(model_path, "codec.pth") + if cached is not None: + codec_path = cached + + if not os.path.exists(codec_path): + raise FileNotFoundError( + f"codec.pth not found for {model_path}. Required for voice cloning with Fish Speech S2 Pro." + ) + + codec = build_dac_codec() + + state_dict = torch.load(codec_path, map_location="cpu", weights_only=True) + if "generator" in state_dict: + state_dict = state_dict["generator"] + codec.load_state_dict(state_dict, strict=False) + codec.eval() + + _codec_cache[model_path] = codec + logger.info("Loaded DAC codec encoder from %s (CPU)", codec_path) + return codec + + +def _resample(wav: np.ndarray, sr: int, target_sr: int) -> np.ndarray: + """Resample audio using torchaudio's polyphase resampling.""" + if sr == target_sr: + return wav + import torchaudio + + wav_t = torch.from_numpy(wav).unsqueeze(0).float() + wav_t = torchaudio.functional.resample(wav_t, sr, target_sr) + return wav_t.squeeze(0).numpy() + + +@torch.no_grad() +def encode_reference_audio( + model_path: str, + wav_samples: list[float] | np.ndarray, + sample_rate: int, +) -> list[int]: + """Encode reference audio into semantic token IDs for prompt conditioning. + + Args: + model_path: HuggingFace model path (for locating codec.pth). + wav_samples: Audio waveform samples (mono, float). + sample_rate: Sample rate of the input audio. + + Returns: + List of semantic token IDs (151678 + code_value for each frame). + """ + codec = _load_dac_codec(model_path) + + wav = np.asarray(wav_samples, dtype=np.float32) + if wav.ndim > 1: + wav = np.mean(wav, axis=-1) + + # Resample to DAC sample rate (44100). + wav = _resample(wav, sample_rate, DAC_SAMPLE_RATE) + + # Encode: [1, 1, T] -> codes [1, num_codebooks, num_frames] + wav_tensor = torch.from_numpy(wav).unsqueeze(0).unsqueeze(0).float() + feature_lengths = torch.tensor([wav_tensor.shape[-1]]) + codes, feature_lengths_out = codec.encode(wav_tensor, feature_lengths) + + # Extract semantic codebook (index 0) - shape [num_frames]. + semantic_codes = codes[0, 0, :].tolist() + + # Convert to semantic token IDs: <|semantic:{i}|> = 151678 + i + SEMANTIC_TOKEN_OFFSET = 151678 + semantic_token_ids = [SEMANTIC_TOKEN_OFFSET + int(c) for c in semantic_codes] + + logger.info( + "Encoded reference audio: %d samples @ %dHz -> %d semantic tokens", + len(wav_samples), + sample_rate, + len(semantic_token_ids), + ) + return semantic_token_ids diff --git a/vllm_omni/model_executor/models/fish_speech/dac_utils.py b/vllm_omni/model_executor/models/fish_speech/dac_utils.py new file mode 100644 index 000000000..f1f2590a3 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/dac_utils.py @@ -0,0 +1,92 @@ +"""Shared DAC codec construction for Fish Speech S2 Pro. + +Used by both the encoder (voice cloning, CPU) and the decoder (Stage 1, GPU). +""" + +from __future__ import annotations + +import torch.nn as nn + +# Default DAC codec config matching s2-pro's modded_dac_vq.yaml. +DAC_SAMPLE_RATE = 44100 +DAC_HOP_LENGTH = 2048 # 512 (decoder upsample) * 4 (quantizer upsample) +DAC_NUM_CODEBOOKS = 10 # 1 semantic + 9 residual + + +def build_dac_codec() -> nn.Module: + """Construct a DAC codec model (uninitialized weights). + + Returns the model on CPU in eval mode -- caller is responsible for + loading weights and moving to the target device. + """ + from fish_speech.models.dac.modded_dac import ( + DAC, + ModelArgs, + WindowLimitedTransformer, + ) + from fish_speech.models.dac.rvq import DownsampleResidualVectorQuantize + + base_transformer_kwargs = dict( + block_size=16384, + n_local_heads=-1, + head_dim=64, + rope_base=10000, + norm_eps=1e-5, + dropout_rate=0.0, + attn_dropout_rate=0.0, + channels_first=True, + ) + + def _make_transformer_config(**kw): + return ModelArgs(**{**base_transformer_kwargs, **kw}) + + quantizer_transformer_config = ModelArgs( + block_size=4096, + n_layer=8, + n_head=16, + dim=1024, + intermediate_size=3072, + n_local_heads=-1, + head_dim=64, + rope_base=10000, + norm_eps=1e-5, + dropout_rate=0.0, + attn_dropout_rate=0.0, + channels_first=True, + ) + post_module = WindowLimitedTransformer( + causal=True, + window_size=128, + input_dim=1024, + config=quantizer_transformer_config, + ) + pre_module = WindowLimitedTransformer( + causal=True, + window_size=128, + input_dim=1024, + config=quantizer_transformer_config, + ) + quantizer = DownsampleResidualVectorQuantize( + input_dim=1024, + n_codebooks=9, + codebook_size=1024, + codebook_dim=8, + quantizer_dropout=0.0, + downsample_factor=[2, 2], + post_module=post_module, + pre_module=pre_module, + semantic_codebook_size=4096, + ) + codec = DAC( + encoder_dim=64, + encoder_rates=[2, 4, 8, 8], + decoder_dim=1536, + decoder_rates=[8, 8, 4, 2], + quantizer=quantizer, + sample_rate=DAC_SAMPLE_RATE, + causal=True, + encoder_transformer_layers=[0, 0, 0, 4], + decoder_transformer_layers=[4, 0, 0, 0], + transformer_general_config=_make_transformer_config, + ) + return codec diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py new file mode 100644 index 000000000..a889f66fd --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py @@ -0,0 +1,272 @@ +"""Fish Speech S2 Pro -- DAC Decoder (Stage 1). + +Loads the DAC codec from ``codec.pth`` and decodes codebook indices +[num_codebooks, T] β†’ audio waveform at 44.1 kHz. + +Analogous to ``Qwen3TTSCode2Wav`` in qwen3_tts. + +Requires the ``fish-speech`` package for the DAC model architecture. +Install with: ``pip install fish-speech`` +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context, is_forward_context_available +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.fish_speech.dac_utils import ( + DAC_HOP_LENGTH, + DAC_NUM_CODEBOOKS, + DAC_SAMPLE_RATE, + build_dac_codec, +) +from vllm_omni.model_executor.models.output_templates import OmniOutput + +logger = init_logger(__name__) + + +class FishSpeechDACDecoder(nn.Module): + """Stage-1 DAC decoder for Fish Speech S2 Pro (GenerationModelRunner). + + Consumes frame-aligned codec tokens from input_ids and decodes waveform + via the DAC codec decoder. + """ + + input_modalities = "audio" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + self._codec: nn.Module | None = None + self._num_codebooks: int = DAC_NUM_CODEBOOKS + self._output_sample_rate: int = DAC_SAMPLE_RATE + self._hop_length: int = DAC_HOP_LENGTH + self._logged_codec_stats = False + + def _ensure_codec_loaded(self) -> None: + if self._codec is not None: + return + + codec_path = os.path.join(self.model_path, "codec.pth") + if not os.path.exists(codec_path): + # Try HuggingFace cache. + try: + from transformers.utils.hub import cached_file + + cached = cached_file(self.model_path, "codec.pth") + if cached is not None: + codec_path = cached + except Exception: + pass + + if not os.path.exists(codec_path): + raise FileNotFoundError( + f"codec.pth not found at {codec_path}. Make sure the Fish Speech S2 Pro model includes codec.pth." + ) + + codec = build_dac_codec() + + # Load weights. + state_dict = torch.load(codec_path, map_location="cpu", weights_only=True) + # Some checkpoints wrap under "generator" key. + if "generator" in state_dict: + state_dict = state_dict["generator"] + codec.load_state_dict(state_dict, strict=False) + + device = self.vllm_config.device_config.device + codec = codec.to(device=device, dtype=torch.float32) + codec.eval() + self._codec = codec + + logger.info( + "Fish Speech DAC codec loaded from %s (device=%s, sample_rate=%d)", + codec_path, + device, + self._output_sample_rate, + ) + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + if input_ids.numel() == 0: + return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None: + return None + + def _split_request_ids( + self, + ids: torch.Tensor, + seq_token_counts: list[int] | None = None, + ) -> list[torch.Tensor]: + if seq_token_counts is not None and len(seq_token_counts) > 1: + boundaries = [0] + for count in seq_token_counts: + boundaries.append(boundaries[-1] + count) + n = ids.numel() + return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))] + if is_forward_context_available(): + slices = get_forward_context().ubatch_slices + if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices): + boundaries = [0] + for s in slices: + boundaries.append(boundaries[-1] + s) + return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)] + return [ids] + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + """Decode codec codes into audio waveform. + + input_ids layout per request: flat codes [num_codebooks * num_frames]. + Codes are codebook-major: [cb0_f0, cb0_f1, ..., cb0_fN, cb1_f0, ...]. + """ + self._ensure_codec_loaded() + assert self._codec is not None + + q = self._num_codebooks + sr_val = self._output_sample_rate + sr_tensor = torch.tensor(sr_val, dtype=torch.int32) + empty = torch.zeros((0,), dtype=torch.float32) + + if input_ids is None or input_ids.numel() == 0: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]}, + ) + + ids = input_ids.reshape(-1).to(dtype=torch.long) + request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts")) + + parsed: list[tuple[int, int]] = [] + valid_codes_qf: list[torch.Tensor] = [] + valid_indices: list[int] = [] + left_context_size = [0] * len(request_ids_list) + if runtime_additional_information is not None: + for i, info in enumerate(runtime_additional_information): + if i >= len(left_context_size): + break + if "left_context_size" in info: + left_context_size[i] = info["left_context_size"] + + for i, req_ids in enumerate(request_ids_list): + if req_ids.numel() < 1: + parsed.append((0, 0)) + continue + ctx_frames = left_context_size[i] + flat = req_ids + n = flat.numel() + if n == 0 or n % q != 0: + if n > 0: + logger.warning( + "DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.", + n, + q, + ) + parsed.append((0, 0)) + continue + frames = n // q + codes_qf = flat.reshape(q, frames) + parsed.append((ctx_frames, frames)) + valid_codes_qf.append(codes_qf) + valid_indices.append(i) + + num_req = len(request_ids_list) + if not valid_codes_qf: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": [empty] * num_req, + "sr": [sr_tensor] * num_req, + }, + ) + + if not self._logged_codec_stats: + self._logged_codec_stats = True + try: + c = valid_codes_qf[0] + logger.info( + "DAC decoder: frames=%d q=%d uniq=%d range=[%d,%d] batch=%d", + c.shape[1], + q, + int(torch.unique(c).numel()), + int(c.min().item()), + int(c.max().item()), + len(valid_codes_qf), + ) + except Exception: + pass + + # Decode each request individually. + wav_tensors: list[torch.Tensor] = [] + for codes_qf in valid_codes_qf: + codes_bqf = codes_qf.unsqueeze(0) # [1, num_codebooks, num_frames] + num_frames = codes_qf.shape[1] + feature_lengths = torch.tensor([num_frames], device=codes_bqf.device) + with torch.cuda.amp.autocast(enabled=False): + wav, audio_lengths = self._codec.decode(codes_bqf, feature_lengths) + # wav shape: [1, 1, wav_len] + wav_tensors.append(wav.squeeze(0).squeeze(0)) # [wav_len] + + audios: list[torch.Tensor] = [empty] * num_req + srs = [sr_tensor] * num_req + + for j, idx in enumerate(valid_indices): + ctx_frames, actual_frames = parsed[idx] + wav = wav_tensors[j] + # Trim context frames (left overlap for streaming). + if ctx_frames > 0: + cut = ctx_frames * self._hop_length + if cut < wav.shape[0]: + wav = wav[cut:] + else: + logger.warning( + "Context trim %d >= decoded length %d; returning empty audio.", + cut, + wav.shape[0], + ) + continue + if wav.shape[0] > 0: + audios[idx] = wav.to(dtype=torch.float32).reshape(-1) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audios, "sr": srs}, + ) + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + if not (isinstance(model_outputs, tuple) and len(model_outputs) == 2): + raise TypeError(f"FishSpeechDACDecoder expected (audio_tensor, sr), got {type(model_outputs)}") + audio_tensor, sr = model_outputs + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audio_tensor, "sr": sr}, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # DAC codec weights are loaded lazily from codec.pth, not from the main checkpoint. + return set() diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py new file mode 100644 index 000000000..f73914c28 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py @@ -0,0 +1,470 @@ +"""Fast AR (audio_decoder) for Fish Speech S2 Pro. + +4-layer transformer that predicts residual codebook codes (1..num_codebooks-1) +autoregressively after each Slow AR step. Analogous to Qwen3 TTS's +CodePredictor but with its own embedding table, RoPE, and output head. + +Uses re-prefill (no KV cache): each AR step forwards the full growing +sequence through the 4-layer transformer. This trades ~O(T^2) attention +FLOPs (negligible for T=10, 4 layers) for zero KV cache management. + +Optimisations: + - torch.compile on model forward (kernel fusion) + - Pre-allocated projection / embedding buffer + - Pre-allocated position_ids + - Inline sampling +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.config import VllmConfig +from vllm.config.vllm import set_current_vllm_config +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.utils import is_pp_missing_parameter + +from .configuration_fish_speech import FishSpeechFastARConfig, FishSpeechSlowARConfig + +logger = init_logger(__name__) + + +# =================================================================== +# Standalone Fast AR Layers (no vLLM paged attention) +# =================================================================== + + +class _FastARAttention(nn.Module): + """Multi-head attention using F.scaled_dot_product_attention (SDPA). + + Supports fused QKV, RoPE, optional q/k normalization, and native GQA. + Input: [B, seq_len, hidden_size]. + """ + + def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + 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._use_gqa = self.num_kv_heads != self.num_heads + + self.qkv_proj = QKVParallelLinear( + hidden_size=self.hidden_size, + head_size=self.head_dim, + total_num_heads=self.num_heads, + total_num_kv_heads=self.num_kv_heads, + bias=False, + prefix=f"{prefix}.qkv_proj", + disable_tp=True, + ) + self.o_proj = RowParallelLinear( + input_size=self.num_heads * self.head_dim, + output_size=self.hidden_size, + bias=False, + prefix=f"{prefix}.o_proj", + disable_tp=True, + ) + self.rotary_emb = get_rope( + self.head_dim, + max_position=config.max_position_embeddings, + is_neox_style=False, # Fish Speech uses interleaved (GPT-J) RoPE + rope_parameters={"rope_theta": config.rope_theta, "rope_type": "default"}, + ) + # Fish Speech Fast AR typically has attention_qk_norm=false. + if config.attention_qk_norm: + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + else: + self.q_norm = None + self.k_norm = None + + def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + bsz, seq_len, _ = hidden_states.shape + + qkv, _ = self.qkv_proj(hidden_states.reshape(bsz * seq_len, -1)) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + if self.q_norm is not None: + q = self.q_norm(q.view(-1, self.num_heads, self.head_dim)).view(q.shape) + k = self.k_norm(k.view(-1, self.num_kv_heads, self.head_dim)).view(k.shape) + + q, k = self.rotary_emb(position_ids, q, k) + + q = q.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) + + attn_out = F.scaled_dot_product_attention( + q, + k, + v, + scale=self.scaling, + is_causal=True, + enable_gqa=self._use_gqa, + ) + attn_out = attn_out.transpose(1, 2).reshape(bsz * seq_len, -1) + output, _ = self.o_proj(attn_out) + return output.view(bsz, seq_len, -1) + + +class _FastARMLP(nn.Module): + """SiLU-gated MLP, matching Qwen3/LLaMA MLP structure.""" + + def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=config.hidden_size, + output_sizes=[config.intermediate_size] * 2, + bias=False, + prefix=f"{prefix}.gate_up_proj", + disable_tp=True, + ) + self.down_proj = RowParallelLinear( + input_size=config.intermediate_size, + output_size=config.hidden_size, + bias=False, + prefix=f"{prefix}.down_proj", + disable_tp=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + gate, up = gate_up.chunk(2, dim=-1) + x = F.silu(gate) * up + x, _ = self.down_proj(x) + return x + + +class _FastARDecoderLayer(nn.Module): + """Transformer decoder layer for Fast AR (SDPA, no KV cache).""" + + def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.self_attn = _FastARAttention(config, prefix=f"{prefix}.self_attn") + self.mlp = _FastARMLP(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, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_ids) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +# =================================================================== +# Fast AR Transformer Model +# =================================================================== + + +class FishSpeechFastARModel(nn.Module): + """4-layer transformer for residual codebook prediction (re-prefill).""" + + def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None: + super().__init__() + self.config = config + self.layers = nn.ModuleList( + [_FastARDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)] + ) + # NOTE: final norm is handled by FishSpeechFastAR.fast_norm (one norm weight + # in checkpoint: audio_decoder.norm.weight β†’ fast_ar.fast_norm.weight). + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + hidden_states = inputs_embeds + for layer in self.layers: + hidden_states = layer(hidden_states, position_ids) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + ("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 + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped.endswith(".bias") and mapped not in params_dict: + continue + if is_pp_missing_parameter(mapped, self): + continue + if mapped.endswith("scale"): + mapped = maybe_remap_kv_scale_name(mapped, params_dict) + if mapped is None: + continue + param = params_dict.get(mapped) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(mapped) + break + else: + mapped = maybe_remap_kv_scale_name(name, params_dict) + if mapped is None: + continue + if name.endswith(".bias") and mapped not in params_dict: + continue + if is_pp_missing_parameter(mapped, self): + continue + param = params_dict.get(mapped) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(mapped) + return loaded_params + + +# =================================================================== +# Fast AR Wrapper (optimised re-prefill + torch.compile) +# =================================================================== + + +class FishSpeechFastAR(nn.Module): + """vLLM-native Fast AR for Fish Speech (residual codebooks). + + Re-prefill approach: each AR step forwards the full growing sequence + through the 4-layer transformer. No KV cache needed. + + Optimisations over baseline: + 1. torch.compile on model forward -- kernel fusion. + 2. Pre-allocated embedding buffer [B, max_seq, H]. + 3. Pre-allocated position_ids. + 4. Inline sampling. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + config: FishSpeechFastARConfig, + slow_ar_config: FishSpeechSlowARConfig, + prefix: str = "fast_ar", + ) -> None: + super().__init__() + self._vllm_config = vllm_config + self.config = config + self.slow_ar_config = slow_ar_config + + self.model = FishSpeechFastARModel(config, prefix=f"{prefix}.model") + + # Codebook embeddings for Fast AR (separate from Slow AR's codebook_embeddings). + self.fast_embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + + # Output projection for residual codes. + self.fast_output = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Layer norm for Fast AR output. + self.fast_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # Projection from Slow AR hidden dim β†’ Fast AR hidden dim. + # s2-pro has same dim for both (2560) and no fast_project_in in checkpoint, + # so use identity when dims match. + if slow_ar_config.hidden_size != config.hidden_size: + self.fast_project_in = nn.Linear(slow_ar_config.hidden_size, config.hidden_size, bias=True) + else: + self.fast_project_in = nn.Identity() + + self._num_codebooks = config.num_codebooks + self._fast_dim = config.hidden_size + + # Pre-allocated buffers (lazily initialised on first forward). + self._embed_buf: torch.Tensor | None = None + self._pos_ids: torch.Tensor | None = None + self._compiled_model_fwd: object | None = None + + def _ensure_buffers(self, bsz: int, device: torch.device, dtype: torch.dtype) -> None: + max_seq = self._num_codebooks + 1 # hidden_state + num_codebooks codes + if ( + self._embed_buf is not None + and self._embed_buf.shape[0] >= bsz + and self._embed_buf.device == device + and self._embed_buf.dtype == dtype + ): + return + self._embed_buf = torch.zeros(bsz, max_seq, self._fast_dim, dtype=dtype, device=device) + self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device) + + def _setup_compile(self) -> None: + if self._compiled_model_fwd is not None: + return + # TODO: Enable torch.compile for performance. Eager for now to avoid + # potential graph-break issues during initial bring-up. + self._compiled_model_fwd = self.model.forward + + @torch.inference_mode() + def forward( + self, + slow_ar_hidden: torch.Tensor, + semantic_token_id: torch.Tensor, + *, + do_sample: bool = True, + temperature: float = 0.8, + top_k: int = 30, + top_p: float = 0.9, + ) -> torch.Tensor: + """Predict residual codebook codes 0..num_codebooks-1 autoregressively. + + Args: + slow_ar_hidden: [B, hidden_size] last hidden state from Slow AR. + semantic_token_id: [B] or [B, 1] sampled semantic token IDs (in vocab space). + + Returns: + all_codes: [B, num_codebooks] codes for codebooks 0..num_codebooks-1. + Code 0 is the semantic code (token - semantic_begin_id). + """ + bsz = int(slow_ar_hidden.shape[0]) + num_cb = self._num_codebooks + device = slow_ar_hidden.device + dtype = slow_ar_hidden.dtype + + semantic_begin = self.slow_ar_config.semantic_begin_id + semantic_end = self.slow_ar_config.semantic_end_id + codebook_size = semantic_end - semantic_begin + 1 # 4096 + # Convert vocab-space semantic token to codebook index. + # Clamp to valid range: im_end or other non-semantic tokens map to 0 (pad). + semantic_code = (semantic_token_id.reshape(bsz) - semantic_begin).clamp(min=0, max=codebook_size - 1) + + all_codes = torch.empty(bsz, num_cb, dtype=torch.long, device=device) + all_codes[:, 0] = semantic_code + + self._ensure_buffers(bsz, device, dtype) + self._setup_compile() + + embed_buf = self._embed_buf + pos_ids = self._pos_ids + model_fwd = self._compiled_model_fwd + + # Position 0: projected Slow AR hidden state. + projected = self.fast_project_in(slow_ar_hidden.reshape(bsz, -1)) + embed_buf[:bsz, 0, :] = projected + + # Position 1: embedding of semantic code. + code_embed = self.fast_embeddings(semantic_code.clamp(min=0)) + embed_buf[:bsz, 1, :] = code_embed + + use_sampling = do_sample and temperature > 0 + inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0 + + # Residual codebook size (1024) vs semantic codebook size (4096). + # The fast_output head has codebook_size (4096) outputs, but residual + # codebooks only have 1024 entries. Truncate logits for steps > 0. + residual_codebook_size = 1024 + + for step in range(1, num_cb): + seq_len = step + 1 + step_input = embed_buf[:bsz, :seq_len, :] + step_pos_ids = pos_ids[:seq_len] if bsz == 1 else pos_ids[:seq_len].repeat(bsz) + + hidden_out = model_fwd(step_input, step_pos_ids) + logits = self.fast_output(self.fast_norm(hidden_out[:, -1, :])) + + # Residual codebooks (step >= 1) only have 1024 entries. + if step >= 1: + logits = logits[:, :residual_codebook_size] + + if use_sampling: + scaled = logits * inv_temperature + if top_k > 0: + topk_vals, _ = scaled.topk(min(top_k, scaled.shape[-1]), dim=-1) + scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) + if top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(scaled, descending=True) + cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) + sorted_indices_to_remove = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= top_p + sorted_logits[sorted_indices_to_remove] = float("-inf") + scaled = sorted_logits.scatter(1, sorted_indices, sorted_logits) + probs = F.softmax(scaled, dim=-1) + next_ids = torch.multinomial(probs, num_samples=1) + else: + next_ids = logits.argmax(dim=-1, keepdim=True) + + all_codes[:, step] = next_ids.reshape(bsz) + + if step < num_cb - 1: + new_embed = self.fast_embeddings(next_ids.reshape(bsz)) + embed_buf[:bsz, step + 1, :] = new_embed + + return all_codes + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights -- Fast AR weights have already been remapped by the parent.""" + with set_current_vllm_config(self._vllm_config): + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded: set[str] = set() + + stacked_params_mapping = [ + ("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), + ] + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped not in params_dict: + continue + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded.add(mapped) + handled = True + break + + if handled: + continue + + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded.add(name) + + return loaded diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py new file mode 100644 index 000000000..9c41d87f9 --- /dev/null +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_slow_ar.py @@ -0,0 +1,669 @@ +"""Fish Speech S2 Pro -- Slow AR model (Stage 0). + +Uses vLLM's ``Qwen3Model`` as the transformer backbone. Adds: + - Multi-codebook input embedding (text + summed codebook embeddings at + semantic-token positions). + - Semantic logit masking. + - Nested Fast AR for residual codebook prediction (``talker_mtp``). + - ``preprocess`` / ``postprocess`` hooks for vLLM-omni's AR scheduler. + +Analogous to ``Qwen3TTSTalkerForConditionalGeneration`` in qwen3_tts. +""" + +from __future__ import annotations + +import dataclasses +import math +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn as nn +from transformers import AutoTokenizer +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.model_executor.models.utils import PPMissingLayer, maybe_prefix +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .configuration_fish_speech import FishSpeechConfig, FishSpeechFastARConfig, FishSpeechSlowARConfig +from .fish_speech_fast_ar import FishSpeechFastAR + +logger = init_logger(__name__) + + +def _remap_fish_speech_weights( + weights: Iterable[tuple[str, torch.Tensor]], + n_head: int, + n_local_heads: int, + head_dim: int, + fast_n_head: int, + fast_n_local_heads: int, + fast_head_dim: int, +) -> Iterable[tuple[str, torch.Tensor]]: + """Transform Fish Speech HF weight names/values to Qwen3-compatible format. + + Key transformations: + - ``wqkv`` β†’ split into ``q_proj``, ``k_proj``, ``v_proj`` + - ``wo`` β†’ ``o_proj`` + - ``w1`` β†’ ``gate_proj``, ``w3`` β†’ ``up_proj``, ``w2`` β†’ ``down_proj`` + - ``attention_norm`` β†’ ``input_layernorm`` + - ``ffn_norm`` β†’ ``post_attention_layernorm`` + - ``text_model.model.embeddings`` β†’ ``model.embed_tokens`` + - ``audio_decoder.*`` β†’ ``fast_ar.*`` with similar transforms + """ + q_size_text = n_head * head_dim + kv_size_text = n_local_heads * head_dim + q_size_fast = fast_n_head * fast_head_dim + kv_size_fast = fast_n_local_heads * fast_head_dim + + for name, tensor in weights: + # --- Text model (Slow AR) --- + if name.startswith("text_model.model."): + suffix = name[len("text_model.model.") :] + + # Embeddings + if suffix == "embeddings.weight": + yield "model.embed_tokens.weight", tensor + continue + + # Norm + if suffix == "norm.weight": + yield "model.norm.weight", tensor + continue + + # Layer weights + if suffix.startswith("layers."): + # layers.{N}.attention.wqkv.weight β†’ split into q/k/v + if ".attention.wqkv.weight" in suffix: + layer_prefix = suffix.split(".attention.wqkv.weight")[0] + q = tensor[:q_size_text, :] + k = tensor[q_size_text : q_size_text + kv_size_text, :] + v = tensor[q_size_text + kv_size_text :, :] + yield f"model.{layer_prefix}.self_attn.q_proj.weight", q + yield f"model.{layer_prefix}.self_attn.k_proj.weight", k + yield f"model.{layer_prefix}.self_attn.v_proj.weight", v + continue + + new_suffix = suffix + new_suffix = new_suffix.replace(".attention.wo.", ".self_attn.o_proj.") + new_suffix = new_suffix.replace(".attention.q_norm.", ".self_attn.q_norm.") + new_suffix = new_suffix.replace(".attention.k_norm.", ".self_attn.k_norm.") + new_suffix = new_suffix.replace(".attention_norm.", ".input_layernorm.") + new_suffix = new_suffix.replace(".feed_forward.w1.", ".mlp.gate_proj.") + new_suffix = new_suffix.replace(".feed_forward.w3.", ".mlp.up_proj.") + new_suffix = new_suffix.replace(".feed_forward.w2.", ".mlp.down_proj.") + new_suffix = new_suffix.replace(".ffn_norm.", ".post_attention_layernorm.") + yield f"model.{new_suffix}", tensor + continue + + # Fallback for any other text_model.model.* weights + yield f"model.{suffix}", tensor + continue + + # --- Audio decoder (Fast AR) --- + if name.startswith("audio_decoder."): + suffix = name[len("audio_decoder.") :] + + # Codebook embeddings (belongs to the main model, not Fast AR). + if suffix == "codebook_embeddings.weight": + yield "codebook_embeddings.weight", tensor + continue + + # Fast AR embeddings, output, norm. + if suffix == "embeddings.weight": + yield "fast_ar.fast_embeddings.weight", tensor + continue + if suffix == "output.weight": + yield "fast_ar.fast_output.weight", tensor + continue + if suffix == "norm.weight": + yield "fast_ar.fast_norm.weight", tensor + continue + + # Fast AR projection in. + if suffix.startswith("fast_project_in."): + yield f"fast_ar.fast_project_in.{suffix[len('fast_project_in.') :]}", tensor + continue + + # Fast AR layer weights. + if suffix.startswith("layers."): + if ".attention.wqkv.weight" in suffix: + layer_prefix = suffix.split(".attention.wqkv.weight")[0] + q = tensor[:q_size_fast, :] + k = tensor[q_size_fast : q_size_fast + kv_size_fast, :] + v = tensor[q_size_fast + kv_size_fast :, :] + yield f"fast_ar.model.{layer_prefix}.self_attn.q_proj.weight", q + yield f"fast_ar.model.{layer_prefix}.self_attn.k_proj.weight", k + yield f"fast_ar.model.{layer_prefix}.self_attn.v_proj.weight", v + continue + + new_suffix = suffix + new_suffix = new_suffix.replace(".attention.wo.", ".self_attn.o_proj.") + new_suffix = new_suffix.replace(".attention.q_norm.", ".self_attn.q_norm.") + new_suffix = new_suffix.replace(".attention.k_norm.", ".self_attn.k_norm.") + new_suffix = new_suffix.replace(".attention_norm.", ".input_layernorm.") + new_suffix = new_suffix.replace(".feed_forward.w1.", ".mlp.gate_proj.") + new_suffix = new_suffix.replace(".feed_forward.w3.", ".mlp.up_proj.") + new_suffix = new_suffix.replace(".feed_forward.w2.", ".mlp.down_proj.") + new_suffix = new_suffix.replace(".ffn_norm.", ".post_attention_layernorm.") + yield f"fast_ar.model.{new_suffix}", tensor + continue + + yield f"fast_ar.{suffix}", tensor + continue + + # Pass through any other weights. + yield name, tensor + + +class FishSpeechSlowARForConditionalGeneration(nn.Module): + """vLLM-AR Slow AR model for Fish Speech S2 Pro. + + Stage 0: text β†’ semantic tokens (+ residual codebook codes via Fast AR). + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + + config: FishSpeechConfig = vllm_config.model_config.hf_config # type: ignore[assignment] + self.config = config + self.text_config: FishSpeechSlowARConfig = config.text_config + self.fast_ar_config: FishSpeechFastARConfig = config.audio_decoder_config + + self._semantic_begin_id = int(config.semantic_start_token_id) + self._semantic_end_id = int(config.semantic_end_token_id) + self._audio_pad_token_id = int(config.audio_pad_token_id) + self._codebook_size = int(self.text_config.codebook_size) + self._num_codebooks = int(self.text_config.num_codebooks) + + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + self.mtp_hidden_size = int(self.text_config.hidden_size) + self.talker_mtp_output_key = "audio_codes" + + # Qwen3 transformer backbone. + self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + + # Fish Speech uses interleaved (GPT-J) RoPE, not NeoX style. + # vLLM's Qwen3Attention defaults to NeoX (is_neox_style=True). + # Replace with interleaved RoPE to match training. + self._fix_rope_style() + + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + self.text_config.vocab_size, + self.text_config.hidden_size, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(self.text_config.vocab_size) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + # Multi-codebook embedding table: codebook_size * num_codebooks entries. + self.codebook_embeddings = nn.Embedding( + self._codebook_size * self._num_codebooks, + self.text_config.hidden_size, + ) + + # Fast AR (residual codebook predictor). + predictor_compilation = dataclasses.replace(vllm_config.compilation_config) + predictor_compilation.static_forward_context = {} + self._fast_ar_vllm_config = dataclasses.replace(vllm_config, compilation_config=predictor_compilation) + from vllm.config.vllm import set_current_vllm_config as _set_cfg + + with _set_cfg(self._fast_ar_vllm_config): + self.fast_ar = FishSpeechFastAR( + vllm_config=self._fast_ar_vllm_config, + config=self.fast_ar_config, + slow_ar_config=self.text_config, + prefix="fast_ar", + ) + + # Constant logit mask: allow only semantic tokens + im_end. + vocab = int(self.text_config.vocab_size) + semantic_mask = torch.zeros((vocab,), dtype=torch.bool) + lo = self._semantic_begin_id + hi = min(self._semantic_end_id + 1, vocab) + if hi > lo: + semantic_mask[lo:hi] = True + # Also allow <|im_end|> (token 151645 in Qwen3 tokeniser). + im_end_id = 151645 + if im_end_id < vocab: + semantic_mask[im_end_id] = True + self.register_buffer("_semantic_allowed_mask", semantic_mask, persistent=False) + + # Tokeniser (lazy). + self._tokenizer = None + + def _fix_rope_style(self) -> None: + """Replace NeoX-style RoPE with interleaved (GPT-J) style. + + Fish Speech was trained with interleaved RoPE (complex-number pairs), + but vLLM's Qwen3Attention defaults to NeoX style. We rebuild the + rotary embedding with ``is_neox_style=False`` for each attention layer. + """ + from vllm.model_executor.layers.rotary_embedding import get_rope + + for layer in self.model.layers: + attn = layer.self_attn + # Extract parameters from the existing RoPE to rebuild it. + head_dim = attn.head_dim + max_position = self.text_config.max_position_embeddings + rope_params = getattr(self.text_config, "rope_scaling", None) or {} + rope_params.setdefault("rope_theta", getattr(self.text_config, "rope_theta", 1000000.0)) + attn.rotary_emb = get_rope( + head_size=head_dim, + max_position=max_position, + is_neox_style=False, + rope_parameters=rope_params, + ) + logger.info("Fixed RoPE style to interleaved (GPT-J) for %d layers", len(self.model.layers)) + + # -------------------- vLLM required hooks -------------------- + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor | IntermediateTensors: + return self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor | None: + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + return None + logits = self.logits_processor(self.lm_head, hidden_states) + if logits is None: + return None + + # Mask to semantic tokens + im_end only. + logits = logits.masked_fill(~self._semantic_allowed_mask, float("-inf")) + return logits + + # -------------------- Omni multimodal output plumbing -------------------- + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + + hidden = model_outputs + info_dicts = kwargs.get("model_intermediate_buffer") + if info_dicts is None: + info_dicts = kwargs.get("runtime_additional_information") or [] + + audio_codes_list: list[torch.Tensor] = [] + for info in info_dicts: + if not isinstance(info, dict): + continue + ac = info.get("audio_codes") + if isinstance(ac, torch.Tensor): + audio_codes_list.append(ac) + + if not audio_codes_list: + logger.debug("make_omni_output: no audio_codes found in info_dicts (len=%d)", len(info_dicts)) + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + audio_codes = torch.cat(audio_codes_list, dim=0) + span_len = int(audio_codes.shape[0]) + hidden = hidden[:span_len] + mm: dict[str, torch.Tensor] = {"audio_codes": audio_codes} + return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + + # -------------------- preprocess / postprocess -------------------- + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + additional_information = info_dict.get("additional_information") + if isinstance(additional_information, dict): + merged: dict[str, Any] = {k: v for k, v in info_dict.items() if k != "additional_information"} + for k, v in additional_information.items(): + merged.setdefault(k, v) + info_dict = merged + + span_len = int(input_ids.shape[0]) + if span_len <= 0: + return input_ids, input_embeds if input_embeds is not None else self.embed_input_ids(input_ids), {} + + if span_len > 1: + # --- Prefill --- + prompt_embeds_cpu = info_dict.get("slow_ar_prompt_embeds") + is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2 + dev = input_ids.device + + if is_first_prefill: + prompt_embeds = self._build_prefill_embeds(input_ids, info_dict) + prompt_embeds_cpu = prompt_embeds.detach().to("cpu").contiguous() + + info_update: dict[str, Any] = { + "slow_ar_prompt_embeds": prompt_embeds_cpu, + "prefill_offset": 0, + } + + take = prompt_embeds_cpu[:span_len] + if int(take.shape[0]) < span_len: + pad_n = span_len - int(take.shape[0]) + pad_embed = self.embed_input_ids( + torch.tensor([self._audio_pad_token_id], device=dev, dtype=torch.long) + ).reshape(1, -1) + take = torch.cat([take, pad_embed.detach().cpu().expand(pad_n, -1)], dim=0) + prompt_embeds = take.to(device=dev, dtype=torch.bfloat16) + info_update["prefill_offset"] = span_len + + zeros = torch.zeros( + (prompt_embeds.shape[0], self._num_codebooks), + device=dev, + dtype=torch.long, + ) + info_update["audio_codes"] = zeros + + input_ids_out = input_ids.clone() + input_ids_out[:] = self._audio_pad_token_id + return input_ids_out, prompt_embeds, info_update + + else: + # Subsequent prefill chunk. + offset = int(info_dict.get("prefill_offset", 0) or 0) + s = max(0, min(offset, int(prompt_embeds_cpu.shape[0]))) + e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0]))) + take = prompt_embeds_cpu[s:e] + if int(take.shape[0]) < span_len: + pad_n = span_len - int(take.shape[0]) + pad_embed = self.embed_input_ids( + torch.tensor([self._audio_pad_token_id], device=dev, dtype=torch.long) + ).reshape(1, -1) + take = torch.cat([take, pad_embed.detach().cpu().expand(pad_n, -1)], dim=0) + prompt_embeds = take.to(device=dev, dtype=torch.bfloat16) + + zeros = torch.zeros((prompt_embeds.shape[0], self._num_codebooks), device=dev, dtype=torch.long) + return ( + input_ids.clone().fill_(self._audio_pad_token_id), + prompt_embeds, + { + "prefill_offset": offset + span_len, + "audio_codes": zeros, + }, + ) + + # --- Decode: span_len == 1 --- + dev = input_ids.device + + last_hidden_cpu = info_dict.get("last_slow_ar_hidden") + if not isinstance(last_hidden_cpu, torch.Tensor): + # First decode step after prefill -- just embed the token directly. + logger.warning( + "preprocess decode: last_slow_ar_hidden not found (keys=%s), " + "returning plain embed (mtp_inputs will NOT be set)", + list(info_dict.keys()), + ) + embeds = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).reshape(1, -1) + return input_ids, embeds.to(dtype=torch.bfloat16), {} + + token_embed = self.embed_input_ids(input_ids.reshape(1, 1).to(torch.long)).to( + device=dev, dtype=torch.bfloat16 + ) # [1, 1, H] + + # Codebook embeddings are added by talker_mtp (using the CURRENT step's + # FastAR hidden state) instead of here (which would use the PREVIOUS step's + # codes, one step too old). Preprocess always returns plain text embed. + inputs_embeds_out = token_embed.reshape(1, -1) + + info_update = { + "mtp_inputs": ( + last_hidden_cpu.to(device=dev, dtype=torch.bfloat16).reshape(1, -1), + torch.zeros(1, self.text_config.hidden_size, device=dev, dtype=torch.bfloat16), + ), + } + return input_ids, inputs_embeds_out, info_update + + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: + if hidden_states.numel() == 0: + logger.debug("postprocess: empty hidden_states") + return {} + last = hidden_states[-1, :].detach().to("cpu").contiguous() + logger.debug("postprocess: saved last_slow_ar_hidden shape=%s", tuple(last.shape)) + return {"last_slow_ar_hidden": last} + + # -------------------- prompt construction -------------------- + + def _get_tokenizer(self): + if self._tokenizer is None: + self._tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True) + return self._tokenizer + + def _build_prefill_embeds( + self, + input_ids: torch.Tensor, + info_dict: dict[str, Any], + ) -> torch.Tensor: + """Build prefill embeddings, adding codebook embeddings at semantic positions. + + For text-only prefill (no reference audio), this is just embed_input_ids. + For voice cloning, reference codes are embedded with codebook offsets. + """ + dev = input_ids.device + # Basic text embeddings. + base_embeds = self.embed_input_ids(input_ids.reshape(1, -1).to(torch.long)) # [1, T, H] + + # Check for reference codebook codes (for voice cloning). + ref_codes = info_dict.get("ref_codes") + if not isinstance(ref_codes, torch.Tensor) or ref_codes.numel() == 0: + return base_embeds.squeeze(0).to(dtype=torch.bfloat16) + + # ref_codes: [T_ref, num_codebooks] -- codebook codes for reference audio positions. + ref_codes = ref_codes.to(device=dev, dtype=torch.long) + ref_positions = info_dict.get("ref_positions") + if not isinstance(ref_positions, torch.Tensor): + return base_embeds.squeeze(0).to(dtype=torch.bfloat16) + + ref_positions = ref_positions.to(device=dev, dtype=torch.long).reshape(-1) + seq_len = int(input_ids.shape[0]) + codebook_sum = torch.zeros_like(base_embeds) # [1, T, H] + + for pos_idx in range(int(ref_positions.shape[0])): + pos = int(ref_positions[pos_idx].item()) + if pos < 0 or pos >= seq_len: + continue + for cb_idx in range(min(int(ref_codes.shape[1]), self._num_codebooks)): + code = ref_codes[pos_idx, cb_idx].clamp(min=0) + code_with_offset = code + cb_idx * self._codebook_size + emb = self.codebook_embeddings(code_with_offset.unsqueeze(0)) + codebook_sum[0, pos, :] += emb.squeeze(0).to(dtype=base_embeds.dtype) + + result = base_embeds + codebook_sum + return result.squeeze(0).to(dtype=torch.bfloat16) + + # -------------------- GPU-side MTP fast-path -------------------- + + @torch.inference_mode() + def talker_mtp( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor, + last_talker_hidden: torch.Tensor, + text_step: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """GPU fast-path: run Fast AR to predict residual codebook codes. + + Returns (inputs_embeds, audio_codes). + + The embedding is: text_embed(token) + sum(codebook_embed(code_i + i * codebook_size)) + where codes come from FastAR(last_talker_hidden). + + This matches the reference Fish Speech inference flow: + - At step t, the Slow AR embedding includes codes from FastAR(hidden_{t-1}) + - last_talker_hidden IS hidden_{t-1} (from postprocess of the previous step) + - Preprocess provides the plain text embed; we add codebook embeddings here + """ + bsz = int(input_ids.shape[0]) + dev = input_embeds.device + + input_ids = input_ids.reshape(bsz, 1).to(dtype=torch.long, device=dev) + past_hidden = last_talker_hidden.reshape(bsz, -1).to(dtype=torch.bfloat16, device=dev) + + # Run Fast AR to predict all num_codebooks codes. + audio_codes = self.fast_ar( + slow_ar_hidden=past_hidden, + semantic_token_id=input_ids.reshape(bsz), + do_sample=True, + temperature=0.8, + top_k=30, + top_p=0.9, + ) # [B, num_codebooks] + + # Add codebook embeddings to the input embedding (from preprocess). + # This ensures the Slow AR sees codes from FastAR(hidden_{t-1}). + inputs_embeds_out = input_embeds.reshape(bsz, -1).clone() + + for b in range(bsz): + token_id = int(input_ids[b, 0].item()) + is_semantic = self._semantic_begin_id <= token_id <= self._semantic_end_id + if is_semantic: + codes = audio_codes[b] # [num_codebooks] + codebook_sum = torch.zeros(self.text_config.hidden_size, device=dev, dtype=torch.bfloat16) + for i in range(self._num_codebooks): + code_with_offset = codes[i].clamp(min=0) + i * self._codebook_size + emb = self.codebook_embeddings(code_with_offset.unsqueeze(0)) + codebook_sum += emb.squeeze(0).to(dtype=torch.bfloat16) + + # Normalize by sqrt(num_codebooks + 1) as in the reference model + # (scale_codebook_embeddings=True for fish_qwen3_omni). + inputs_embeds_out[b] = (inputs_embeds_out[b] + codebook_sum) / math.sqrt(self._num_codebooks + 1) + + return inputs_embeds_out, audio_codes.to(dtype=torch.long) + + # -------------------- Prompt length estimation -------------------- + + @staticmethod + def estimate_prompt_len_from_additional_information( + additional_information: dict[str, Any] | None, + **kwargs: Any, + ) -> int: + """Estimate prompt length for placeholder allocation.""" + info = additional_information or {} + text = info.get("text", [""])[0] if isinstance(info.get("text"), list) else info.get("text", "") + # Conservative estimate: tokenize text length + overhead. + return max(2, len(str(text)) // 2 + 64) + + # -------------------- Weight loading -------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights with Fish Speech β†’ Qwen3 format transformation. + + Transforms weight names (wqkv β†’ q/k/v split, w1/w2/w3 β†’ gate/up/down, + etc.) and routes to the correct sub-modules. + """ + n_head = self.text_config.num_attention_heads + n_local_heads = self.text_config.num_key_value_heads + head_dim = self.text_config.head_dim + fast_n_head = self.fast_ar_config.num_attention_heads + fast_n_local_heads = self.fast_ar_config.num_key_value_heads + fast_head_dim = self.fast_ar_config.head_dim + + remapped = _remap_fish_speech_weights( + weights, + n_head, + n_local_heads, + head_dim, + fast_n_head, + fast_n_local_heads, + fast_head_dim, + ) + + # Qwen3Model uses stacked_params_mapping for q/k/v β†’ qkv_proj + # and gate/up β†’ gate_up_proj. Feed the remapped weights through + # the standard Qwen3 loading path. + stacked_params_mapping = [ + ("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 remapped: + if "rotary_emb.inv_freq" in name: + continue + + # Handle tied embeddings β†’ lm_head. + if name == "model.embed_tokens.weight" and self.text_config.tie_word_embeddings: + # Also load into lm_head if present. + lm_key = "lm_head.weight" + if lm_key in params_dict: + p = params_dict[lm_key] + wl = getattr(p, "weight_loader", default_weight_loader) + wl(p, loaded_weight) + loaded_params.add(lm_key) + + # Try stacked params mapping (q/k/v β†’ qkv_proj, gate/up β†’ gate_up_proj). + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + mapped = name.replace(weight_name, param_name) + if mapped not in params_dict: + continue + param = params_dict[mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if weight_loader == default_weight_loader: + weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(mapped) + handled = True + break + + if handled: + continue + + # Direct parameter mapping. + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + logger.info("Loaded %d weights for FishSpeechSlowARForConditionalGeneration", len(loaded_params)) + + # Truncate RoPE cos/sin caches to bf16 precision to match training. + # Without this, f32 RoPE values cause logit divergence and premature EOS. + truncated = 0 + for module in self.modules(): + if hasattr(module, "cos_sin_cache") and isinstance(module.cos_sin_cache, torch.Tensor): + cache = module.cos_sin_cache + module.cos_sin_cache = cache.to(torch.bfloat16).to(cache.dtype) + truncated += 1 + if truncated: + logger.info("Truncated %d RoPE cos_sin_cache buffers to bf16 precision", truncated) + + return loaded_params diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index b886aec69..5a0ae9965 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -1054,6 +1054,79 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc ] +class HunyuanImage3RotaryEmbedding(nn.Module): + """Custom interleaved 2D Rotary Embedding for HunyuanImage3. + + The original HunyuanImage3 ``build_2d_rope`` interleaves y (height) and + x (width) positions across consecutive frequencies:: + + theta = inv_freq.reshape(n_elem // 4, 2) + idx_theta[2k] = y_pos * theta[k, 0] # even freq -> y + idx_theta[2k+1] = x_pos * theta[k, 1] # odd freq -> x + + vLLM's standard ``MRotaryEmbedding`` instead assigns *contiguous* blocks + of frequencies to each position dimension, producing different encodings + for the same positions. This class re-implements the original interleaved + pattern so that pre-trained weights work correctly. + + Reference: https://kexue.fm/archives/10352 + """ + + def __init__(self, head_dim: int, rope_theta: float = 10000.0) -> None: + super().__init__() + self.head_dim = head_dim + assert head_dim % 4 == 0, f"head_dim must be divisible by 4, got {head_dim}" + inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if positions.dim() == 2 and positions.shape[0] == 3: + y_pos = positions[1].float() + x_pos = positions[2].float() + else: + # 1D fallback: both dims get the same position β†’ standard RoPE + y_pos = positions.float() + x_pos = positions.float() + + num_tokens = y_pos.shape[0] + dtype = query.dtype + query_shape = query.shape + key_shape = key.shape + + inv_freq = self.inv_freq.to(device=y_pos.device, dtype=torch.float32) + + inv_freq_y = inv_freq[0::2] # even indices -> y + inv_freq_x = inv_freq[1::2] # odd indices -> x + + y_freqs = y_pos.unsqueeze(-1) * inv_freq_y.unsqueeze(0) + x_freqs = x_pos.unsqueeze(-1) * inv_freq_x.unsqueeze(0) + + # Interleave: [y*ΞΈβ‚€, x*θ₁, y*ΞΈβ‚‚, x*θ₃, ...] + freqs = torch.stack([y_freqs, x_freqs], dim=-1).reshape(num_tokens, -1) + + emb = torch.cat([freqs, freqs], dim=-1) + cos = emb.cos().to(dtype).unsqueeze(1) + sin = emb.sin().to(dtype).unsqueeze(1) + + query = query.view(num_tokens, -1, self.head_dim) + key = key.view(num_tokens, -1, self.head_dim) + + query = query * cos + self._rotate_half(query) * sin + key = key * cos + self._rotate_half(key) * sin + + return query.view(query_shape), key.view(key_shape) + + def _rotate_half(self, x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + @MULTIMODAL_REGISTRY.register_processor( HunyuanImage3MultiModalProcessor, info=HunyuanImage3ProcessingInfo, @@ -1153,6 +1226,38 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo self._mrope_joint_img_sep_token_id = tokenizer.convert_tokens_to_ids("") self._mrope_max_num_patches = config.vit_processor.get("max_num_patches", 729) + self._replace_rotary_embeddings() + + def _replace_rotary_embeddings(self): + """Replace vLLM's standard MRotaryEmbedding with the custom + interleaved 2D RoPE that matches the original HunyuanImage3 model.""" + rope_theta = getattr(self.config, "rope_theta", 10000.0) + head_dim = getattr( + self.config, + "head_dim", + getattr( + self.config, + "attention_head_dim", + self.config.hidden_size // self.config.num_attention_heads, + ), + ) + custom_rope = HunyuanImage3RotaryEmbedding( + head_dim=head_dim, + rope_theta=rope_theta, + ) + replaced = 0 + for layer in self.model.layers: + if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "rotary_emb"): + layer.self_attn.rotary_emb = custom_rope + replaced += 1 + logger.info( + "Replaced %d rotary embeddings with HunyuanImage3RotaryEmbedding " + "(interleaved 2D RoPE, head_dim=%d, rope_theta=%.1f)", + replaced, + head_dim, + rope_theta, + ) + def _parse_and_validate_image_input( self, **kwargs: dict[str, Any], @@ -1414,19 +1519,20 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo """Compute mRoPE positions for HunyuanImage-3. Maps the original model's build_2d_rope logic into vLLM's 3-dim - mRoPE position tensor [3, seq_len] where dim-0 is temporal (unused, - kept equal to 1D), dim-1 is height, and dim-2 is width. + mRoPE position tensor ``[3, seq_len]`` where dim-1 is height and + dim-2 is width. dim-0 is unused (temporal) and kept equal to 1D. - For text tokens and auxiliary image tokens (timestep, ViT): + For text tokens and auxiliary image tokens (timestep, separators): All three dims get the same flat 1D position id. - For VAE image tokens: - dim-0 (T): flat 1D position id at the image start + For VAE / ViT image tokens: + dim-0 (T): flat 1D position id at the region start dim-1 (H): 2D y-position using build_2d_rope centering dim-2 (W): 2D x-position using build_2d_rope centering """ - # Extract per-image VAE grid dims from mm_features + # Extract per-image VAE and ViT grid dims from mm_features vae_grids: list[tuple[int, int]] = [] + vit_grids: list[tuple[int, int]] = [] if mm_features is not None: for mm_feature in mm_features: mm_item = mm_feature.data @@ -1437,20 +1543,22 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo if vae_hw is not None: grid = vae_hw.tolist() vae_grids.append((int(grid[0]), int(grid[1]))) + vit_hw = mm_input.get("vit_spatial_shapes") + if vit_hw is not None: + grid = vit_hw.tolist() + vit_grids.append((int(grid[0]), int(grid[1]))) - # Identify image token ids (cached from __init__) img_token_id = self._mrope_img_token_id boi_token_id = self._mrope_boi_token_id eoi_token_id = self._mrope_eoi_token_id joint_img_sep_token_id = self._mrope_joint_img_sep_token_id - max_num_patches = self._mrope_max_num_patches # Build position arrays - t_pos = [] # temporal (same as 1D for this model) - h_pos = [] # height - w_pos = [] # width + t_pos: list[int] = [] # temporal (same as 1D for this model) + h_pos: list[int] = [] # height + w_pos: list[int] = [] # width - pos = 0 # current 1D position counter + pos = 0 image_idx = 0 i = 0 n = len(input_tokens) @@ -1462,7 +1570,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo # Found start of image block. # Structure: *timestep *vae # *vit - # Assign a flat position + # token t_pos.append(pos) h_pos.append(pos) w_pos.append(pos) @@ -1493,16 +1601,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo pos += 1 i += 1 - # VAE tokens: get grid dims + # VAE tokens: 2D grid positions if image_idx < len(vae_grids): vae_h, vae_w = vae_grids[image_idx] else: vae_h, vae_w = 0, 0 - image_idx += 1 - - # Assign 2D positions to VAE tokens using build_2d_rope - # centering logic - L = pos # position at start of VAE region + L = pos wh = vae_w * vae_h beta_y = L + (wh - vae_h) / 2 beta_x = L + (wh - vae_w) / 2 @@ -1510,12 +1614,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo for row in range(vae_h): for col in range(vae_w): if i < n and input_tokens[i] == img_token_id: - t_pos.append(L) # temporal stays flat + t_pos.append(L) h_pos.append(int(beta_y + row)) w_pos.append(int(beta_x + col)) i += 1 - pos = L + wh # advance past VAE region + pos = L + wh # token if i < n and input_tokens[i] == joint_img_sep_token_id: @@ -1525,15 +1629,34 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo pos += 1 i += 1 - # ViT tokens (max_num_patches tokens) β€” flat 1D - vit_consumed = 0 - while i < n and input_tokens[i] == img_token_id and vit_consumed < max_num_patches: + # ViT tokens: 2D grid positions + if image_idx < len(vit_grids): + vit_h, vit_w = vit_grids[image_idx] + else: + vit_h, vit_w = 0, 0 + L = pos + wh = vit_w * vit_h + beta_y = L + (wh - vit_h) / 2 + beta_x = L + (wh - vit_w) / 2 + + for row in range(vit_h): + for col in range(vit_w): + if i < n and input_tokens[i] == img_token_id: + t_pos.append(L) + h_pos.append(int(beta_y + row)) + w_pos.append(int(beta_x + col)) + i += 1 + + pos = L + wh + + # Remaining ViT padding tokens (when max_num_patches > + # vit_h * vit_w) get flat 1D positions. + while i < n and input_tokens[i] == img_token_id: t_pos.append(pos) h_pos.append(pos) w_pos.append(pos) pos += 1 i += 1 - vit_consumed += 1 # token if i < n and input_tokens[i] == eoi_token_id: @@ -1543,6 +1666,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo pos += 1 i += 1 + image_idx += 1 else: # Regular text token β€” flat 1D position t_pos.append(pos) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 9bf2b4fdb..f708b6ff9 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -129,6 +129,17 @@ _OMNI_MODELS = { "hunyuan_image3", "HunyuanImage3ForConditionalGeneration", ), + ## fish_speech (Fish Speech S2 Pro) + "FishSpeechSlowARForConditionalGeneration": ( + "fish_speech", + "fish_speech_slow_ar", + "FishSpeechSlowARForConditionalGeneration", + ), + "FishSpeechDACDecoder": ( + "fish_speech", + "fish_speech_dac_decoder", + "FishSpeechDACDecoder", + ), } diff --git a/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml b/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml new file mode 100644 index 000000000..7f5ba1202 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml @@ -0,0 +1,97 @@ +async_chunk: true +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 4 + engine_args: + model_stage: fish_speech_slow_ar + model_arch: FishSpeechSlowARForConditionalGeneration + hf_overrides: + architectures: [FishSpeechSlowARForConditionalGeneration] + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: false + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.4 + distributed_executor_backend: "mp" + max_num_batched_tokens: 2048 + max_model_len: 32768 + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.fish_speech.slow_ar_to_dac_decoder_async_chunk + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.8 + top_k: 30 + top_p: 0.9 + max_tokens: 200 + seed: 42 + detokenize: false + repetition_penalty: 1.0 + # <|im_end|> token -- stop when model emits end-of-turn. + stop_token_ids: [151645] + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: dac_decoder + model_arch: FishSpeechDACDecoder + hf_overrides: + architectures: [FishSpeechDACDecoder] + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.2 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 32768 + engine_input_source: [0] + final_output: true + final_output_type: audio + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + detokenize: true + repetition_penalty: 1.0 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + # Chunk sizes for streaming -- ~21 Hz codec. + # 25 frames β‰ˆ 1.16s of audio at 21.5 Hz. + codec_chunk_frames: 25 + codec_left_context_frames: 25 + initial_codec_chunk_frames: 0 + + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml index b4b918288..e8a603af4 100644 --- a/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml @@ -22,6 +22,10 @@ stage_args: max_num_batched_tokens: 32768 tensor_parallel_size: 8 pipeline_parallel_size: 1 + hf_overrides: + rope_parameters: + mrope_section: [0, 32, 32] + rope_type: default is_comprehension: true final_output: true final_output_type: text diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml index 095aafe08..570ed710c 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts.yaml @@ -15,7 +15,7 @@ stage_args: scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler enforce_eager: false trust_remote_code: true - async_scheduling: false + async_scheduling: true enable_prefix_caching: false engine_output_type: latent gpu_memory_utilization: 0.3 @@ -50,7 +50,7 @@ stage_args: scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true - async_scheduling: false + async_scheduling: true enable_prefix_caching: false engine_output_type: audio gpu_memory_utilization: 0.2 diff --git a/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml b/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml index d73739109..5236668dc 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml @@ -19,7 +19,7 @@ stage_args: scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler enforce_eager: false trust_remote_code: true - async_scheduling: false + async_scheduling: true enable_prefix_caching: false engine_output_type: latent gpu_memory_utilization: 0.3 @@ -54,7 +54,7 @@ stage_args: scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true - async_scheduling: false + async_scheduling: true enable_prefix_caching: false engine_output_type: audio gpu_memory_utilization: 0.2 diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py new file mode 100644 index 000000000..5be70bc18 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py @@ -0,0 +1,151 @@ +"""Stage input processor for Fish Speech S2 Pro: Slow AR β†’ DAC Decoder.""" + +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_NUM_CODEBOOKS = 10 # 1 semantic + 9 residual + + +def _extract_last_frame(pooling_output: dict[str, Any]) -> torch.Tensor | None: + """Extract the last frame of audio codes from the pooling output.""" + audio_codes = pooling_output.get("audio_codes") + if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: + return None + if audio_codes.ndim == 2: + frame = audio_codes[-1] + if frame.numel() == 0 or not bool(frame.any().item()): + return None + return frame.to(torch.long).reshape(-1) + if audio_codes.ndim == 1: + return audio_codes.to(torch.long).reshape(-1) + raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}") + + +def slow_ar_to_dac_decoder( + stage_list: list[Any], + engine_input_source: list[int], + prompt: Any = None, + requires_multimodal_data: bool = False, +) -> list[Any]: + """Non-async processor: wait for Slow AR to finish, then pass all codes to DAC decoder.""" + from vllm_omni.inputs.data import OmniTokensPrompt + from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _validate_stage_inputs + + slow_ar_outputs = _validate_stage_inputs(stage_list, engine_input_source) + dac_inputs: list[OmniTokensPrompt] = [] + + for output in slow_ar_outputs: + out = output.outputs[0] + # audio_codes shape: [num_frames, num_codebooks] + audio_codes = out.multimodal_output["audio_codes"].to(torch.long) + # Filter zero-padded frames. + valid_mask = audio_codes.any(dim=1) + audio_codes = audio_codes[valid_mask] + # Codebook-major flat: [num_codebooks * num_frames] + codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist() + dac_inputs.append( + OmniTokensPrompt( + prompt_token_ids=codec_codes, + multi_modal_data=None, + mm_processor_kwargs=None, + additional_information=None, + ) + ) + return dac_inputs + + +def slow_ar_to_dac_decoder_async_chunk( + transfer_manager: Any, + pooling_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> dict[str, Any] | None: + """Async streaming processor: emit code chunks as they are produced. + + Accumulates per-step codes and emits fixed-size chunks with left context + overlap for smooth audio transitions, analogous to + ``talker2code2wav_async_chunk`` in Qwen3 TTS. + """ + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + + if isinstance(pooling_output, dict): + frame = _extract_last_frame(pooling_output) + if frame is not None: + codec_codes = frame.cpu().tolist() + transfer_manager.code_prompt_token_ids[request_id].append(codec_codes) + elif not finished: + return None + + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + chunk_size = int(cfg.get("codec_chunk_frames", 25)) + left_context_size_config = int(cfg.get("codec_left_context_frames", 25)) + initial_chunk_size = int(cfg.get("initial_codec_chunk_frames", 0)) + + # Per-request override. + additional_information = getattr(request, "additional_information", None) + if ( + additional_information is not None + and hasattr(additional_information, "entries") + and "initial_codec_chunk_frames" in additional_information.entries + ): + entry = additional_information.entries["initial_codec_chunk_frames"] + if entry.list_data is not None and len(entry.list_data) == 1: + initial_chunk_size = int(entry.list_data[0]) + + if chunk_size <= 0 or left_context_size_config < 0 or initial_chunk_size < 0: + raise ValueError( + f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, " + f"codec_left_context_frames={left_context_size_config}, " + f"initial_codec_chunk_frames={initial_chunk_size}" + ) + if initial_chunk_size > chunk_size: + initial_chunk_size = chunk_size + + length = len(transfer_manager.code_prompt_token_ids[request_id]) + + if length <= 0: + if finished: + return { + "code_predictor_codes": [], + "finished": torch.tensor(True, dtype=torch.bool), + } + return None + + in_initial_phase = initial_chunk_size > 0 and length <= chunk_size + + if in_initial_phase: + already_sent = transfer_manager.put_req_chunk[request_id] * initial_chunk_size + pending = length - already_sent + if pending <= 0: + return None + if pending < initial_chunk_size and not finished: + return None + context_length = min(pending, initial_chunk_size) + left_context_size = max(0, length - context_length) + window_frames = transfer_manager.code_prompt_token_ids[request_id][:length] + else: + initial_coverage = (chunk_size // initial_chunk_size) * initial_chunk_size if initial_chunk_size > 0 else 0 + adjusted = length - initial_coverage + chunk_length = adjusted % chunk_size + if chunk_length != 0 and not finished: + return None + context_length = chunk_length if chunk_length != 0 else chunk_size + end_index = min(length, left_context_size_config + context_length) + left_context_size = max(0, int(end_index - context_length)) + window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] + + # Pack into codebook-major flat codes. + code_predictor_codes = torch.tensor(window_frames).transpose(0, 1).reshape(-1).tolist() + + return { + "code_predictor_codes": code_predictor_codes, + "left_context_size": left_context_size, + "finished": torch.tensor(finished, dtype=torch.bool), + } diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index f90c69e3a..314cb3219 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum +from typing import Any import torch from vllm.platforms import Platform @@ -52,6 +53,16 @@ class OmniPlatform(Platform): def get_default_stage_config_path(cls) -> str: raise NotImplementedError + @classmethod + def get_diffusion_model_impl_qualname(cls, op_name: str) -> str: + if op_name == "hunyuan_fused_moe": + return "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault" + raise NotImplementedError(f"Unsupported diffusion model op: {op_name}") + + @classmethod + def prepare_diffusion_op_runtime(cls, op_name: str, **kwargs: Any) -> None: + return None + @classmethod def get_diffusion_attn_backend_cls( cls, diff --git a/vllm_omni/platforms/npu/models/__init__.py b/vllm_omni/platforms/npu/models/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/platforms/npu/models/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py b/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py new file mode 100644 index 000000000..46c76a652 --- /dev/null +++ b/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import torch +import vllm.forward_context as _vllm_fc +from vllm.config import VllmConfig +from vllm.distributed import get_ep_group +from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size +from vllm.distributed.parallel_state import ( + init_model_parallel_group as vllm_init_model_parallel_group, +) +from vllm_ascend.ascend_forward_context import MoECommType +from vllm_ascend.ops.fused_moe.fused_moe import AscendSharedFusedMoE +from vllm_ascend.ops.fused_moe.moe_comm_method import _MoECommMethods +from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type + +from vllm_omni.diffusion.distributed.parallel_state import ( + get_data_parallel_world_size, + get_world_group, +) +from vllm_omni.diffusion.forward_context import get_forward_context as omni_get_ctx + + +def _init_mc2_group_for_diffusion( + world_size: int, + data_parallel_size: int, + tensor_parallel_size: int, + backend: str, + local_rank: int, +) -> None: + import vllm_ascend.distributed.parallel_state as vllm_ascend_parallel_state + + if getattr(vllm_ascend_parallel_state, "_MC2", None) is not None: + return + all_ranks = torch.arange(world_size).reshape(-1, data_parallel_size * tensor_parallel_size) + group_ranks = all_ranks.unbind(0) + group_ranks = [x.tolist() for x in group_ranks] + + vllm_ascend_parallel_state._MC2 = vllm_init_model_parallel_group( + group_ranks, + local_rank, + backend, + group_name="mc2", + ) + + +def _select_moe_comm_method(vllm_config: VllmConfig) -> MoECommType | None: + soc_version = get_ascend_device_type() + if not vllm_config.parallel_config.enable_expert_parallel or get_ep_group().world_size == 1: + moe_comm_type = MoECommType.ALLGATHER + elif soc_version in {AscendDeviceType.A2}: + moe_comm_type = MoECommType.ALLGATHER + elif soc_version in {AscendDeviceType.A3}: + moe_comm_type = MoECommType.ALLTOALL + elif soc_version in {AscendDeviceType._310P}: + moe_comm_type = MoECommType.ALLGATHER + elif soc_version in {AscendDeviceType.A5}: + moe_comm_type = MoECommType.ALLTOALL + else: + raise ValueError(f"Unsupported soc_version: {soc_version}") + return moe_comm_type + + +def prepare_hunyuan_fused_moe_runtime() -> None: + world_size = torch.distributed.get_world_size() + data_parallel_size = get_data_parallel_world_size() + tensor_parallel_size = get_tensor_model_parallel_world_size() + backend = torch.distributed.get_backend(get_world_group().device_group) + local_rank = get_world_group().local_rank + _init_mc2_group_for_diffusion( + world_size=world_size, + data_parallel_size=data_parallel_size, + tensor_parallel_size=tensor_parallel_size, + backend=backend, + local_rank=local_rank, + ) + + if not hasattr(_vllm_fc.ForwardContext, "moe_comm_method"): + _vllm_fc.ForwardContext.__annotations__["in_profile_run"] = bool + _vllm_fc.ForwardContext.in_profile_run = False + + _vllm_fc.ForwardContext.moe_comm_type = _select_moe_comm_method(vllm_config=omni_get_ctx().vllm_config) + _vllm_fc.ForwardContext.moe_comm_method = _MoECommMethods.get(_vllm_fc.ForwardContext.moe_comm_type) + _vllm_fc.ForwardContext.flash_comm_v1_enabled = False + + +class AscendHunyuanFusedMoE(AscendSharedFusedMoE): + def __init__(self, *, prefix: str = "", **kwargs: Any) -> None: + super().__init__(prefix=prefix, **kwargs) + self._prefix = prefix + self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True) + + def _initialize_kernel_hook(self, module: Any, args: Any, kwargs: Any) -> None: + if self.quant_method: + self.quant_method.process_weights_after_loading(self) + self._init_hook_handle.remove() + + def forward(self, hidden_states: Any, router_logits: Any) -> Any: + return super().forward(hidden_states, router_logits) + + def __del__(self): + import vllm_ascend.distributed.parallel_state as vllm_ascend_parallel_state + + if vllm_ascend_parallel_state._MC2: + vllm_ascend_parallel_state._MC2.destroy() + vllm_ascend_parallel_state._MC2 = None diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 3c2495c3d..bda4e4f61 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + import torch from vllm.logger import init_logger from vllm_ascend.platform import NPUPlatform @@ -33,6 +35,23 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): def get_default_stage_config_path(cls) -> str: return "vllm_omni/platforms/npu/stage_configs" + @classmethod + def get_diffusion_model_impl_qualname(cls, op_name: str) -> str: + if op_name == "hunyuan_fused_moe": + return "vllm_omni.platforms.npu.models.hunyuan_fused_moe.AscendHunyuanFusedMoE" + return super().get_diffusion_model_impl_qualname(op_name) + + @classmethod + def prepare_diffusion_op_runtime(cls, op_name: str, **kwargs: Any) -> None: + if op_name != "hunyuan_fused_moe": + return + + from vllm_omni.platforms.npu.models.hunyuan_fused_moe import ( + prepare_hunyuan_fused_moe_runtime, + ) + + prepare_hunyuan_fused_moe_runtime() + @classmethod def get_diffusion_attn_backend_cls( cls, diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 136a680a3..59b23f914 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -14,6 +14,9 @@ _CLASS_TO_MODULE: dict[str, str] = { "Mammothmoda2Qwen2_5_VLConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", "Mammothmoda2Qwen2_5_VLTextConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", "Mammothmoda2Qwen2_5_VLVisionConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", + "FishSpeechConfig": "vllm_omni.transformers_utils.configs.fish_speech", + "FishSpeechSlowARConfig": "vllm_omni.transformers_utils.configs.fish_speech", + "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", } __all__ = [ @@ -21,6 +24,9 @@ __all__ = [ "Mammothmoda2Qwen2_5_VLConfig", "Mammothmoda2Qwen2_5_VLTextConfig", "Mammothmoda2Qwen2_5_VLVisionConfig", + "FishSpeechConfig", + "FishSpeechSlowARConfig", + "FishSpeechFastARConfig", ] @@ -39,4 +45,5 @@ def __dir__(): # Eagerly import all config modules so their AutoConfig.register() side-effects # run as soon as `vllm_omni.transformers_utils.configs` is imported. +from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/fish_speech.py b/vllm_omni/transformers_utils/configs/fish_speech.py new file mode 100644 index 000000000..44b5a67eb --- /dev/null +++ b/vllm_omni/transformers_utils/configs/fish_speech.py @@ -0,0 +1,24 @@ +"""Fish Speech S2 Pro config registration with transformers AutoConfig. + +Registers FishSpeechConfig (model_type="fish_qwen3_omni") and sub-configs +so that ``AutoConfig.from_pretrained("fishaudio/s2-pro")`` returns the +correct config class. +""" + +from transformers import AutoConfig + +from vllm_omni.model_executor.models.fish_speech.configuration_fish_speech import ( + FishSpeechConfig, + FishSpeechFastARConfig, + FishSpeechSlowARConfig, +) + +AutoConfig.register("fish_qwen3_omni", FishSpeechConfig) +AutoConfig.register("fish_qwen3", FishSpeechSlowARConfig) +AutoConfig.register("fish_qwen3_audio_decoder", FishSpeechFastARConfig) + +__all__ = [ + "FishSpeechConfig", + "FishSpeechSlowARConfig", + "FishSpeechFastARConfig", +]