diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index e02d49472..01897db96 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -20,19 +20,7 @@ steps: - label: "Simple Unit Test" depends_on: image-build commands: - - | - pytest -v -s \ - tests/entrypoints/ \ - tests/diffusion/cache/ \ - tests/diffusion/lora/ \ - tests/model_executor/models/qwen2_5_omni/test_audio_length.py \ - tests/worker/ \ - tests/distributed/omni_connectors/test_kv_flow.py \ - --cov=vllm_omni \ - --cov-branch \ - --cov-report=term-missing \ - --cov-report=html \ - --cov-report=xml + - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml" agents: queue: "gpu_1_queue" plugins: @@ -118,7 +106,7 @@ steps: timeout_in_minutes: 15 depends_on: image-build commands: - - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py + - pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4' agents: queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU plugins: @@ -135,7 +123,7 @@ steps: timeout_in_minutes: 20 depends_on: image-build commands: - - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py + - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model agents: queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index fac5c7268..773748c17 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -44,7 +44,7 @@ steps: - export GPU_ARCHS=gfx942 - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py + - pytest -s -v -m 'core_model and cache and diffusion and not distributed_rocm and MI325' - label: "Diffusion Sequence Parallelism Test" timeout_in_minutes: 20 @@ -56,7 +56,7 @@ steps: - export GPU_ARCHS=gfx942 - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py + - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model - label: "Diffusion Tensor Parallelism Test" timeout_in_minutes: 20 diff --git a/.buildkite/test-nightly.yaml b/.buildkite/test-nightly.yaml index 27fdc8378..d985c36f6 100644 --- a/.buildkite/test-nightly.yaml +++ b/.buildkite/test-nightly.yaml @@ -10,11 +10,13 @@ steps: queue: "cpu_queue_premerge" - label: "Omni Model Test with H100" - timeout_in_minutes: 180 + timeout_in_minutes: 60 depends_on: image-build + if: build.env("NIGHTLY") == "1" commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py + - pytest -s -v tests/examples/online_serving/test_qwen3_omni.py agents: queue: "mithril-h100-pool" plugins: @@ -43,3 +45,24 @@ steps: hostPath: path: /mnt/hf-cache type: DirectoryOrCreate + + + + - label: "Omni Model Test" + timeout_in_minutes: 60 + depends_on: image-build + commands: + - export VLLM_LOGGING_LEVEL=DEBUG + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py + agents: + queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index e84b399df..2d76c69b2 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -7,6 +7,7 @@ on: paths-ignore: - 'docs/**' - 'examples/**' + - 'apps/**' - '.github/**' - '.buildkite/**' - '**.md' @@ -20,6 +21,7 @@ on: paths-ignore: - 'docs/**' - 'examples/**' + - 'apps/**' - '.github/**' - '.buildkite/**' - '*.yml' diff --git a/.github/workflows/comfyui-publish.yml b/.github/workflows/comfyui-publish.yml new file mode 100644 index 000000000..c282bcada --- /dev/null +++ b/.github/workflows/comfyui-publish.yml @@ -0,0 +1,27 @@ +name: ComfyUI Integration -- Publish to Comfy registry +on: + workflow_dispatch: + push: + tags: + - "*" + paths: + - "apps/ComfyUI-vLLM-Omni/**" + +permissions: + issues: write + +defaults: + run: + working-directory: apps/ComfyUI-vLLM-Omni + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + steps: + - name: ♻️ Check out code + uses: actions/checkout@v4 + - name: 📦 Publish Custom Node + uses: Comfy-Org/publish-node-action@main + with: + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/.github/workflows/comfyui-validate.yml b/.github/workflows/comfyui-validate.yml new file mode 100644 index 000000000..06b9b1f91 --- /dev/null +++ b/.github/workflows/comfyui-validate.yml @@ -0,0 +1,40 @@ +name: ComfyUI Integration -- Validate backwards compatibility + +on: + pull_request: + branches: + - master + - main + paths: + - "apps/ComfyUI-vLLM-Omni/**" + +defaults: + run: + working-directory: apps/ComfyUI-vLLM-Omni + +jobs: + check-base-path: + runs-on: ubuntu-latest + outputs: + exists: ${{ steps.check.outputs.exists }} + defaults: + run: + working-directory: . + steps: + - uses: actions/checkout@v4 + with: + ref: ${{ github.event.pull_request.base.sha }} + fetch-depth: 1 + - id: check + run: | + if [ -d "apps/ComfyUI-vLLM-Omni" ]; then + echo "exists=true" >> "$GITHUB_OUTPUT" + else + echo "exists=false" >> "$GITHUB_OUTPUT" + fi + validate: + needs: check-base-path + if: needs.check-base-path.outputs.exists == 'true' + runs-on: ubuntu-latest + steps: + - uses: comfy-org/node-diff@main diff --git a/.gitignore b/.gitignore index f7b771bc7..7b0a9c368 100644 --- a/.gitignore +++ b/.gitignore @@ -153,6 +153,14 @@ dmypy.json # Cython debug symbols cython_debug/ +# Claude +CLAUDE.md +.claude/ + +# Codex +AGENTS.md +.codex/ + # cursor .cursor/ @@ -214,6 +222,7 @@ data/ datasets/ *.csv *.json +!apps/ComfyUI-vLLM-Omni/example_workflows/*.json *.jsonl *.parquet diff --git a/apps/ComfyUI-vLLM-Omni/.gitignore b/apps/ComfyUI-vLLM-Omni/.gitignore new file mode 100644 index 000000000..5704ad153 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/.gitignore @@ -0,0 +1,115 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# OSX useful to ignore +*.DS_Store +.AppleDouble +.LSOverride + +# Thumbnails +._* + +# Files that might appear in the root of a volume +.DocumentRevisions-V100 +.fseventsd +.Spotlight-V100 +.TemporaryItems +.Trashes +.VolumeIcon.icns +.com.apple.timemachine.donotpresent + +# Directories potentially created on remote AFP share +.AppleDB +.AppleDesktop +Network Trash Folder +Temporary Items +.apdisk + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +venv/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*,cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log + +# Sphinx documentation +docs/_build/ + +# IntelliJ Idea +.idea +*.iml +*.ipr +*.iws + +# PyBuilder +target/ + +# Cookiecutter +output/ +python_boilerplate/ +cookiecutter-pypackage-env/ + +# vscode settings +.history/ +*.code-workspace + +# Frontend extension +node_modules/ +.env +.env.local +.env.development.local +.env.test.local +.env.production.local +npm-debug.log* +yarn-debug.log* +yarn-error.log* +node.zip +.vscode/ +.claude/ +.codemate/ diff --git a/apps/ComfyUI-vLLM-Omni/LICENSE b/apps/ComfyUI-vLLM-Omni/LICENSE new file mode 100644 index 000000000..b3c346397 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/LICENSE @@ -0,0 +1,15 @@ +Apache Software License 2.0 + +Copyright (c) 2026, Zeyu Huang + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/apps/ComfyUI-vLLM-Omni/README.md b/apps/ComfyUI-vLLM-Omni/README.md new file mode 100644 index 000000000..dd530174f --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/README.md @@ -0,0 +1,170 @@ +# vLLM-Omni + +vLLM-Omni offers a ComfyUI integration on top of its online serving API. +It can send model inference requests to either a locally running vLLM-Omni service or a remote one. + +## Requirement + +- Python 3.12 or above +- [ComfyUI installed](https://docs.comfy.org/installation/system_requirements) +- [vLLM-Omni installed](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/) on either the same device or another device discoverable via the internet. +- No need to install additional packages apart from those already required by ComfyUI. + +> [!TIP] +> If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them. + + +## Installation + +Copy this folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`. + +If you are running ComfyUI during copying, you should restart ComfyUI to load this extension. + +> [!TIP] +> You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info. + +On the device and virtual environment you run ComfyUI, launch ComfyUI with +```bash +cd ComfyUI + +# The regular way +python main.py + +# If you are mainly using this node, launch it faster with +python main.py --cpu +``` + +On the device and virtual environment you run vLLM-Omni, start a model service with +```bash +vllm serve The_Model_ID_to_Serve --omni --port 8000 +``` + +Check **ComfyUI's sidebar -> Node Library**. There should be a new folder named **vLLM-Omni**. +If no, check your shell running the ComfyUI process. There may be some error messages before the line `Import times for custom nodes:` and the line `To see the GUI go to: http://127.0.0.1:8188`. + +## Quickstart + +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 +- **Multimodality Comprehension** 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**) + +> [!INFO] +> The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support. + +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. +- 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. + - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages. + +**The following features are tested**: + +- Single-node workflows for + - Multimodal Comprehension (e.g., Qwen Omni, BAGEL) + - Text-to-Image Generation (e.g., Qwen-Image) + - Image-to-Image Generation (e.g., Qwen-Image-Edit) + - TTS (e.g., Qwen TTS, including VoiceDesign, VoiceClone, CustomVoice) + +**The following features are not currently tested**. They may work or break. You are welcomed to test it out and offer comments. + +- Multi-node workflow that connects multiple model services together. + +## Screenshots and Examples + +### Multimodal comprehension (e.g., Qwen Omni series, BAGEL) + +(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Annotated Example**) + +

+ + + vLLM-Omni Main Architecture + +

+ +> [!TIP] +> Although this node enables all-modality input, you should check whether the specific model you host and request for supports the modalities you connect to the node. + +You can configure per-stage sampling parameters for multi-stage models. + +

+ + + vLLM-Omni Main Architecture + +

+ +### Text-to-image and image-to-image generation (e.g., Z-Image-Turbo, Qwen-Image-Edit, BAGEL) + +(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Image Generation**) + +

+ + + vLLM-Omni Main Architecture + +

+ +> [!TIP] +> The node automatically choose text-to-image or image-to-image 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**) + +

+ + + vLLM-Omni Main Architecture + +

+ +> [!TIP] +> There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node. + +## Develop + +Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/). + +## Limitation and Non-Goals + +- Single server mode only. No automatic load balancing or failover. +- Features set is bounded to vLLM-Omni's online service capability, including + - The types of models supported in online mode, + - The types of sampling parameters supported in the online mode, + - The ways to send files (primarily through full-length base64 in JSON payload), + - Figuring out errors in the payload (such as unsupported fields by a specific model) if the endpoint does not explicitly return an error, + - (The lack of) Authentication + - (The lack of) Progress indicator + +## Support + +If you are new to ComfyUI, please check out [its documentation](https://docs.comfy.org/) for usage instructions. + +If you are new to vLLM-Omni, please also check out [its documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/) for usage instructions. + +Whenever you find an issue or problem, please + +- First find out if this is an upstream limitation of vLLM-Omni's online serving mode, by [checking their documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/examples/). +- [Open an issue](https://github.com/vllm-project/vllm-omni/issues) that clearly describes this ComfyUI or online service problem. + +## Acknowledgements + +Features + +- https://github.com/dougbtv/comfyui-vllm-omni/ The official reference implementation for ComfyUI integration with vLLM-Omni's DALL-E compatible image generation API. +- https://github.com/Comfy-Org/ComfyUI/tree/master/comfy_extras ComfyUI's built-in node implementations. + +UI/UX design references + +- https://github.com/sgl-project/sglang/pull/15271 SGLang Diffusion's official ComfyUI integration for image and video generation. +- https://github.com/SXQBW/ComfyUI-Qwen-Omni A third party ComfyUI integration for Qwen Omni series. +- https://github.com/flybirdxx/ComfyUI-Qwen-TTS https://github.com/DarioFT/ComfyUI-Qwen3-TTS Third party ComfyUI integrations for Qwen TTS series. diff --git a/apps/ComfyUI-vLLM-Omni/__init__.py b/apps/ComfyUI-vLLM-Omni/__init__.py new file mode 100644 index 000000000..ffb86ddb3 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/__init__.py @@ -0,0 +1,52 @@ +"""Top-level package for comfyui_vllm_omni.""" # noqa: N999 # This is not a python library intended to be imported + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "WEB_DIRECTORY", +] + +__author__ = """Zeyu Huang""" +__email__ = "11222265+fhfuih@users.noreply.github.com" +__version__ = "0.0.1" + +from .comfyui_vllm_omni.nodes import ( + VLLMOmniARSampling, + VLLMOmniComprehension, + VLLMOmniDiffusionSampling, + VLLMOmniGenerateImage, + VLLMOmniQwenTTSParams, + VLLMOmniSamplingParamsList, + VLLMOmniTTS, + VLLMOmniVoiceClone, +) + +# A dictionary that contains all nodes you want to export with their names +NODE_CLASS_MAPPINGS = { + # === Generation === + "VLLMOmniGenerateImage": VLLMOmniGenerateImage, + "VLLMOmniComprehension": VLLMOmniComprehension, + "VLLMOmniTTS": VLLMOmniTTS, + "VLLMOmniVoiceClone": VLLMOmniVoiceClone, + # === Params === + "VLLMOmniARSampling": VLLMOmniARSampling, + "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling, + "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList, + "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams, +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + # === Generation === + "VLLMOmniGenerateImage": "Generate Image", + "VLLMOmniComprehension": "Multimodality Comprehension", + "VLLMOmniTTS": "TTS (Text to Speech)", + "VLLMOmniVoiceClone": "TTS Voice Cloning", + # === Params === + "VLLMOmniARSampling": "AR Sampling Params", + "VLLMOmniDiffusionSampling": "Diffusion Sampling Params", + "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List", + "VLLMOmniQwenTTSParams": "Qwen TTS Params", +} + +WEB_DIRECTORY = "./web" diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/__init__.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/__init__.py new file mode 100644 index 000000000..ebc9c5a59 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/__init__.py @@ -0,0 +1 @@ +# noqa: N999 # This is not a python library intended to be imported diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py new file mode 100644 index 000000000..ebc2822df --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py @@ -0,0 +1,569 @@ +from typing import Literal, cast + +import torch +from comfy_api.input import AudioInput, VideoInput + +from .utils.api_client import VLLMOmniClient +from .utils.logger import get_logger +from .utils.models import lookup_model_spec +from .utils.types import AudioFormat +from .utils.validators import ( + add_sampling_parameters_to_stage, + validate_model_and_sampling_params_types, +) + +logger = get_logger(__name__) + + +class _VLLMOmniGenerateBase: + """Base class for vLLM-Omni generation nodes with shared functionality.""" + + CATEGORY = "vLLM-Omni" + + @classmethod + def VALIDATE_INPUTS(cls, url, model) -> str | Literal[True]: + """ + Can only validate this model's own input. Cannot check inputs from other nodes. + See: https://docs.comfy.org/custom-nodes/backend/server_overview#validate_inputs + """ + if not url: + return "URL must not be empty" + if not model: + return "Model must not be empty" + return True + + +class VLLMOmniGenerateImage(_VLLMOmniGenerateBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ("STRING", {"default": "http://localhost:8000/v1"}), + "model": ("STRING", {"default": "Tongyi-MAI/Z-Image-Turbo"}), + "prompt": ("STRING", {"multiline": True}), + "negative_prompt": ("STRING", {"multiline": True, "default": ""}), + "width": ("INT", {"default": 512, "min": 64, "max": 2048}), + "height": ("INT", {"default": 512, "min": 64, "max": 2048}), + }, + "optional": { + "image": ("IMAGE",), + "mask": ("MASK",), + # "video": ("VIDEO",), + # "audio": ("AUDIO",), + "sampling_params": ("SAMPLING_PARAMS",), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "generate" + + async def generate( + self, + url: str, + model: str, + prompt: str, + width: int, + height: int, + negative_prompt: str | None = None, + image: torch.Tensor | None = None, + mask: torch.Tensor | None = None, + audio: AudioInput | None = None, # Hidden & unused + video: VideoInput | None = None, # Hidden & unused + sampling_params: dict | list[dict] | None = None, + **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: + raise ValueError("Mask input provided without an image input.") + + client = VLLMOmniClient(url) + + spec, pattern = lookup_model_spec(model) + is_bagel = pattern is not None and "bagel" in pattern.lower() + + # Prefer DALL-E compatible API for simple (one-stage) diffusion models + if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel: + sampling_params = cast(dict | None, sampling_params) + if audio is None and image is None and video is None: + # No multimodal input --- use DALL-E image generation + logger.info("Using DALL-E image generation endpoint") + output = await client.generate_image( + model=model, + prompt=prompt, + width=width, + height=height, + negative_prompt=negative_prompt, + sampling_params=sampling_params, + ) + return (output,) + elif image is not None and audio is None and video is None: + # Image and text input --- use DALL-E image edit + logger.info("Using DALL-E image edit endpoint") + output = await client.edit_image( + model=model, + prompt=prompt, + image=image, + width=width, + height=height, + negative_prompt=negative_prompt, + mask=mask, + sampling_params=sampling_params, + ) + return (output,) + + logger.info("Using chat completion endpoint") + sampling_params = add_sampling_parameters_to_stage( + model, sampling_params, "diffusion", width=width, height=height + ) + logger.debug("Edited sampling params: %s", sampling_params) + + output = await client.generate_image_chat_completion( + model=model, + prompt=prompt, + negative_prompt=negative_prompt, + image=image, + audio=audio, + video=video, + sampling_params=sampling_params, + ) + + return (output,) + + +class VLLMOmniComprehension(_VLLMOmniGenerateBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ("STRING", {"default": "http://localhost:8000/v1"}), + "model": ("STRING", {"default": "Qwen/Qwen2.5-Omni-7B"}), + "prompt": ("STRING", {"multiline": True}), + "output_text": ("BOOLEAN", {"default": True}), + "output_audio": ("BOOLEAN", {"default": True}), + "use_audio_in_video": ("BOOLEAN", {"default": True}), + }, + "optional": { + "image": ("IMAGE",), + "video": ("VIDEO",), + "audio": ("AUDIO",), + "sampling_params": ("SAMPLING_PARAMS",), + }, + } + + RETURN_TYPES = ("STRING", "AUDIO") + RETURN_NAMES = ("text_response", "audio_response") + FUNCTION = "generate" + + @classmethod + def VALIDATE_INPUTS(cls, url, model, output_text, output_audio) -> str | Literal[True]: # type: ignore[reportIncompatibleMethodOverride] + super_validation = super().VALIDATE_INPUTS(url, model) + if isinstance(super_validation, str): + return super_validation + if not output_text and not output_audio: + return "At least one of output_text or output_audio must be True." + return True + + async def generate( + self, + url: str, + model: str, + prompt: str, + image: torch.Tensor | None = None, + audio: AudioInput | None = None, + video: VideoInput | None = None, + sampling_params: dict | list[dict] | None = None, + output_text: bool = True, + output_audio: bool = True, + use_audio_in_video: bool = True, + **kwargs, + ) -> tuple[str, AudioInput]: + logger.info("Uncaught kwargs: %s", kwargs) + logger.debug("Got sampling params: %s", sampling_params) + validate_model_and_sampling_params_types(model, sampling_params) + + client = VLLMOmniClient(url) + spec, pattern = lookup_model_spec(model) + is_bagel = pattern is not None and "bagel" in pattern.lower() + + if is_bagel: + # A lot of special handlings here... + if output_audio: + raise ValueError("BAGEL models do not support audio output.") + if audio is not None or video is not None: + raise ValueError("BAGEL models do not support audio or video input.") + ( + text_response, + _, + ) = await client.generate_comprehension_chat_completion( + model=model, + prompt=prompt, + image=image, + audio=None, + video=None, + sampling_params=sampling_params, + modalities=["text"], + ) + else: + modalities = [] + if output_text: + modalities.append("text") + if output_audio: + modalities.append("audio") + + if use_audio_in_video and video is not None: + use_audio_in_video = True + else: + use_audio_in_video = False + + ( + text_response, + audio, + ) = await client.generate_comprehension_chat_completion( + model=model, + prompt=prompt, + image=image, + audio=audio, + video=video, + sampling_params=sampling_params, + modalities=modalities, + # == extra kwargs == + mm_processor_kwargs={"use_audio_in_video": use_audio_in_video}, + ) + + if text_response is None: + text_response = "" + if audio is None: + channels = 1 + duration = 1 + sample_rate = 44100 + num_samples = int(round(duration * sample_rate)) + waveform = torch.zeros((1, channels, num_samples), dtype=torch.float32) + audio = {"waveform": waveform, "sample_rate": sample_rate} + + return (text_response, audio) + + +class VLLMOmniTTS(_VLLMOmniGenerateBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ("STRING", {"default": "http://localhost:8000/v1"}), + "model": ( + "STRING", + {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"}, + ), + "input": ("STRING", {"multiline": True}), + "voice": ("STRING", {"default": "Vivian"}), + "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],), + "speed": ( + "FLOAT", + {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01}, + ), + }, + "optional": { + "model_specific_params": ("TTS_PARAMS",), + }, + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "generate" + + async def generate( + self, + url: str, + model: str, + input: str, + voice: str, + response_format: AudioFormat, + speed: float, + model_specific_params: dict | None, + **kwargs, + ) -> tuple[AudioInput]: + logger.info("Got extra kwargs in TTS: %s", kwargs) + + is_qwen_tts = "qwen3-tts" in model.lower() + extra_params_type = None if model_specific_params is None else model_specific_params["type"] + if not is_qwen_tts and extra_params_type == "qwen-tts": + raise ValueError( + "You have provided Qwen-specific TTS params." + "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)." + ) + + combined_params = {**kwargs, **(model_specific_params or {})} + combined_params.pop("type", None) # Internal fields in model_specific_params + + client = VLLMOmniClient(url) + + audio = await client.generate_speech( + model=model, + input=input, + voice=voice, + response_format=response_format, + speed=speed, + **combined_params, + ) + return (audio,) + + +class VLLMOmniVoiceClone(_VLLMOmniGenerateBase): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "url": ("STRING", {"default": "http://localhost:8000/v1"}), + "model": ("STRING", {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-Base"}), + "input": ("STRING", {"multiline": True}), + "voice": ("STRING", {"default": "Vivian"}), + "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],), + "speed": ( + "FLOAT", + {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01}, + ), + "ref_audio": ("AUDIO",), + "ref_text": ("STRING", {"multiline": True}), + "x_vector_only_mode": ("BOOLEAN", {"default": False}), + }, + "optional": { + "model_specific_params": ("TTS_PARAMS",), + }, + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "generate" + + async def generate( + self, + url: str, + model: str, + input: str, + voice: str, + response_format: AudioFormat, + speed: float, + ref_audio: AudioInput, + ref_text: str, + x_vector_only_mode: bool, + model_specific_params: dict | None, + **kwargs, + ): + is_qwen_tts = "qwen3-tts" in model.lower() + extra_params_type = None if model_specific_params is None else model_specific_params["type"] + if not is_qwen_tts and extra_params_type == "qwen-tts": + raise ValueError( + "You have provided Qwen-specific TTS params." + "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)." + ) + + combined_params = { + "ref_audio": ref_audio, + "ref_text": ref_text, + "x_vector_only_mode": x_vector_only_mode, + **kwargs, + **(model_specific_params or {}), + } + combined_params.pop("type", None) # Internal fields in model_specific_params + + client = VLLMOmniClient(url) + + audio = await client.generate_speech( + model=model, + input=input, + voice=voice, + response_format=response_format, + speed=speed, + **combined_params, + ) + return (audio,) + + +class VLLMOmniARSampling: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "max_tokens": ("INT", {"default": 100, "min": 1, "max": 10000}), + "temperature": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}, + ), + "top_p": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}, + ), + "repetition_penalty": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01}, + ), + # === Put seed at last. === + # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate" + "seed": ( + "INT", + { + "default": -1, + "min": -1, + "step": 1, + "tooltip": "-1 means to not provide a seed.", + }, + ), + } + } + + RETURN_TYPES = ("SAMPLING_PARAMS",) + RETURN_NAMES = ("AR sampling params",) + FUNCTION = "get_params" + CATEGORY = "vLLM-Omni/Sampling Params" + + def get_params(self, seed, **kwargs): + params = { + "type": "autoregression", # for internal use, removed before sending the request + **kwargs, + } + if seed >= 0: + params["seed"] = seed + return (params,) + + +class VLLMOmniDiffusionSampling: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "n": ( + "INT", + { + "default": 1, + "min": 0, + "max": 10, + "step": 1, + "tooltip": "Number of images to generate", + }, + ), + "num_inference_steps": ( + "INT", + { + "default": 50, + "min": 1, + "max": 1000, + "tooltip": "Number of denoising steps (higher = better quality, slower).", + }, + ), + "guidance_scale": ( + "FLOAT", + { + "default": 7.5, + "min": 0.0, + "max": 20.0, + "step": 0.1, + "tooltip": "Classifier-free guidance scale (higher = more prompt adherence).", + }, + ), + "true_cfg_scale": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 20.0, + "step": 0.5, + "tooltip": "True CFG scale for advanced control (model-specific).", + }, + ), + "vae_use_slicing": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)", + }, + ), + # === Put seed at last. === + # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate" + "seed": ( + "INT", + { + "default": -1, + "min": -1, + "step": 1, + "tooltip": "-1 means to not provide a seed.", + }, + ), + } + } + + RETURN_TYPES = ("SAMPLING_PARAMS",) + RETURN_NAMES = ("diffusion sampling params",) + FUNCTION = "get_params" + CATEGORY = "vLLM-Omni/Sampling Params" + + def get_params(self, seed, **kwargs): + params = { + "type": "diffusion", # for internal use, removed before sending the request + **kwargs, + } + if seed >= 0: + params["seed"] = seed + return (params,) + + +class VLLMOmniSamplingParamsList: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "param1": ("SAMPLING_PARAMS",), + }, + "optional": { + "param2": ("SAMPLING_PARAMS",), + "param3": ("SAMPLING_PARAMS",), + }, + } + + RETURN_TYPES = ("SAMPLING_PARAMS",) + RETURN_NAMES = ("param list",) + FUNCTION = "aggregate" + CATEGORY = "vLLM-Omni/Sampling Params" + + def aggregate(self, param1: dict, param2: dict | None = None, param3: dict | None = None): + for i, p in enumerate((param1, param2, param3)): + if isinstance(p, list): + raise ValueError( + f"Input {i} is a Multi-Stage Sampling Params List. " + f"Expected a single sampling parameters node (either AR or Diffusion)." + ) + + params = [param1] + if param2 is not None: + params.append(param2) + if param3 is not None: + params.append(param3) + return (params,) + + +class VLLMOmniQwenTTSParams: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "task_type": ( + ["CustomVoice", "VoiceDesign", "Base"], + {"default": "CustomVoice"}, + ), + "language": ( + ["Auto", "Chinese", "English", "Japanese", "Korean"], + {"default": "Auto"}, + ), + "instructions": ("STRING", {"multiline": True}), + "max_new_tokens": ("INT", {"default": 2048, "min": 1}), + } + } + + RETURN_TYPES = ("TTS_PARAMS",) + RETURN_NAMES = ("Qwen TTS params",) + FUNCTION = "get_params" + CATEGORY = "vLLM-Omni/TTS Params" + + def get_params(self, **kwargs): + return ({"type": "qwen-tts", **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 new file mode 100644 index 000000000..c5372641f --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py @@ -0,0 +1,473 @@ +""" +An high-level API client adapter that forwards ComfyUI inputs to vLLM-Omni's REST API, +and transforms the API responses back to ComfyUI formats. + +The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv). +Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License. +""" + +from typing import Any + +import aiohttp +import av.error +import torch +from comfy_api.input import AudioInput, VideoInput + +from .format import ( + audio_to_base64, + base64_to_audio, + base64_to_image_tensor, + bytes_to_audio, + image_tensor_to_base64, + image_tensor_to_png_bytes, + video_to_base64, +) +from .logger import get_logger, pretty_printer +from .models import lookup_model_spec +from .types import AudioFormat + +logger = get_logger(__name__) + + +class VLLMOmniClient: + def __init__(self, base_url: str, timeout: float = 300.0): + self.base_url = base_url + self.timeout = aiohttp.ClientTimeout(total=timeout) + + async def generate_image( + self, + *, + model: str, + prompt: str, + width: int, + height: int, + negative_prompt: str | None = None, + sampling_params: 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 = { + "model": model, + "prompt": prompt, + "size": size, + "response_format": "b64_json", + } + if negative_prompt: + payload["negative_prompt"] = negative_prompt + if sampling_params is not None: + # Only select specific sampling params + for k in ( + "n", + "num_inference_steps", + "guidance_scale", + "true_cfg_scale", + "vae_use_slicing", + ): + if k in sampling_params and sampling_params[k] is not None: + payload[k] = sampling_params[k] + logger.debug("img gen payload: %s", payload) + + url = self.base_url + "/images/generations" + async with aiohttp.ClientSession(timeout=self.timeout) as session: + try: + async with session.post( + url, + json=payload, + headers={"Content-Type": "application/json"}, + ) 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}") + 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") + + image_tensors = [] + for idx, img in enumerate(data["data"]): + if "b64_json" not in img: + raise RuntimeError(f"API returned image #{idx} without 'b64_json' field") + base64_str = img["b64_json"] + tensor = base64_to_image_tensor(base64_str) + image_tensors.append(tensor) + logger.debug("Image #%d has shape %s", idx, tensor.shape) + + batch_tensor = torch.stack(image_tensors, dim=0) + logger.debug("batch_tensor output has shape: %s", batch_tensor.shape) + return batch_tensor + + except aiohttp.ClientError as e: + raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}") + + async def edit_image( + self, + *, + model: str, + prompt: str, + image: torch.Tensor, + width: int, + height: int, + negative_prompt: str | None = None, + mask: torch.Tensor | None = None, + sampling_params: dict | None = None, + ) -> torch.Tensor: + """Run image editing via DALLE API""" + await self._check_model_exist(model) + + size = f"{width}x{height}" + image_filename = "image.png" # Required for multipart form + form = aiohttp.FormData() + form.add_field("model", model) + form.add_field( + "image", + image_tensor_to_png_bytes(image, image_filename), + filename=image_filename, + content_type="image/png", + ) + form.add_field("prompt", prompt) + form.add_field("size", size) + if negative_prompt: + form.add_field("negative_prompt", negative_prompt) + if sampling_params is not None: + # Only select specific sampling params + for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"): + if k in sampling_params and sampling_params[k] is not None: + form.add_field(k, str(sampling_params[k])) + if mask is not None: + mask_filename = "mask.png" + form.add_field( + "mask", + image_tensor_to_png_bytes(mask, mask_filename), + filename=mask_filename, + content_type="image/png", + ) + + url = self.base_url + "/images/edits" + 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}") + + 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") + + image_tensors = [] + for idx, img in enumerate(data["data"]): + if "b64_json" not in img: + raise RuntimeError(f"API returned image #{idx} without 'b64_json' field") + base64_str = img["b64_json"] + tensor = base64_to_image_tensor(base64_str) + image_tensors.append(tensor) + + return torch.stack(image_tensors, dim=0) + + except aiohttp.ClientError as e: + raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}") + + async def generate_image_chat_completion( + self, + *, + model: str, + prompt: str, + negative_prompt: str | None = None, + image: torch.Tensor | None = None, + audio: AudioInput | None = None, + video: VideoInput | None = None, + sampling_params: dict | list[dict] | None = None, + ) -> torch.Tensor: + payload = VLLMOmniClient._prepare_chat_completion_messages( + model=model, + prompt=prompt, + negative_prompt=negative_prompt, + image=image, + audio=audio, + video=video, + sampling_params=sampling_params, + modalities=["image"], + ) + choices = await self._generate_base_chat_completion(model, payload) + + image_tensors = [] + for idx, img_content in enumerate(choices[0]["message"]["content"]): + base64_str = img_content.get("image_url", {}).get("url", "") + if not base64_str: + raise RuntimeError(f"API returned image #{idx} without image url") + tensor = base64_to_image_tensor(base64_str) + image_tensors.append(tensor) + + return torch.stack(image_tensors, dim=0) + + async def generate_comprehension_chat_completion( + self, + *, + model: str, + prompt: str, + image: torch.Tensor | None = None, + audio: AudioInput | None = None, + video: VideoInput | None = None, + sampling_params: dict | list[dict] | None = None, + modalities: list[str] = ["text", "audio"], + **extra_body, + ) -> tuple[str | None, AudioInput | None]: + # Response may contain two choices: one with text, one with audio + payload = VLLMOmniClient._prepare_chat_completion_messages( + model=model, + prompt=prompt, + negative_prompt=None, + image=image, + audio=audio, + video=video, + sampling_params=sampling_params, + modalities=modalities, + **extra_body, + ) + + choices = await self._generate_base_chat_completion(model, payload) + text_response = None + audio_base64 = None + for choice in choices: + try: + text_response = choice["message"]["content"] + except (KeyError, TypeError): + # Either this case (text response) or the audio response case will be hit. Checking None's later. + pass + try: + audio_base64 = choice["message"]["audio"]["data"] + except (KeyError, TypeError): + # Either this case (text response) or the audio response case will be hit. Checking None's later. + pass + if audio_base64 is None and text_response is None: + raise RuntimeError( + "API response missing both '.message.audio' and 'message.content' fields." + f"The choices object is {choices}" + ) + if audio_base64 is not None: + audio = base64_to_audio(audio_base64) + logger.debug( + "audio sample rate %d, audio shape %s, duration in second %f", + audio["sample_rate"], + audio["waveform"].shape, + audio["waveform"].shape[2] / audio["sample_rate"], + ) + else: + audio = None + return text_response, audio + + async def generate_speech( + self, + *, + model: str, + input: str, + voice: str, + response_format: AudioFormat, + speed: float, + **extra_params, + ) -> AudioInput: + await self._check_model_exist(model) + + ref_audio: AudioInput | None = extra_params.pop("ref_audio", None) + + payload = { + "model": model, + "input": input, + "voice": voice, + "response_format": response_format, + "speed": speed, + **extra_params, + } + + if ref_audio is not None: + audio_base64 = audio_to_base64(ref_audio) + payload["ref_audio"] = audio_base64 + + logger.debug("Omni TTS payload: %s", pretty_printer.pformat(payload)) + + url = self.base_url + "/audio/speech" + async with aiohttp.ClientSession(timeout=self.timeout) as session: + try: + async with session.post( + url, + json=payload, + headers={"Content-Type": "application/json"}, + ) 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: + audio_bytes = await response.read() + except aiohttp.ContentTypeError as e: + raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}") + + try: + audio = bytes_to_audio(audio_bytes) + except av.error.InvalidDataError as e: + raise ValueError( + f"Invalid audio data received from vLLM-Omni: {e}" + "Check if you have input unsupported arguments (such as 'voice')" + ) + return audio + + except aiohttp.ClientError as e: + raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}") + + async def _generate_base_chat_completion(self, model: str, payload: dict[str, Any]) -> list[dict[str, Any]]: + logger.debug("Omni payload: %s", pretty_printer.pformat(payload)) + await self._check_model_exist(model) + + url = self.base_url + "/chat/completions" + async with aiohttp.ClientSession(timeout=self.timeout) as session: + try: + async with session.post( + url, + json=payload, + headers={"Content-Type": "application/json"}, + ) 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}") + + logger.debug( + "chat completion response: %s", + pretty_printer.pformat(data), + ) + + try: + return data["choices"] + except (KeyError, TypeError): + raise RuntimeError("Invalid JSON response from vLLM-Omni: missing 'choices' field") + + except aiohttp.ClientError as e: + raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}") + + async def _check_model_exist(self, model: str): + url = self.base_url + "/models" + async with aiohttp.ClientSession(timeout=self.timeout) as session: + try: + async with session.get( + url, + headers={"Content-Type": "application/json"}, + ) 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} " + f"when getting hosted model list: {error_text}" + ) + + try: + data = await response.json() + except aiohttp.ContentTypeError as e: + raise RuntimeError(f"Invalid JSON response when getting hosted model list from vLLM-Omni: {e}") + + except aiohttp.ClientError as e: + raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}") + try: + model_list = data["data"] + model_found = next((True for m in model_list if m["id"] == model), False) + except (KeyError, TypeError): + raise RuntimeError(f"Invalid JSON response of the hosted model list: {data}") + + if not model_found: + raise ValueError(f"Model {model} not served at {self.base_url}.") + + @staticmethod + def _prepare_chat_completion_messages( + *, + model: str, + prompt: str, + negative_prompt: str | None, + image: torch.Tensor | None = None, + audio: AudioInput | None = None, + video: VideoInput | None = None, + sampling_params: dict | list[dict] | None = None, + modalities: list[str] | None = None, # diffusion don't have this field + **extra_body, + ): + message_content: list[dict] = [{"type": "text", "text": prompt}] + if image is not None: + message_content.append( + { + "type": "image_url", + "image_url": {"url": image_tensor_to_base64(image)}, + } + ) + if audio is not None: + message_content.append({"type": "audio_url", "audio_url": {"url": audio_to_base64(audio)}}) + if video is not None: + message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}}) + messages = [{"role": "user", "content": message_content}] + + combined_extra_body: dict[str, Any] = {} + if sampling_params is not None: + spec, _ = lookup_model_spec(model) + is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1 + + # Exclude internal key + if isinstance(sampling_params, dict): + sampling_params = {k: v for k, v in sampling_params.items() if k != "type"} + else: + sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params] + + if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]): + # Diffusion format: extra_body directly contains sampling params. + # Validation should have taken care of matching sampling params' types. + # * Use this mode if the model is a simple one-stage diffusion model. + # * Fallback to this mode if model is not registered and a single sampling param is provided. + sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0] + combined_extra_body: dict[str, Any] = {**sampling_params} + if "n" in combined_extra_body: + combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"] + del combined_extra_body["n"] + else: + # Use AR style payload, extra_body has a sampling_params_list field + combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params} + + if negative_prompt: + combined_extra_body["negative_prompt"] = negative_prompt + + if extra_body: + combined_extra_body.update(extra_body) + + payload: dict[str, Any] = {"messages": messages, "model": model} + if combined_extra_body: + payload["extra_body"] = combined_extra_body + if modalities: + payload["modalities"] = modalities + + spec, _ = lookup_model_spec(model) + if spec: + preprocessor = spec.get("payload_preprocessor", None) + if preprocessor is not None: + payload = preprocessor(payload) + + return payload diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py new file mode 100644 index 000000000..0ddf86f9c --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py @@ -0,0 +1,230 @@ +"""Image/tensor format helpers. + +The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv). +Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License. +""" + +import base64 +import mimetypes +from io import BytesIO + +import av +import numpy as np +import torch +from comfy_api.input import AudioInput, VideoInput +from comfy_extras import nodes_audio +from PIL import Image + +from .logger import get_logger + +logger = get_logger(__name__) + + +def base64_to_image_tensor(base64_str: str, mode: str = "RGB") -> torch.Tensor: + """ + Convert base64-encoded image to ComfyUI image tensor. + + Args: + base64_str: Base64-encoded image string + mode: PIL image mode (default RGB for transparency support) + + Returns: + torch.Tensor with shape (1, H, W, C) in float32 [0, 1] range + + Raises: + ValueError: If base64 string is invalid or image cannot be decoded + """ + if base64_str.startswith("data:image"): + _, base64_str = base64_str.split(",", 1) + + try: + # Decode base64 to bytes + image_bytes = base64.b64decode(base64_str) + except Exception as e: + raise ValueError(f"Invalid base64 string: {e}") + + # Create BytesIO object for PIL + image_bytesio = BytesIO(image_bytes) + + # Open with PIL and convert to desired mode + try: + pil_image = Image.open(image_bytesio) + pil_image = pil_image.convert(mode) + except Exception as e: + raise RuntimeError(f"Failed to open image: {e}") + + image_array = np.asarray(pil_image).astype(np.float32) / 255.0 + image_tensor = torch.from_numpy(image_array) + return image_tensor + + +def image_tensor_to_png_bytes(tensor: torch.Tensor, filename: str = "image.png") -> BytesIO: + """ + Convert ComfyUI image tensor to PNG BytesIO for multipart upload. + + This function converts a ComfyUI IMAGE tensor to a PNG-encoded BytesIO object + suitable for multipart/form-data upload. The BytesIO object has its .name + attribute set, which is required by aiohttp for file uploads. + + Args: + tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1] + filename: Name attribute to set on BytesIO (default: "image.png") + + Returns: + BytesIO object containing PNG-encoded image with .name attribute set + + Raises: + ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.) + """ + if tensor.ndim != 4: + raise ValueError(f"Expected 4D tensor with shape (B, H, W, C), got {tensor.ndim}D tensor") + + image_tensor = tensor[0] # Shape: (H, W, C) + image_np = (image_tensor.cpu().numpy() * 255.0).astype(np.uint8) + pil_image = Image.fromarray(image_np) + + # Save to BytesIO as image file + img_bytes = BytesIO() + # Set name attribute (required for multipart upload and mimetype detection) + img_bytes.name = filename + try: + pil_image.save(img_bytes) + except Exception as e: + raise RuntimeError(f"Failed to save image as file: {e}") + + # Reset position to beginning + img_bytes.seek(0) + + return img_bytes + + +def image_tensor_to_base64(tensor: torch.Tensor, filename: str = "image.png") -> str: + """ + Convert ComfyUI image tensor to base64-encoded image string. + + Args: + tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1] + filename: Name attribute to set on BytesIO (default: "image.png") + format: File format of the output image file buffer (default: "PNG") + + Returns: + Base64-encoded image string + + Raises: + ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.) + """ + img_bytes = image_tensor_to_png_bytes(tensor, filename) + img_bytes.seek(0) + byte_data = img_bytes.read() + base64_str = base64.b64encode(byte_data).decode("utf-8") + mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" + return f"data:{mime_type};base64,{base64_str}" + + +def video_to_bytes(video: VideoInput, filename: str = "video.mp4") -> BytesIO: + output_buffer = BytesIO() + output_buffer.name = filename + video.save_to(output_buffer) + output_buffer.seek(0) + return output_buffer + + +def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str: + video_buffer = video_to_bytes(video, filename) + video_buffer.seek(0) + byte_data = video_buffer.read() + base64_str = base64.b64encode(byte_data).decode("utf-8") + mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" + return f"data:{mime_type};base64,{base64_str}" + + +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"] + format = filename.rsplit(".", maxsplit=1)[1] + layout = "mono" if waveform.shape[0] == 1 else "stereo" + + output_buffer = BytesIO() + output_buffer.name = filename + output_container = av.open(output_buffer, mode="w", format=format) + if format == "opus": + out_stream = output_container.add_stream("libopus", rate=sample_rate, layout=layout) + if quality == "64k": + out_stream.bit_rate = 64000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + elif quality == "96k": + out_stream.bit_rate = 96000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + elif quality == "128k": + out_stream.bit_rate = 128000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + elif quality == "192k": + out_stream.bit_rate = 192000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + elif quality == "320k": + out_stream.bit_rate = 320000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + elif format == "mp3": + out_stream = output_container.add_stream("libmp3lame", rate=sample_rate, layout=layout) + if quality == "V0": + out_stream.codec_context.qscale = 1 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + elif quality == "128k": + out_stream.bit_rate = 128000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + elif quality == "320k": + out_stream.bit_rate = 320000 # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + else: # format == "flac": + out_stream = output_container.add_stream("flac", rate=sample_rate, layout=layout) + + frame = av.AudioFrame.from_ndarray( + waveform.movedim(0, 1).reshape(1, -1).float().numpy(), + format="flt", + layout=layout, + ) + frame.sample_rate = sample_rate + frame.pts = 0 + output_container.mux(out_stream.encode(frame)) # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + # Flush encoder + output_container.mux(out_stream.encode(None)) # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py + output_container.close() + output_buffer.seek(0) + + return output_buffer + + +def audio_to_base64(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> str: + audio_buffer = audio_to_bytes(audio, filename, quality) + audio_buffer.seek(0) + byte_data = audio_buffer.read() + base64_str = base64.b64encode(byte_data).decode("utf-8") + mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" + return f"data:{mime_type};base64,{base64_str}" + + +def bytes_to_audio(audio_bytes: bytes) -> AudioInput: + """ + Convert audio bytes to ComfyUI audio tensor. + + Args: + audio_bytes: Audio file bytes + Returns: + torch.Tensor with shape (B, C, T) in float32 range [-1, 1] + """ + audio_buffer = BytesIO(audio_bytes) + waveform, sample_rate = nodes_audio.load(audio_buffer) # type: ignore # Although expect string argument, it calls av.open underneath, which supports BytesIO (file-like) + return {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate} + + +def base64_to_audio(base64_str: str) -> AudioInput: + """ + Convert base64-encoded audio to ComfyUI audio tensor. + + Args: + base64_str: Base64-encoded audio string + Returns: + torch.Tensor with shape (B, C, T) in float32 range [-1, 1] + """ + if base64_str.startswith("data:audio"): + _, base64_str = base64_str.split(",", 1) + + try: + # Decode base64 to bytes + audio_bytes = base64.b64decode(base64_str) + except Exception as e: + raise ValueError(f"Invalid base64 string: {e}") + + return bytes_to_audio(audio_bytes) diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/logger.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/logger.py new file mode 100644 index 000000000..cad640f4d --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/logger.py @@ -0,0 +1,125 @@ +"""Centralized logger configuration for vLLM-Omni ComfyUI.""" + +import logging +import pprint +import sys +from typing import Any + + +def get_logger(name: str) -> logging.Logger: + """ + Get or create a logger with proper formatting. + + Args: + name: Logger name (typically __name__ of the calling module) + + Returns: + Configured logger instance + """ + logger = logging.getLogger(name) + + # Only configure if not already configured + if not logger.handlers: + logger.setLevel(logging.DEBUG) + + # Create console handler + handler = logging.StreamHandler(sys.stdout) + handler.setLevel(logging.INFO) + + # Create formatter + formatter = logging.Formatter( + fmt="(ComfyUI-vLLM-Omni) [%(levelname)s] %(asctime)s [%(filename)s:%(lineno)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + handler.setFormatter(formatter) + + # Add handler to logger + logger.addHandler(handler) + + # Prevent propagation to root logger + logger.propagate = False + + return logger + + +class OmitBase64PrettyPrinter(pprint.PrettyPrinter): + """ + A PrettyPrinter that redacts specific field names with '...' + wherever they appear in nested structures. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def _format(self, obj: Any, stream, indent: int, allowance: int, context, level: int) -> None: + # Check if this is a dict with redacted keys + if isinstance(obj, dict): + # Create a copy with redacted values + display_obj = {} + for key, value in obj.items(): + if key == "data" or key == "url": + if value.startswith("data:"): + base64_header = value.split(",", 1)[0] + display_obj[key] = f"{base64_header},***" + elif value.startswith("http://") or value.startswith("https://"): + display_obj[key] = value + elif len(value) > 10: + display_obj[key] = f"{value[:10]}***" + else: + display_obj[key] = value + else: + display_obj[key] = value + obj = display_obj + + # Handle list/tuple/set containing dicts that might have redacted keys + # (pprint will recursively call _format on nested items, so this + # handles arbitrary nesting automatically) + + super()._format(obj, stream, indent, allowance, context, level) + + +pretty_printer = OmitBase64PrettyPrinter() + + +# ========== PPrint EXAMPLES ========== + +if __name__ == "__main__": + data = { + "messages": [ + { + "role": "system", + "content": ( + "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group," + "capable of perceiving auditory and visual inputs, as well as generating text and speech." + ), + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What sound is it, and what is the drawing about?", + }, + { + "type": "text", + "text": "What sound is it, and what is the drawing about?", + }, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAwAAAA"}, + }, + { + "type": "audio_url", + "audio_url": {"url": "data:audio/mpeg;base64,SUQzBAAAAAAAIlRTU0UAAAAOAAADT="}, + }, + ], + }, + ], + "extra_body": {"mm_processor_kwargs": {"use_audio_in_video": False}}, + "modalities": ["text"], + } + + # Create printer that redacts 'password' and 'token' fields + + print("\nRedactingPrettyPrinter (hides secrets):") + pretty_printer.pprint(data) diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py new file mode 100644 index 000000000..83bf64fe7 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/models.py @@ -0,0 +1,119 @@ +import re + +from .types import Modality, ModelMode, Spec + + +def _bagel_payload_preprocessor(payload: dict) -> dict: + try: + for message in payload["messages"]: + for content in message["content"]: + if content["type"] == "text": + content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>" + except (KeyError, TypeError): + raise RuntimeError("Internal Error: malformatted BAGEL payload") + extra_body = payload.pop("extra_body", {}) + return {**payload, **extra_body} + + +def _qwen25_payload_preprocessor(payload: dict) -> dict: + if payload["messages"][0]["role"] != "system": + payload["messages"] = [ + { + "role": "system", + "content": ( + "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group," + "capable of perceiving auditory and visual inputs, as well as generating text and speech." + ), + }, + *payload["messages"], + ] + return payload + + +_MODEL_PIPELINE_SPECS: dict[str, Spec] = { + r"BAGEL-7B-MoT": { + "stages": [ + "diffusion" # The vLLM-Omni interface treats it as a single-stage diffusion model + ], + "modes": [ + { + "mode": ModelMode.COMPREHENSION, + "input_modalities": [Modality.TEXT, Modality.IMAGE], + } + ], + "payload_preprocessor": _bagel_payload_preprocessor, + }, + r"Qwen2.5-Omni*": { + "stages": ["autoregression", "autoregression", "autoregression"], + "payload_preprocessor": _qwen25_payload_preprocessor, + "modes": [ + { + "mode": ModelMode.COMPREHENSION, + "input_modalities": [ + Modality.TEXT, + Modality.IMAGE, + Modality.VIDEO, + Modality.AUDIO, + ], + } + ], + }, + r"Qwen3-Omni*": { + "stages": ["autoregression", "autoregression", "autoregression"], + "modes": [ + { + "mode": ModelMode.COMPREHENSION, + "input_modalities": [ + Modality.TEXT, + Modality.IMAGE, + Modality.VIDEO, + Modality.AUDIO, + ], + } + ], + }, +} +# Convert dict keys to regex patterns +MODEL_PIPELINE_SPECS: dict[re.Pattern, Spec] = {} +for k, v in _MODEL_PIPELINE_SPECS.items(): + MODEL_PIPELINE_SPECS[re.compile(k)] = v +del _MODEL_PIPELINE_SPECS + + +def lookup_model_spec(model: str) -> tuple[Spec | None, str | None]: + try: + last_component = model.rstrip("/").rsplit("/", 1)[-1] + except IndexError: + last_component = model + for pattern, spec in MODEL_PIPELINE_SPECS.items(): + if pattern.search(last_component): + return spec, pattern.pattern + return None, None + + +# ============== DEMONSTRATION ============== + +if __name__ == "__main__": + test_paths = [ + "Qwen/Qwen2.5-Omni-7B", + "MyModels/Qwen2.5-Omni-3B", + "/root/home/Qwen2.5-Omni-7B", + "Qwen/Qwen3-Omni", + "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "Custom/Path/UnknownModel-Instruct", + "Not/Matching/Anything", + ] + + test_payload = {"messages": [{"role": "user", "content": "prompt"}]} + + print("Testing registry lookups:\n") + for path in test_paths: + spec, _ = lookup_model_spec(path) + if spec: + if preprocessor := spec.get("payload_preprocessor"): + result = preprocessor(test_payload) + print(f"✓ {path:<40} → {result}") + else: + print(f"✓ {path:<40} → No preprocessor") + else: + print(f"✗ {path:<40} → No match") diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py new file mode 100644 index 000000000..29be9648c --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/types.py @@ -0,0 +1,39 @@ +from collections.abc import Callable +from enum import Enum, auto +from typing import ( + Any, + Literal, + NotRequired, + TypeAlias, + TypedDict, +) + +AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"] + + +class ModelMode(Enum): + IMAGE_GENERATION = auto() + VIDEO_GENERATION = auto() + AUDIO_GENERATION = auto() + COMPREHENSION = auto() + + +class Modality(Enum): + TEXT = auto() # maybe not useful. Prompt is always required + IMAGE = auto() + VIDEO = auto() + AUDIO = auto() + + +class ModelModeSpec(TypedDict): + mode: ModelMode + input_modalities: list[Modality] + + +PayloadPreprocessor: TypeAlias = Callable[[dict[str, Any]], dict[str, Any]] + + +class Spec(TypedDict): + stages: list[Literal["diffusion", "autoregression"]] + modes: list[ModelModeSpec] + payload_preprocessor: NotRequired[PayloadPreprocessor] diff --git a/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py new file mode 100644 index 000000000..f1af7b2e9 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py @@ -0,0 +1,84 @@ +from .logger import get_logger +from .models import lookup_model_spec + +logger = get_logger(__name__) + + +def validate_model_and_sampling_params_types( + model_name: str, + sampling_param_list: dict | list[dict] | None = None, +): + # Check if model name exists + if not model_name: + raise ValueError("Model name must not be empty.") + + # 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") + return + if sampling_param_list is None: + return + + # Check the number of stages and their data types + stages = pipeline_spec["stages"] + if isinstance(sampling_param_list, list): + # Check that the lengths match + if len(stages) != len(sampling_param_list): + raise ValueError( + f"Sampling parameter list length {len(sampling_param_list)} does not match " + f"number of stages {len(stages)} for model {model_name}." + ) + # Check that each stage's type match + for i, sp in enumerate(sampling_param_list): + if "type" not in sp: + raise RuntimeError("Internal error: unknown sampling parameter type") + if sp["type"] != stages[i]: + raise ValueError( + f"Sampling parameter type ({sp['type']}) does not match " + f"stage type ({stages[i]}) at index {i} for model {model_name}." + ) + elif isinstance(sampling_param_list, dict): + if "type" not in sampling_param_list: + raise RuntimeError("Internal error: unknown sampling parameter type") + # Check that the provided single sampling param matches all stages + elif any(stage != sampling_param_list["type"] for stage in stages): + raise ValueError( + f"When passing a single sampling parameter node, all stages of the model must match " + f"the provided sampling parameter's type. " + f"However, the stages of model {model_name} are: {stages}. " + f"The provided sampling parameter is {sampling_param_list['type']}" + ) + + +def add_sampling_parameters_to_stage( + model_name: str, + sampling_param_list: dict | list[dict] | None, + stage_type: str, + /, + **params_to_add, +) -> dict | list[dict]: + """ + Given a model's name and the sampling parameter list to query this model, + add arbitrary additional parameters to the sampling parameters of all stages of the given type. + """ + pipeline_spec, _ = lookup_model_spec(model_name) + if not pipeline_spec: + logger.warning( + f"Since the model {model_name} is not in our list, we cannot ensure if " + f"the fields ({tuple(params_to_add.keys())}) are added to the correct stage's sampling params. " + f"We will do it heuristically." + ) + pipeline_spec = {"stages": ["diffusion"]} + + stages = pipeline_spec["stages"] + if isinstance(sampling_param_list, dict): + sampling_param_list.update(params_to_add) + elif sampling_param_list is None: + sampling_param_list = params_to_add.copy() + else: + for i, stage in enumerate(stages): + if stage == stage_type: + sampling_param_list[i].update(params_to_add) + + return sampling_param_list diff --git a/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg new file mode 100644 index 000000000..f2f6c8de1 Binary files /dev/null and b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg differ diff --git a/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg new file mode 100644 index 000000000..5aec26e8c Binary files /dev/null and b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg differ diff --git a/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg new file mode 100644 index 000000000..54eb91b9e Binary files /dev/null and b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg differ diff --git a/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg new file mode 100644 index 000000000..f446de2a0 Binary files /dev/null and b/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg differ diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Annotated Example.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Annotated Example.json new file mode 100644 index 000000000..16657df0e --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Annotated Example.json @@ -0,0 +1 @@ +{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"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":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list","type":"SAMPLING_PARAMS","links":[14]}],"properties":{"Node name for S&R":"VLLMOmniSamplingParamsList"},"widgets_values":[]},{"id":11,"type":"MarkdownNote","pos":[826.2328280438272,569.1890318701705],"size":[333.8220435590464,261.63596728060656],"flags":{},"order":4,"mode":0,"inputs":[],"outputs":[],"title":"Note: Sampling Parameters","properties":{},"widgets_values":["## Sampling Parameter Types\n\nThere are two types of sampling parameters: one for autoregression and one for diffusion.\nYou should ensure that you have chosen the correct type of sampling parameters for the model you request.\n\n## Stages & Shorthand\n\nFor multi-stage models such as Qwen Omni, you can either\n- connect one sampling parameter node, which is applied to all stages.\n- connect exactly the same number of sampling parameter nodes to a \"Multi-Stage Sampling Parameter List\" node, then connect this node to the primary request node.\n\nNote that this shorthand is intended to stay consistent with the [online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/qwen2_5_omni/)"],"color":"#432","bgcolor":"#000"},{"id":12,"type":"MarkdownNote","pos":[378.9866207003777,152.59550252215752],"size":[319.7287574247016,107.15904081906785],"flags":{},"order":6,"mode":0,"inputs":[],"outputs":[],"title":"Note: Input","properties":{},"widgets_values":["Note that not all models support every modality as input. For example, `ByteDance-Seed/BAGEL-7B-MoT` in Multimodality Comprehension mode only support text and image input.\n\nYou should ensure that the input are supported by the model. You can check the corresponding [online serving documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/) for confirmation."],"color":"#432","bgcolor":"#000"},{"id":2,"type":"LoadImage","pos":[394.4674804308822,-207.6987397548834],"size":[282.798828125,314],"flags":{},"order":5,"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":[1]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":10,"type":"PreviewAudio","pos":[1664.548345556043,297.5921292054054],"size":[270,88],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":9},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null}],"outputs":[],"properties":{"Node name for S&R":"PreviewAudio"},"widgets_values":[]},{"id":9,"type":"ShowText|pysssss","pos":[1649.2506875091847,66.22823888292349],"size":[318.7188464232943,173.38502269972975],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"text","name":"text","type":"STRING","link":8}],"outputs":[{"localized_name":"STRING","name":"STRING","shape":6,"type":"STRING","links":null}],"properties":{"Node name for S&R":"ShowText|pysssss"},"widgets_values":[]}],"links":[[1,2,0,1,0,"IMAGE"],[2,3,0,1,1,"VIDEO"],[3,4,0,1,2,"AUDIO"],[5,7,0,8,0,"SAMPLING_PARAMS"],[8,1,0,9,0,"STRING"],[9,1,1,10,0,"AUDIO"],[12,5,0,8,1,"SAMPLING_PARAMS"],[13,7,0,8,2,"SAMPLING_PARAMS"],[14,8,0,1,3,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Sampling Parameters","bounding":[480.1649301181556,341.08402513937995,692.4113568972277,510.48648853403665],"color":"#3f789e","font_size":24,"flags":{}},{"id":2,"title":"Input","bounding":[344.2364817813528,-287.5850484183313,704.8832238768634,559.5124009832894],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":0.9478575057427204,"offset":[33.81037136029557,307.1974296197726]}},"version":0.4} diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Image Generation.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Image Generation.json new file mode 100644 index 000000000..d1495bc94 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Image Generation.json @@ -0,0 +1 @@ +{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"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}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"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":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint 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":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_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":[276.8427259980847,549.4885226647582]}},"version":0.4} diff --git a/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni TTS.json b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni TTS.json new file mode 100644 index 000000000..761e8d147 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni TTS.json @@ -0,0 +1 @@ +{"id":"c0a11241-f0e6-4f5c-8488-b3d344a9df3d","revision":0,"last_node_id":7,"last_link_id":3,"nodes":[{"id":4,"type":"PreviewAudio","pos":[690.1759063763212,226.76444957981545],"size":[270,88],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":2},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null}],"outputs":[],"properties":{"Node name for S&R":"PreviewAudio"},"widgets_values":[]},{"id":5,"type":"LoadAudio","pos":[-83.61093919682321,103.80510898465587],"size":[282.798828125,136],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":3,"type":"VLLMOmniVoiceClone","pos":[248.63466280064262,168.13112893706375],"size":[400,306],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"ref_audio","name":"ref_audio","type":"AUDIO","link":3},{"localized_name":"model_specific_params","name":"model_specific_params","shape":7,"type":"TTS_PARAMS","link":1},{"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":"input","name":"input","type":"STRING","widget":{"name":"input"},"link":null},{"localized_name":"voice","name":"voice","type":"STRING","widget":{"name":"voice"},"link":null},{"localized_name":"response_format","name":"response_format","type":"COMBO","widget":{"name":"response_format"},"link":null},{"localized_name":"speed","name":"speed","type":"FLOAT","widget":{"name":"speed"},"link":null},{"localized_name":"ref_text","name":"ref_text","type":"STRING","widget":{"name":"ref_text"},"link":null},{"localized_name":"x_vector_only_mode","name":"x_vector_only_mode","type":"BOOLEAN","widget":{"name":"x_vector_only_mode"},"link":null}],"outputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","links":[2]}],"properties":{"Node name for S&R":"VLLMOmniVoiceClone"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen3-TTS-12Hz-1.7B-Base","Someone just spilled a cup of coffee on my jacket this morning in the subway. Now I have to wear this stained jacket for a whole day!","","mp3",1,"",true]},{"id":6,"type":"MarkdownNote","pos":[255.52667408769662,27.26842775454861],"size":[381.58204108637847,88],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: TTS Nodes","properties":{},"widgets_values":["Apart from the Voice Cloning node, there is also another TTS node for simple text-to-speech tasks (without extra inputs for reference audio)."],"color":"#432","bgcolor":"#000"},{"id":7,"type":"MarkdownNote","pos":[-196.99453003423977,598.3693693727904],"size":[382.4539509209102,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Model-Specific Parameters","properties":{},"widgets_values":["TTS models often require some tailor-made parameters. If you need to customize these parameters, grab one that matches the model you are requesting from \"TTS Params\" subfolder."],"color":"#432","bgcolor":"#000"},{"id":2,"type":"VLLMOmniQwenTTSParams","pos":[-207.8925310560497,342.7942867650348],"size":[400,200],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"task_type","name":"task_type","type":"COMBO","widget":{"name":"task_type"},"link":null},{"localized_name":"language","name":"language","type":"COMBO","widget":{"name":"language"},"link":null},{"localized_name":"instructions","name":"instructions","type":"STRING","widget":{"name":"instructions"},"link":null},{"localized_name":"max_new_tokens","name":"max_new_tokens","type":"INT","widget":{"name":"max_new_tokens"},"link":null}],"outputs":[{"localized_name":"Qwen TTS params","name":"Qwen TTS params","type":"TTS_PARAMS","links":[1]}],"properties":{"Node name for S&R":"VLLMOmniQwenTTSParams"},"widgets_values":["Base","Auto","Super angry",2048]}],"links":[[1,2,0,3,1,"TTS_PARAMS"],[2,3,0,4,0,"AUDIO"],[3,5,0,3,0,"AUDIO"]],"groups":[{"id":1,"title":"Model-Specific Parameters","bounding":[-228.38328407738007,264.42790274716276,443.9887293164586,439.6853171248897],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486911,"offset":[898.8819468322333,258.7179979718389]}},"version":0.4} diff --git a/apps/ComfyUI-vLLM-Omni/web/main.js b/apps/ComfyUI-vLLM-Omni/web/main.js new file mode 100644 index 000000000..da99b2539 --- /dev/null +++ b/apps/ComfyUI-vLLM-Omni/web/main.js @@ -0,0 +1,21 @@ +/** + * @file This file is intended to add dynamic fields to vLLM-Omni nodes + * based on widget (in-node form fields) and input (connection link) values and changes. + * However, this functionality is currently disabled/commented out. + * Because it introduces too much complexity, + * and it may even conflict with the current backend (Python) validation for unknown reasons (pending ComfyUI upstream fixes). + */ + +import { app } from "../../scripts/app.js"; +app.registerExtension({ + name: "vllm.vllm_omni", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (!nodeData.name.startsWith("VLLMOmni")) { + return + } + // Stub frontend plugin for now + }, + async setup() { + console.info("vLLM-Omni Setup complete!") + }, +}) diff --git a/benchmarks/qwen3-omni/README.md b/benchmarks/qwen3-omni/README.md index fdc385445..c76fea895 100644 --- a/benchmarks/qwen3-omni/README.md +++ b/benchmarks/qwen3-omni/README.md @@ -53,7 +53,7 @@ bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh ``` What it does: -- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--enable-stats`. +- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`. - Uses `benchmarks/build_dataset/top100.txt` and writes to: - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/` - `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats. diff --git a/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh b/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh index 61e46f8c3..935753605 100644 --- a/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh +++ b/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh @@ -23,7 +23,7 @@ else python $end2end_script_path --output-wav $outputs_dir \ --query-type text \ --txt-prompts $build_dataset_path \ - --enable-stats \ + --log-stats \ --log-dir $log_dir echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:" echo " - omni_llm_pipeline_text run dir/base name" diff --git a/docs/.nav.yml b/docs/.nav.yml index 41754c9c7..16821c595 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -44,6 +44,7 @@ nav: - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md + - ComfyUI: features/comfyui.md - Developer Guide: - General: - contributing/README.md @@ -53,6 +54,12 @@ nav: - contributing/model/README.md - contributing/model/adding_omni_model.md - contributing/model/adding_diffusion_model.md + - Advanced Features: + - contributing/features/cfg_parallel.md + - contributing/features/sequence_parallel.md + - contributing/features/tensor_parallel.md + - contributing/features/cache_dit.md + - contributing/features/teacache.md - CI: contributing/ci - Design Documents: - design/index.md diff --git a/docs/api/README.md b/docs/api/README.md index 6219aee40..eb3a40950 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -15,9 +15,6 @@ Main entry points for vLLM-Omni inference and serving. - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][] - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][] - [vllm_omni.entrypoints.client_request_state.ClientRequestState][] -- [vllm_omni.entrypoints.log_utils.OrchestratorMetrics][] -- [vllm_omni.entrypoints.log_utils.StageRequestMetrics][] -- [vllm_omni.entrypoints.log_utils.StageStats][] - [vllm_omni.entrypoints.omni.Omni][] - [vllm_omni.entrypoints.omni.OmniBase][] - [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][] @@ -114,3 +111,8 @@ Worker classes and model runners for distributed inference. - [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][] - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][] - [vllm_omni.worker.mixins.OmniWorkerMixin][] + + +## Metrics + +- [vllm_omni.metrics.OrchestratorAggregator][] diff --git a/docs/contributing/ci/tests_markers.md b/docs/contributing/ci/tests_markers.md index bf56914f8..2afff8460 100644 --- a/docs/contributing/ci/tests_markers.md +++ b/docs/contributing/ci/tests_markers.md @@ -5,33 +5,33 @@ By adding markers before test functions, tests can later be executed uniformly b ## Current Markers Defined in `pyproject.toml`: -| Marker | Description | -| ------------------ | ------------------------------------------------------- | -| `core_model` | Core model tests (run in each PR) | -| `diffusion` | Diffusion model tests | -| `omni` | Omni model tests | -| `cache` | Cache backend tests | -| `parallel` | Parallelism/distributed tests | -| `cpu` | Tests that run on CPU | -| `gpu` | Tests that run on GPU (auto-added) | -| `cuda` | Tests that run on CUDA (auto-added) | -| `rocm` | Tests that run on AMD/ROCm (auto-added) | -| `npu` | Tests that run on NPU/Ascend (auto-added) | -| `H100` | Tests that require H100 GPU | -| `L4` | Tests that require L4 GPU | -| `MI325` | Tests that require MI325 GPU (AMD/ROCm) | -| `A2` | Tests that require A2 NPU | -| `A3` | Tests that require A3 NPU | -| `distributed_cuda` | Tests that require multi cards on CUDA platform | -| `distributed_rocm` | Tests that require multi cards on ROCm platform | -| `distributed_npu` | Tests that require multi cards on NPU platform | -| `skipif_cuda` | Skip if the num of CUDA cards is less than the required | -| `skipif_rocm` | Skip if the num of ROCm cards is less than the required | -| `skipif_npu` | Skip if the num of NPU cards is less than the required | -| `slow` | Slow tests (may skip in quick CI) | -| `benchmark` | Benchmark tests | +| Marker | Description | +| ------------------ | --------------------------------------------------------- | +| `core_model` | Core model tests (run in each PR) | +| `diffusion` | Diffusion model tests | +| `omni` | Omni model tests | +| `cache` | Cache backend tests | +| `parallel` | Parallelism/distributed tests | +| `cpu` | Tests that run on CPU | +| `gpu` | Tests that run on GPU * | +| `cuda` | Tests that run on CUDA * | +| `rocm` | Tests that run on AMD/ROCm * | +| `npu` | Tests that run on NPU/Ascend * | +| `H100` | Tests that require H100 GPU * | +| `L4` | Tests that require L4 GPU * | +| `MI325` | Tests that require MI325 GPU (AMD/ROCm) * | +| `A2` | Tests that require A2 NPU * | +| `A3` | Tests that require A3 NPU * | +| `distributed_cuda` | Tests that require multi cards on CUDA platform * | +| `distributed_rocm` | Tests that require multi cards on ROCm platform * | +| `distributed_npu` | Tests that require multi cards on NPU platform * | +| `skipif_cuda` | Skip if the num of CUDA cards is less than the required * | +| `skipif_rocm` | Skip if the num of ROCm cards is less than the required * | +| `skipif_npu` | Skip if the num of NPU cards is less than the required * | +| `slow` | Slow tests (may skip in quick CI) | +| `benchmark` | Benchmark tests | -For those markers shown as auto-added, they will be added by the `@hardware_test` decorator. +\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator. ### Example usage for markers @@ -71,10 +71,7 @@ This decorator is intended to make hardware-aware, cross-platform test authoring Support for `skipif_rocm` and `skipif_npu` will be implemented later. -5. **Runs each test in a new process** - Automatically wraps the distributed test with a decorator (`@create_new_process_for_each_test`) to ensure isolation and compatibility with multi-process hardware backends. - -6. **Works with pytest filtering** +5. **Works with pytest filtering** Allows tests to be filtered and selected at runtime using standard pytest marker expressions (e.g., `-m "distributed_cuda and L4"`). #### Example usage for decorator @@ -94,7 +91,6 @@ This decorator is intended to make hardware-aware, cross-platform test authoring ``` - `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), NPU (A2/A3) - `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing -- `hardware_test` automatically applies `@create_new_process_for_each_test` for distributed tests. - Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases - Filtering examples: - CUDA only: `pytest -m "distributed_cuda and L4"` diff --git a/docs/contributing/features/cache_dit.md b/docs/contributing/features/cache_dit.md new file mode 100644 index 000000000..f39b93578 --- /dev/null +++ b/docs/contributing/features/cache_dit.md @@ -0,0 +1,286 @@ +# Support Cache-DiT + +This section describes how to add cache-dit acceleration to a new diffusion pipeline. We use the Qwen-Image pipeline and LongCat-Image pipeline as reference implementations. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Standard Models: Automatic Support](#standard-models-automatic-support) +- [Custom Architectures: Writing Custom Implementation](#custom-architectures-writing-custom-implementation) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is Cache-DiT? + +Cache-DiT is an acceleration library for Diffusion Transformers (DiT) that caches intermediate computation results across denoising steps. The core insight is that adjacent denoising steps often produce similar intermediate features, so we can skip redundant computations by reusing cached results. + +The library supports three main caching strategies: + +- **DBCache:** Dynamic block-level caching that selectively computes or caches transformer blocks based on residual differences +- **TaylorSeer:** Calibration-based prediction that estimates block outputs using Taylor expansion +- **SCM (Step Computation Masking):** Dynamic step skipping based on configurable policies + +### Architecture + +vLLM-omni integrates cache-dit through the `CacheDiTBackend` class, which provides a unified interface for managing cache-dit acceleration on diffusion models. + +| Method/Class | Purpose | Behavior | +|--------------|---------|----------| +| [`CacheDiTBackend`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.CacheBackend) | Unified backend interface | Automatically handles enabler selection and cache refresh | +| [`enable_cache_for_dit()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Apply caching to transformer | Configures DBCache on transformer blocks | + +**Key APIs from Cache-DiT:** + +[Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/) + +| API | Description | +|-----|-------------| +| `BlockAdapter` | Core abstraction for applying cache-dit to transformers. Specifies transformer module(s), block list(s), and forward signature pattern(s). | +| `ForwardPattern` | Defines block forward signature patterns: `Pattern_0`, `Pattern_1`, `Pattern_2` | +| `ParamsModifier` | Per-transformer or per-block-list cache configuration customization | +| `DBCacheConfig` | Configuration for DBCache parameters (warmup steps, cached steps, thresholds) | +| `refresh_context()` | Update cache context | Called when `num_inference_steps` changes | + +--- + +## Standard Models: Automatic Support + +Most DiT models follow this pattern: +- Single transformer with one `ModuleList` of blocks +- Standard forward signature +- Compatible with cache-dit's automatic detection + +**Examples:** Qwen-Image, Z-Image + +For standard single-transformer models, **no code changes are needed**. The `CacheDiTBackend` automatically uses `enable_cache_for_dit()`: + +```python +from vllm_omni import Omni + +# Works automatically for standard models +omni = Omni( + model="Qwen/Qwen-Image", # Standard single-transformer model + cache_backend="cache_dit", + cache_config={ + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + } +) +``` + +**What happens automatically:** + +```python +def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], None]: + """Default enabler for standard single-transformer DiT models.""" + + # Build cache configuration + db_cache_config = DBCacheConfig( + num_inference_steps=None, # Will be set during first inference + Fn_compute_blocks=cache_config.Fn_compute_blocks, + Bn_compute_blocks=cache_config.Bn_compute_blocks, + max_warmup_steps=cache_config.max_warmup_steps, + max_cached_steps=cache_config.max_cached_steps, + max_continuous_cached_steps=cache_config.max_continuous_cached_steps, + residual_diff_threshold=cache_config.residual_diff_threshold, + ) + + # Enable cache-dit on transformer + cache_dit.enable_cache( + pipeline.transformer, + cache_config=db_cache_config, + ) + + # Return refresh function for dynamic num_inference_steps updates + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True): + cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) + + return refresh_cache_context +``` + +--- + +## Custom Architectures: Writing Custom Implementation + +Some models require custom handling: + +- **Dual-transformer:** Models with separate high-noise and low-noise transformers (e.g., Wan2.2) +- **Multi-block-list:** Models with multiple block lists in one transformer (e.g., LongCatImage with `transformer_blocks` + `single_transformer_blocks`) +- **Special forward patterns:** Models with non-standard block execution patterns + +### Example 1: Dual-Transformer Model (Wan2.2) + +Wan2.2 uses two transformers: one for high-noise steps and one for low-noise steps. + +**Key difference:** Use `BlockAdapter` to wrap multiple transformers with separate configurations. + +```python +# Standard: cache_dit.enable_cache(pipeline.transformer, ...) +# Custom: Use BlockAdapter to handle multiple transformers +cache_dit.enable_cache( + BlockAdapter( + transformer=[pipeline.transformer, pipeline.transformer_2], # Multiple transformers + blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks], + forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2], + params_modifiers=[ + ParamsModifier(...), # Config for high-noise transformer + ParamsModifier(...), # Config for low-noise transformer (different params) + ], + ), + cache_config=db_cache_config, +) +``` + +**Key difference:** `refresh_context` must be called on each transformer separately. + +```python +# Standard: cache_dit.refresh_context(pipeline.transformer, num_inference_steps=N) +# Custom: Refresh each transformer with its own step count +def refresh_cache_context(pipeline, num_inference_steps, verbose=True): + high_steps, low_steps = _split_inference_steps(num_inference_steps) + cache_dit.refresh_context(pipeline.transformer, num_inference_steps=high_steps, ...) + cache_dit.refresh_context(pipeline.transformer_2, num_inference_steps=low_steps, ...) +``` + +### Example 2: Multi-Block-List Model (LongCatImage) + +LongCatImage has a single transformer with two block lists: `transformer_blocks` and `single_transformer_blocks`. + +**Key difference:** Use `BlockAdapter` to specify multiple block lists within one transformer. + +```python +# Standard: cache_dit.enable_cache(pipeline.transformer, ...) +# - Automatically detects single block list +# Custom: Use BlockAdapter to specify multiple block lists +cache_dit.enable_cache( + BlockAdapter( + transformer=pipeline.transformer, # Single transformer + blocks=[ + pipeline.transformer.transformer_blocks, # Block list 1 + pipeline.transformer.single_transformer_blocks, # Block list 2 + ], + forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1], + params_modifiers=[modifier], + ), + cache_config=db_cache_config, +) +``` + +> **Note:** For single transformer with multiple block lists, `refresh_context` works the same as standard models. + +### Registering Custom Implementations + +After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`: + +```python +CUSTOM_DIT_ENABLERS = { + "Wan22Pipeline": enable_cache_for_wan22, + "LongCatImagePipeline": enable_cache_for_longcat_image, + "YourCustomPipeline": enable_cache_for_your_model, # Add here +} +``` + +--- + +## Testing + +After adding cache-dit support, test with: + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +# Test your custom model +omni = Omni( + model="your-model-name", + cache_backend="cache_dit", + cache_config={ + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "residual_diff_threshold": 0.24, + } +) + +images = omni.generate( + "a beautiful landscape", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +**Verify:** + +1. Cache is applied (check logs for "Cache-dit enabled successfully on xxx") +2. Performance improvement (should be around 1.5x-2x faster) +3. Image quality (compare with `cache_backend=None`) + +--- + +## Troubleshooting + +### Issue: Cache not applied + +**Symptoms:** No speedup observed, no cache-related log messages. + +**Causes & Solutions:** + +- **Enabler not registered:** + +**Problem:** Pipeline name not in `CUSTOM_DIT_ENABLERS` registry. + +**Solution:** Verify `pipeline.__class__.__name__` matches the registry key and add your enabler to `CUSTOM_DIT_ENABLERS`. + +### Issue: Quality degradation + +**Symptoms:** Generated images have artifacts or lower quality compared to non-cached inference. + +**Causes & Solutions:** + +- **Cache parameters too aggressive:** + +**Solution:** +```python +cache_config={ + "residual_diff_threshold": 0.12, # Lower from 0.24 (try 0.12-0.18) + "max_warmup_steps": 6, # Increase from 4 (try 6-8) + "max_continuous_cached_steps": 2, # Reduce if higher +} +``` + +Check the [user guide for cache_dit](../../user_guide/diffusion/cache_dit_acceleration.md) for more adjustable parameters. + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic | +| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Dual-transformer | Separate high/low noise transformers | +| **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer | +| **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture | + +--- + +## Summary + +Adding cache-dit support: + +1. ✅ **Check model type** - Standard models work automatically, custom architectures need enablers +2. ✅ **Write enabler** (if needed) - Use `BlockAdapter` for complex architectures +3. ✅ **Register enabler** (if needed) - Add to `CUSTOM_DIT_ENABLERS` dictionary +4. ✅ **Return refresh function** (if needed) - Handle `num_inference_steps` changes +5. ✅ **Test** - Verify with `cache_backend="cache_dit"` + +For most models, the default enabler is sufficient. Only write custom enablers for complex architectures! diff --git a/docs/contributing/features/cfg_parallel.md b/docs/contributing/features/cfg_parallel.md new file mode 100644 index 000000000..adfe447c9 --- /dev/null +++ b/docs/contributing/features/cfg_parallel.md @@ -0,0 +1,249 @@ +# Support CFG-Parallel + +This section describes how to add CFG-Parallel (Classifier-Free Guidance Parallel) to a diffusion pipeline. We use the Qwen-Image pipeline as the reference implementation. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Step-by-Step Implementation](#step-by-step-implementation) +- [Customization](#customization) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is CFG-Parallel? + +In standard Classifier-Free Guidance, each diffusion step requires two forward passes through the transformer: + +1. **Positive/Conditional**: Guided by the text prompt +2. **Negative/Unconditional**: Typically using empty or negative prompt + +CFG-Parallel eliminates this bottleneck by distributing the two forward passes across different GPU ranks, allowing them to execute simultaneously rather than sequentially. + +### Architecture + +vLLM-omni provides `CFGParallelMixin` that encapsulates all CFG parallel logic. Pipelines inherit from this mixin and implement a `diffuse()` method that orchestrates the denoising loop. + +| Method | Purpose | Automatic Behavior | +|--------|---------|-------------------| +| [`predict_noise_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Predict noise with CFG | Detects parallel mode, distributes computation, gathers results | +| [`scheduler_step_maybe_with_cfg()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Step scheduler with sync | Rank 0 steps, broadcasts latents to all ranks | +| [`combine_cfg_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Combine positive/negative | Applies CFG formula with optional normalization | +| [`predict_noise()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Forward pass wrapper | Override for custom transformer calls | +| [`cfg_normalize_function()`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/distributed/cfg_parallel/) | Normalize CFG output | Override for custom normalization | + +### How It Works + +`predict_noise_maybe_with_cfg()` automatically detects and switches between two execution modes: + +- **CFG-Parallel mode** (when `cfg_world_size > 1`): + - Rank 0 computes positive prompt prediction + - Rank 1 computes negative prompt prediction + - Results are gathered via `all_gather()` + - Combined on rank 0 using CFG formula + +- **Sequential mode** (when `cfg_world_size == 1`): + - Single rank computes both positive and negative predictions + - Directly combines them with CFG formula + +`scheduler_step_maybe_with_cfg()` ensures consistent latent states across all ranks: + +- **CFG-Parallel mode**: + - Only rank 0 performs the scheduler step (applies noise prediction to update latents) + - Updated latents are broadcast to all other ranks via `broadcast()` + - All ranks maintain synchronized latent states for the next iteration + +- **Sequential mode**: + - Single rank directly performs the scheduler step + - No synchronization needed + +--- + +## Step-by-Step Implementation + +### Step 1: Inherit `CFGParallelMixin` + +Allow your pipeline to inherit from `CFGParallelMixin` and implements the `diffuse()` method for your specific model. + +**Example (Qwen-Image):** + +```python +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +import torch.nn as nn +class YourModelPipeline(nn.Module, CFGParallelMixin): + def diffuse(self, ...) -> torch.Tensor: + for i, t in enumerate(timesteps): + # Prepare positive_kwargs (conditional) and negative_kwargs (unconditional) + positive_kwargs = {...} # hidden_states, encoder_hidden_states, etc. + negative_kwargs = {...} if do_true_cfg else None + + # Key method 1: Predict noise with automatic CFG parallel handling + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=do_true_cfg, + true_cfg_scale=true_cfg_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + ) + + # Key method 2: Step scheduler with automatic CFG synchronization + latents = self.scheduler_step_maybe_with_cfg( + noise_pred, t, latents, do_true_cfg + ) + + return latents +``` + +**Key Points:** + +- `positive_kwargs`: transformer arguments for conditional (text-guided) prediction +- `negative_kwargs`: transformer arguments for unconditional prediction (set to `None` if CFG disabled) +- For image editing pipelines, add `output_slice=image_seq_len` to extract the generative image portion + +### Step 2: Call `diffuse` + +Call `self.diffuse` in your pipeline's forward function: + +```python +import torch.nn as nn +class YourModelPipeline(nn.Module, CFGParallelMixin): + def forward( + self, + prompt: str, + negative_prompt: str | None = None, + guidance_scale: float = 3.5, + num_inference_steps: int = 50, + **kwargs, + ): + # Encode prompts, Initialize latents, Get timesteps + ... + # Run diffusion loop (calls the mixin's diffuse method) + latents = self.diffuse( + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + negative_prompt_embeds=negative_embeds, + negative_prompt_embeds_mask=negative_mask, + latents=latents, + timesteps=timesteps, + do_true_cfg=do_true_cfg, + true_cfg_scale=guidance_scale, + ... + ) +``` + +--- + +## Customization + +### Override `predict_noise()` for Custom Transformer Calls + +If your transformer requires custom prediction function, you can rewrite `predict_noise` function. Taking Wan2.2 as an example, which has two transformer models. The actual transformer to be called is determined by `self.transformer`. + +```python +class Wan22Pipeline(nn.Module, CFGParallelMixin): + def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor: + if current_model is None: + current_model = self.transformer + return current_model(**kwargs)[0] +``` + + +### Override `cfg_normalize_function()` for Custom Normalization + +Some models have their own normalization function. Taking LongCat Image model as an example: + +```python +class LongCatImagePipeline(nn.Module, CFGParallelMixin): + def cfg_normalize_function(self, noise_pred, comb_pred, cfg_renorm_min=0.0): + """ + Normalize the combined noise prediction. + """ + cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True) + noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True) + scale = (cond_norm / (noise_norm + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) + noise_pred = comb_pred * scale + return noise_pred + + # The original cfg_normalize_function function in CFGParallelMixin + # cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True) + # noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True) + # noise_pred = comb_pred * (cond_norm / noise_norm) + # return noise_pred +``` + +--- + +## Testing + +After adding CFG-Parallel support, test with: + +```bash +cd examples/offline_inference/text_to_image +python text_to_image.py \ + --model Your-org/your-model \ + --prompt "a cup of coffee on the table" \ + --negative_prompt "ugly, unclear" \ + --cfg_scale 4.0 \ + --num_inference_steps 50 \ + --output "cfg_enabled.png" \ + --cfg_parallel_size 2 +``` + +**Verify:** + +1. Check logs for CFG parallel being activated +2. Record the `e2e_time_ms` in the log and compare with CFG-Parallel disabled +3. Compare the generated result quality with baseline +4. Record comparison results in your PR + +--- + +## Troubleshooting + +### Issue: CFG parallel not activating + +**Symptoms:** Generation still slow, logs don't show CFG parallel being used. + +**Causes & Solutions:** + +- **CFG is not enabled:** + +**Problem:** Guidance scale too low or negative prompt not provided. + +**Solution:** Ensure `guidance_scale > 1.0` and negative prompt is provided: +```python +images = pipeline( + prompt="a cat", + negative_prompt="", # Must provide (even if empty) + guidance_scale=3.5, # Must be > 1.0 +) +``` + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/cfg_parallel.py` | Mixin | Dual-stream transformer | +| **Qwen-Image-Edit** | `vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py` | Mixin | Image editing with `output_slice` | +| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` | Mixin | Dual-transformer architecture | +| **CFGParallelMixin** | `vllm_omni/diffusion/distributed/cfg_parallel.py` | Base implementation | Core mixin class | + +--- + +## Summary + +Adding CFG-Parallel support: + +1. ✅ **Create mixin** - Inherit from `CFGParallelMixin` and implement `diffuse()` method +2. ✅ **(Optional) Customize** - Override `predict_noise()` or `cfg_normalize_function()` for custom behavior +3. ✅ **Test** - Verify with `--cfg_parallel_size 2` and compare performance diff --git a/docs/contributing/features/sequence_parallel.md b/docs/contributing/features/sequence_parallel.md new file mode 100644 index 000000000..3bcee7345 --- /dev/null +++ b/docs/contributing/features/sequence_parallel.md @@ -0,0 +1,405 @@ +# Support Sequence Parallel + +This section describes how to add Sequence Parallel (SP) to a diffusion transformer model. We use the Qwen-Image transformer and Wan2.2 transformer as reference implementations. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Approach 1: Non-Intrusive `_sp_plan` (Recommended)](#approach-1-non-intrusive-_sp_plan-recommended) +- [Approach 2: Intrusive Modification (For Complex Cases)](#approach-2-intrusive-modification-for-complex-cases) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + + +### What is Sequence Parallel? + +**Terminology Note:** Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py). We use "Sequence Parallelism" to align with vLLM-Omni's terminology. + +Diffusion transformers process long sequences of image patches or video frames. For high-resolution generation, these sequences can become very large. Enabling SP allows each GPU to process only a portion of the sequence, with attention mechanisms (Ulysses/Ring) handling cross-GPU communication transparently. + +### Architecture + +The major APIs for Sequence Parallel: + +```python +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, # For sharding (splitting) tensors + SequenceParallelOutput, # For gathering tensors +) +from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather +``` + +| Method/Class | Purpose | Behavior | +|--------------|---------|----------| +| `SequenceParallelInput` | Declare input sharding in `_sp_plan` | Auto-shards tensors at module input | +| `SequenceParallelOutput` | Declare output gathering in `_sp_plan` | Auto-gathers tensors at module output | +| `sp_shard()` | Manual tensor sharding | Splits tensor across SP workers | +| `sp_gather()` | Manual tensor gathering | Gathers sharded tensors from all workers | + +--- + +## Approach 1: Non-Intrusive `_sp_plan` (Recommended) + +The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries. + +**When to use:** +- Standard transformer architectures +- Tensor operations happen at `nn.Module` boundaries +- Predictable sharding/gathering patterns + +**How it works:** +1. Declare `_sp_plan` dict in your transformer class +2. Framework automatically applies hooks when `sequence_parallel_size > 1` +3. Hooks shard/gather tensors at specified module boundaries +4. Attention layers handle cross-GPU communication internally + +```python +class StandardTransformer(nn.Module): + _sp_plan = { + # Shard hidden_states at first transformer block input + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Gather at final output projection + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } +``` + +`StandardTransformer` has a transformer blocks list `self.blocks = nn.ModuleList([...])`, and a projection output layer `self.proj_out`. The `_sp_plan` above defines that when SP is enabled, sharding the input tensor to the first transformer block, and gathering the sharded tensor at the final output projection layer. + +**Requirements:** +- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** +- Inline Python operations (e.g., `torch.cat`, `pad_sequence`) **cannot be hooked** + +**Solution for inline operations:** Extract into a submodule (see Step 2 below). + +### Step 1: Understand Module Boundaries + +Identify where tensors need to be sharded or gathered in your model's forward pass: + +```python +class MyTransformer(nn.Module): + def __init__(self): + self.patch_embed = PatchEmbed() # ← Boundary 1 + self.pos_embed = RoPE() # ← Boundary 2 + self.blocks = nn.ModuleList([...]) # ← Boundary 3 + self.norm_out = LayerNorm() + self.proj_out = Linear() # ← Boundary 4 + + def forward(self, x): + x = self.patch_embed(x) # ← Shard before this? + pos = self.pos_embed(x) # ← Shard RoPE outputs? + for block in self.blocks: + x = block(x, pos) # ← Blocks process sharded x + x = self.norm_out(x) + output = self.proj_out(x) # ← Gather after this? + return output +``` + +### Step 2: Handle Inline Operations + +If your `forward()` contains inline tensor operations, **extract them into submodules**. + +**Example: Z-Image concatenates image + text features inline** + +```python +# ❌ BAD: Inline operation - hooks cannot intercept +class ZImageTransformer(nn.Module): + def forward(self, x, cap_feats): + # This concatenation happens inline - _sp_plan can't shard it! + unified = torch.cat([x, cap_feats], dim=1) + + for layer in self.layers: + unified = layer(unified) + + return unified + +# ✅ GOOD: Extract into submodule +class UnifiedPrepare(nn.Module): + """Submodule to concatenate image and text features.""" + def forward(self, x, cap_feats): + return torch.cat([x, cap_feats], dim=1) + +class ZImageTransformer(nn.Module): + def __init__(self): + super().__init__() + self.unified_prepare = UnifiedPrepare() # Now a module! + self.layers = nn.ModuleList([...]) + + def forward(self, x, cap_feats): + # Now _sp_plan can shard the output of unified_prepare! + unified = self.unified_prepare(x, cap_feats) + + for layer in self.layers: + unified = layer(unified) + + return unified +``` + +**Other common cases:** +- `pad_sequence()` → `PadSequenceModule` +- `torch.cat()` → `ConcatModule` +- `tensor.reshape()` → `ReshapeModule` +- Complex preprocessing → `PreprocessModule` + +### Step 3: Write `_sp_plan` for Your Model + +Create a class-level `_sp_plan` dictionary specifying where to shard/gather tensors. + +Typically, there are two patterns for diffusion models: + +**Pattern 1: Shard at first block, gather at output projection** + +Most common pattern for standard transformers: + +```python +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, # For sharding (splitting) tensors + SequenceParallelOutput, # For gathering tensors +) +class StandardTransformer(nn.Module): + _sp_plan = { + # Shard hidden_states at first transformer block input + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Gather at final output projection + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } +``` + +**Pattern 2: Shard RoPE embeddings separately** + +When RoPE is computed in a separate module: + +```python +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, # For sharding (splitting) tensors + SequenceParallelOutput, # For gathering tensors +) +class TransformerWithRoPE(nn.Module): + _sp_plan = { + # Shard RoPE module OUTPUTS (returns tuple of cos, sin) + "rope": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # cos + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # sin + }, + # Shard transformer block INPUT + "blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + # Gather at output + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } +``` + +### API Reference + +**SequenceParallelInput Parameters:** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `split_dim` | int | Dimension to split (usually `1` for sequence) | +| `expected_dims` | int \| None | Expected tensor rank for validation (optional) | +| `split_output` | bool | `False`: shard **input** params; `True`: shard **output** tensors | +| `auto_pad` | bool | Auto-pad if sequence not divisible by world_size (default: `False`) | + +**SequenceParallelOutput Parameters:** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `gather_dim` | int | Dimension to gather (usually `1` for sequence) | +| `expected_dims` | int \| None | Expected tensor rank for validation (optional) | + +**Module Naming Conventions:** + +| Key | Meaning | Python equivalent | +|-----|---------|-------------------| +| `""` | Root model | `model` | +| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` | +| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` | +| `"rope"` | Named submodule | `model.rope` | +| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` | + +**Dictionary Value Types:** + +| Key type | `split_output` | Description | +|----------|----------------|-------------| +| `"param_name"` (str) | `False` | Shard **input parameter** by name | +| `0`, `1`, ... (int) | `True` | Shard **output tuple** by index | + +--- + +## 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. + +**When to use:** +- Dynamic/conditional sharding logic +- Complex tensor manipulations that can't be encapsulated +- Temporary workaround during development + +```python +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 +``` + +--- + +## Testing + +After implementing Sequence Parallel support, thoroughly test your implementation to ensure correctness and performance across different configurations. + +**Test Different `sp_size`:** + +Test your model with various sequence parallel world sizes to verify correctness and identify optimal configurations: + +```bash +cd examples/offline_inference/text_to_image +python text_to_image.py \ + --model Your-org/your-model \ + --prompt "a cup of coffee on the table" \ + --num_inference_steps 50 \ + --ulysses_degree 2 \ + --ring_degree 2 \ + --output sp_test_image_ulysses=2_ring=2.png +``` + +**Verify:** + +1. **Correctness:** Output should be identical across all `sp_size` values +2. **Speed:** Throughput should remain stable or improve (especially for large sequences) +3. **Logs:** Check for any shape mismatch or communication errors + +**Test with Tensor Parallel:** + +Sequence Parallel can be combined with other parallelism strategies: + +```bash +cd examples/offline_inference/text_to_image +python text_to_image.py \ + --model Your-org/your-model \ + --prompt "a cup of coffee on the table" \ + --num_inference_steps 50 \ + --ulysses_degree 2 \ + --tensor_parallel_size 2 \ + --output sp_test_image_ulysses=2_tp=2.png +``` + +--- + +## Troubleshooting + +### Issue: Shape mismatch errors + +**Symptoms:** `RuntimeError: shape mismatch` during forward pass. + +**Causes & Solutions:** + +- **RoPE dimension mismatch:** + +**Problem:** RoPE embeddings not sharded, but hidden_states is sharded. + +**Solution:** Shard RoPE outputs in `_sp_plan`: +```python +_sp_plan = { + "rope": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + }, + ... +} +``` + +- **Sequence Length not divisible by sp_size:** + +**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding. + +**Solution:** In `SequenceParallelInput`, set `auto_pad=True`: +```python +"blocks.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True) +} +``` + +### Issue: Inline operations not sharded + +**Symptoms:** Some tensors remain full-sized, not sharded. + +**Causes & Solutions:** + +- **Operations happen inline in `forward()`, not at module boundaries:** + +**Problem:** +```python +def forward(self, x, cap): + unified = torch.cat([x, cap], dim=1) # ← Inline operation! + # _sp_plan can't hook this +``` + +**Solution:** Extract into submodule: +```python +class ConcatModule(nn.Module): + def forward(self, x, cap): + return torch.cat([x, cap], dim=1) + +class MyModel(nn.Module): + _sp_plan = { + "concat": { + 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), + }, + ... + } + def __init__(self): + self.concat = ConcatModule() # Now hookable! + + def forward(self, x, cap): + unified = self.concat(x, cap) # ← Can be sharded via _sp_plan +``` + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **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 | +| **SP Plan Types** | `vllm_omni/diffusion/distributed/sp_plan.py` | Type definitions | SequenceParallelInput/Output | +| **Hook Implementation** | `vllm_omni/diffusion/hooks/sequence_parallel.py` | Hook mechanics | How hooks work | +| **Tests** | `tests/diffusion/distributed/test_sp_plan_hooks.py` | Test examples | Validation patterns | + +--- + +## Summary + +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? +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 +6. ✅ **Test** - Verify with different `ulysses_degree` and `ring_degree` combinations diff --git a/docs/contributing/features/teacache.md b/docs/contributing/features/teacache.md new file mode 100644 index 000000000..6c29c0f99 --- /dev/null +++ b/docs/contributing/features/teacache.md @@ -0,0 +1,459 @@ +# Support TeaCache + +This section describes how to add TeaCache to a diffusion transformer model. We use the Qwen-Image transformer as the reference implementation. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Step-by-Step Implementation](#step-by-step-implementation) +- [Customization](#customization) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is TeaCache? + +TeaCache speeds up diffusion inference by caching transformer block computations when consecutive timesteps are similar. It provides **1.5x-2.0x speedup** with minimal quality loss. + +The core insight is that the modulated input (after normalization and timestep conditioning) changes gradually across timesteps. By measuring the L1 distance between consecutive modulated inputs and comparing it to a threshold, TeaCache decides whether to execute the full transformer blocks or reuse the cached residual from the previous step. + +vLLM-omni provides a **hook-based** TeaCache system that requires **zero changes to model code**. The hook completely intercepts the transformer's forward pass and implements adaptive caching transparently. This design allows easy integration with any transformer model by simply writing an extractor function. + +### Architecture + +The TeaCache system consists of three main components: + +| Component | Purpose | Location | +|-----------|---------|----------| +| [`CacheContext`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.CacheContext) | Dataclass containing model-specific information for caching | `vllm_omni/diffusion/cache/teacache/context.py` | +| [`EXTRACTOR_REGISTRY`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/teacache/extractors/#vllm_omni.diffusion.cache.teacache.extractors.EXTRACTOR_REGISTRY) | Maps transformer class names to extractor functions | `vllm_omni/diffusion/cache/teacache/extractors.py` | +| [`TeaCacheConfig`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/#vllm_omni.diffusion.cache.TeaCacheConfig) | Configuration including thresholds and polynomial coefficients | `vllm_omni/diffusion/cache/teacache/config.py` | + +The hook handles all caching logic automatically, including: + +- CFG-aware state management (separate states for positive/negative branches) +- CFG-parallel compatibility +- L1 distance computation with polynomial rescaling +- Residual caching and reuse + + +--- + +## Step-by-Step Implementation + +To add TeaCache support for a new model, you need to: + +1. Write an **extractor function** that returns a `CacheContext` object +2. Register the extractor in the `EXTRACTOR_REGISTRY` +3. Add model-specific polynomial coefficients to `TeaCacheConfig` + +### Step 1: Model-Specific Preprocessing + +Extract and process model inputs. This typically involves: +- Embedding image/latent inputs +- Processing text encoder outputs (if dual-stream) +- Creating timestep embeddings +- Computing positional embeddings + +**Example (Qwen-Image):** + +```python +def extract_qwen_context( + module: nn.Module, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_mask: torch.Tensor, + timestep: torch.Tensor, + img_shapes: torch.Tensor, + txt_seq_lens: torch.Tensor, + guidance: torch.Tensor | None = None, + **kwargs: Any, +) -> CacheContext: + # Validate model structure + if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0: + raise ValueError("Module must have transformer_blocks") + + # Preprocessing: embed inputs + hidden_states = module.img_in(hidden_states) + timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) + encoder_hidden_states = module.txt_norm(encoder_hidden_states) + encoder_hidden_states = module.txt_in(encoder_hidden_states) + + # Create timestep embedding + if guidance is not None: + guidance = guidance.to(hidden_states.dtype) * 1000 + temb = ( + module.time_text_embed(timestep, hidden_states) + if guidance is None + else module.time_text_embed(timestep, guidance, hidden_states) + ) + + # Compute position embeddings + image_rotary_emb = module.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device) +``` + +### Step 2: Extract Modulated Input + +The modulated input is used for cache decisions. Extract it from the **first transformer block** after normalization and modulation. + +**Example (Qwen-Image):** + +```python + # Extract modulated input from first transformer block + block = module.transformer_blocks[0] + img_mod_params = block.img_mod(temb) + img_mod1, _ = img_mod_params.chunk(2, dim=-1) + img_modulated, _ = block.img_norm1(hidden_states, img_mod1) +``` + +**Key Points:** + +- Use the **first block** to extract modulated input early +- Apply the same normalization and modulation as the actual forward pass +- The tensor should represent the processed features that will change across timesteps + +### Step 3: Define Transformer Execution + +Create a callable that executes all transformer blocks. This encapsulates the main computation loop. + +**Example (Qwen-Image dual-stream):** + +```python + def run_transformer_blocks(): + """Execute all Qwen transformer blocks.""" + h = hidden_states + e = encoder_hidden_states + + for block in module.transformer_blocks: + e, h = block( + hidden_states=h, + encoder_hidden_states=e, + encoder_hidden_states_mask=encoder_hidden_states_mask, + temb=temb, + image_rotary_emb=image_rotary_emb, + ) + return (h, e) # Return both image and text hidden states +``` + +**Example (Single-stream model like Flux):** + +```python + def run_transformer_blocks(): + """Execute all Flux transformer blocks.""" + h = hidden_states + + for block in module.transformer_blocks: + h = block(h, temb=temb) + return (h,) # Return only image hidden states +``` + +**Key Points:** + +- Return format: +- For single-stream models: return `(hidden_states,)` +- For dual-stream models: return `(hidden_states, encoder_hidden_states)` + +### Step 4: Define Postprocessing + +Create a callable that applies final transformations to produce the model output. + +**Example (Qwen-Image):** + +```python + return_dict = kwargs.get("return_dict", True) + + def postprocess(h): + """Apply Qwen-specific output postprocessing.""" + h = module.norm_out(h, temb) + output = module.proj_out(h) + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) +``` + +### Step 5: Return CacheContext + +Package all information into a `CacheContext` object. + +```python + return CacheContext( + modulated_input=img_modulated, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, # or None for single-stream + temb=temb, + run_transformer_blocks=run_transformer_blocks, + postprocess=postprocess, + ) +``` + +**CacheContext Fields:** + +| Field | Type | Purpose | +|-------|------|---------| +| `modulated_input` | `torch.Tensor` | Tensor used for cache decision (similarity comparison) | +| `hidden_states` | `torch.Tensor` | Current hidden states (will be modified by caching) | +| `encoder_hidden_states` | `torch.Tensor | None` | Encoder states for dual-stream models, `None` for single-stream | +| `temb` | `torch.Tensor` | Timestep embedding tensor | +| `run_transformer_blocks` | `Callable[[], tuple]` | Executes transformer blocks, returns `(hidden_states, [encoder_hidden_states])` | +| `postprocess` | `Callable[[torch.Tensor], Any]` | Applies final transformations to produce model output | +| `extra_states` | `dict | None` | Optional dict for additional model-specific state | + +### Step 6: Register the Extractor + +Add your extractor to the `EXTRACTOR_REGISTRY` in `vllm_omni/diffusion/cache/teacache/extractors.py`: + +```python +EXTRACTOR_REGISTRY: dict[str, Callable] = { + "QwenImageTransformer2DModel": extract_qwen_context, + "Bagel": extract_bagel_context, + "ZImageTransformer2DModel": extract_zimage_context, + "YourModelTransformer2DModel": extract_your_model_context, # Add here +} +``` + +**Key:** Use the transformer class name (`module.__class__.__name__`) + +### Step 7: Add Model Coefficients + +Add polynomial rescaling coefficients to `vllm_omni/diffusion/cache/teacache/config.py`: + +```python +_MODEL_COEFFICIENTS = { + "QwenImageTransformer2DModel": [ + -4.50000000e02, + 2.80000000e02, + -4.50000000e01, + 3.20000000e00, + -2.00000000e-02, + ], + "YourModelTransformer2DModel": [ # Add your model's coefficients + # 5 polynomial coefficients (can reuse similar model's coefficients initially) + ], +} +``` + + +**Initial approach:** Start with coefficients from a similar model architecture, then tune empirically following [Customization](#customization) section. + +--- + +## Customization + +### Coefficient Estimation + +While you can start with coefficients from a similar model architecture, estimating custom coefficients for your specific model typically improves TeaCache performance. + +**Why Estimate Coefficients?** + +The polynomial coefficients rescale L1 distances between consecutive modulated inputs to better predict when cached residuals can be reused. Model-specific coefficients account for: + +- Architecture differences (layer count, hidden size, attention patterns) +- Training data characteristics +- Noise prediction behavior across timesteps + +| Approach | Performance | Effort | +|----------|-------------|--------| +| Using defaults from similar model | Within 5-10% of optimal | Low | +| Estimating custom coefficients | Best performance | Medium | + +#### Implement Data Collection Adapter + +Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`: + +```python +class YourModelAdapter: + """Adapter for coefficient estimation on your model.""" + + @staticmethod + def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any: + """Load your diffusion pipeline.""" + from your_model_package import YourModelPipeline + + pipeline = YourModelPipeline.from_pretrained( + model_path, + torch_dtype=dtype, + ) + pipeline = pipeline.to(device) + return pipeline + + @staticmethod + def get_transformer(pipeline: Any) -> tuple[Any, str]: + """Extract transformer from pipeline.""" + return pipeline.transformer, "YourTransformer2DModel" + + @staticmethod + def install_hook(transformer: Any, hook: DataCollectionHook) -> None: + """Install data collection hook on transformer.""" + from vllm_omni.diffusion.hooks import HookRegistry + + registry = HookRegistry.get_or_create(transformer) + registry.register_hook(hook._HOOK_NAME, hook) + + +# Register your adapter +_MODEL_ADAPTERS["YourModel"] = YourModelAdapter +``` + +#### Collect Data and Estimate + +```python +from vllm_omni.diffusion.cache.teacache.coefficient_estimator import ( + TeaCacheCoefficientEstimator, +) +from datasets import load_dataset +from tqdm import tqdm + +# Initialize estimator +estimator = TeaCacheCoefficientEstimator( + model_path="/path/to/your/model", + model_type="YourModel", +) + +# Load diverse prompts (paper recommends ~70 prompts) +dataset = load_dataset("nateraw/parti-prompts", split="train") +prompts = dataset["Prompt"][:70] + +# Collect data +for prompt in tqdm(prompts, desc="Collecting data"): + estimator.collect_from_prompt(prompt=prompt, num_inference_steps=50) + +# Estimate coefficients +coeffs = estimator.estimate(poly_order=4) +print(f"Estimated coefficients: {coeffs.tolist()}") +``` + +**Data Statistics Guide:** + +| Metric | Good Range | Warning Signs | +|--------|------------|---------------| +| **Count** | 2000-5000+ | < 1000: too few prompts | +| **Input Diffs (x)** | 0.01-0.10 | Very small (<0.001): model may not modulate properly | +| **Output Diffs (y)** | Should correlate with x | No correlation: check extractor | +| **Coefficient magnitude** | -1e6 to 1e6 | > 1e8: numerical instability | + +--- + +## Testing + +After adding TeaCache support, test with: + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni( + model="your-model-name", + cache_backend="tea_cache", + cache_config={ + "rel_l1_thresh": 0.2, + "coefficients": [1.33e6, -1.69e5, 7.95e3, -1.64e2, 1.26], # Your coefficients + } +) + +images = omni.generate( + "a beautiful landscape", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +**Verify:** + +1. **Check logs** - Look for TeaCache initialization messages +2. **Compare performance** - Measure speedup vs baseline (expect 1.5x-2.0x) +3. **Verify output quality** - Visually compare cached vs uncached outputs (should be nearly identical) + +See more detailed examples in [user guide for teacache](../../user_guide/diffusion/teacache.md). + +--- + +## Troubleshooting + +### Issue: "Unknown model type" + +**Symptoms:** Error message indicating the model type is not recognized when enabling TeaCache. + +**Causes & Solutions:** + +- **Extractor not registered:** + +**Problem:** The transformer class name doesn't exist in `EXTRACTOR_REGISTRY`. + +**Solution:** Check the class name and add to registry: +```python +# Check transformer class name +print(pipeline.transformer.__class__.__name__) + +# Add to EXTRACTOR_REGISTRY +EXTRACTOR_REGISTRY["YourTransformer2DModel"] = extract_your_context +``` + +- **Transformer class name mismatch:** + +**Solution:** Ensure the registry key matches exactly with `module.__class__.__name__`. + +### Issue: "Cannot find coefficients" + +**Symptoms:** Error when initializing TeaCache about missing model coefficients. + +**Causes & Solutions:** + +- **Missing coefficients in config:** + +**Solution:** Add coefficients to `_MODEL_COEFFICIENTS` in `config.py`, or pass custom coefficients: +```python +omni = Omni( + model="your-model", + cache_backend="tea_cache", + cache_config={"coefficients": [1.0, -0.5, 0.1, -0.01, 0.001]} +) +``` + +### Issue: Quality Degradation + +**Symptoms:** Output images look noticeably different or have artifacts compared to baseline. + +**Causes & Solutions:** + +- **Threshold too high:** + +**Problem:** `rel_l1_thresh` is too aggressive, causing cache reuse when outputs differ significantly. + +**Solution:** Lower the threshold: +```python +cache_config={"rel_l1_thresh": 0.1} # Try 0.1-0.2 +``` + +- **Coefficients not tuned:** + +**Solution:** Estimate model-specific coefficients using the coefficient estimation process described above. + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **Qwen-Image** | `vllm_omni/diffusion/cache/teacache/extractors.py` | Dual-stream | `extract_qwen_context` | +| **Bagel** | `vllm_omni/diffusion/cache/teacache/extractors.py` | Omni model | `extract_bagel_context` | +| **TeaCache Core** | `vllm_omni/diffusion/cache/teacache/` | Base implementation | Hook and config | +| **Coefficient Estimator** | `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py` | Estimation tool | Adapter pattern | + +--- + +## Summary + +Adding TeaCache support: + +1. ✅ **Write extractor** - Create function returning `CacheContext` with model-specific preprocessing +2. ✅ **Register extractor** - Add to `EXTRACTOR_REGISTRY` with transformer class name +3. ✅ **Add coefficients** - Add polynomial coefficients to `_MODEL_COEFFICIENTS` +4. ✅ **Test** - Verify with `cache_backend="tea_cache"` diff --git a/docs/contributing/features/tensor_parallel.md b/docs/contributing/features/tensor_parallel.md new file mode 100644 index 000000000..aecc1aedc --- /dev/null +++ b/docs/contributing/features/tensor_parallel.md @@ -0,0 +1,279 @@ +# Support Tensor Parallel + +This section describes how to add Tensor Parallel (TP) to a diffusion transformer model. We use the Z-Image transformer as the reference implementation. + +--- + +## Table of Contents + +- [Overview](#overview) +- [Step-by-Step Implementation](#step-by-step-implementation) +- [Testing](#testing) +- [Troubleshooting](#troubleshooting) +- [Reference Implementations](#reference-implementations) +- [Summary](#summary) + +--- + +## Overview + +### What is Tensor Parallel? + +Tensor Parallel (TP) is a model parallelism technique that **shards model weights** across multiple GPUs. Each GPU holds only a portion of the model's parameters and computes only part of each layer's output. + +Diffusion transformers contain large attention and MLP layers. We can use Tensor Parallel to shard the model dimension across multiple GPUs, allowing larger models to fit in memory while achieving near-linear speedup. + +### Architecture + +The Tensor Parallel implementation relies vLLM's Parallel Layers: + +[vLLM Parallel Layers API Reference](https://docs.vllm.ai/en/latest/contributing/model/basic/?h=column#3-optional-implement-tensor-parallelism-and-quantization-support) + +**Parallel Layer Types:** + +| Layer Type | Purpose | Weight Partitioning | +|------------|---------|---------------------| +| `ColumnParallelLinear` | First FFN layer, separated QKV | Columns (output dimension) | +| `RowParallelLinear` | Second FFN layer, attention output | Rows (input dimension) | +| `QKVParallelLinear` | Multi-head/grouped-query attention QKV | Handles head replication automatically | +| `ReplicatedLinear` | Layers that shouldn't be sharded | No partitioning (replicated) | + +--- + +## Step-by-Step Implementation + + +### Step 1: Identify Linear Layers + +Find all `nn.Linear` layers in your transformer that need to be sharded. + +**Key questions:** +- Which layers should be column parallel (weight split by columns)? +- Which layers should be row parallel (weight split by rows)? + +### Step 2: Replace Linear Layers with Parallel Equivalents + +Replace `nn.Linear` with parallel layers from `vllm.model_executor.layers.linear`. + +**Example (MLP Block - Up-Down Pattern):** + +```python +class FeedForward(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + super().__init__() + # Column parallel: weight split by columns [hidden_dim/N, dim] + self.w1 = ColumnParallelLinear( + dim, + hidden_dim, + bias=False, + return_bias=False, + ) + self.act = nn.GELU() + + self.w2 = RowParallelLinear( + hidden_dim, + dim, + bias=False, + input_is_parallel=True, # Input already sharded from w1 + return_bias=False, + ) + + def forward(self, x): + # x: [batch, seq, dim] (replicated on all GPUs) + # w1 outputs sharded [batch, seq, hidden_dim/N] + x = self.w1(x) + # act operates on sharded tensors (no communication) + x = self.act(x) + # w2 outputs full dim [batch, seq, dim] via all-reduce + x = self.w2(x) + return x +``` + +**Example (Attention - QKV-Out Pattern):** + +```python +from vllm_omni.diffusion.attention.layer import Attention +class YourModelAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, num_kv_heads: int): + super().__init__() + self.head_dim = dim // num_heads + + # Column parallel: QKV weight split by columns + # Each GPU gets num_heads/N heads + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_heads, + total_num_kv_heads=num_kv_heads, + bias=False, + return_bias=False, + ) + + # Row parallel: output weight split by rows + self.to_out = RowParallelLinear( + dim, + dim, + bias=False, + input_is_parallel=True, # Input sharded from attention + return_bias=False, + ) + + self.attn = Attention( + num_heads=self.to_qkv.num_heads, # Each GPU gets num_heads/N heads + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + num_kv_heads=self.to_qkv.num_kv_heads, + ) + + def forward(self, x): + # x: [batch, seq, dim] (replicated) + # to_qkv outputs sharded [batch, seq, (q+k+v) * head_dim/N] + qkv = self.to_qkv(x) + # Split into Q, K, V (each sharded on heads) + q, k, v = qkv.split([...], dim=-1) + # Attention computed independently on each GPU + out = self.attn(q, k, v) + # to_out all-reduces to full dim + out = self.to_out(out) + return out +``` + +**Key Points:** + +- `ColumnParallelLinear` → `RowParallelLinear` is the standard pairing +- Set `input_is_parallel=True` on `RowParallelLinear` when input comes from `ColumnParallelLinear` +- Use `QKVParallelLinear` for attention projections (handles head replication automatically) + +### Step 3: Validate TP Constraints + +For correct TP operation, these dimensions **must be divisible** by `tensor_parallel_size`: + +| Dimension | Reason | Example Error | +|-----------|--------|---------------| +| `num_heads` | Heads sharded by QKVParallelLinear | `num_heads=30, tp=4` ❌ (30 % 4 ≠ 0) | +| `num_kv_heads` | KV heads sharded by QKVParallelLinear | `num_kv_heads=30, tp=4` ❌ (30 % 4 ≠ 0) | + +--- + +## Testing + +After adding Tensor Parallel support, test with: + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +parallel_config = DiffusionParallelConfig(tensor_parallel_size=2) +omni = Omni(model="your-model-name", parallel_config=parallel_config) + +output = omni.generate( + "a cup of coffee on the table", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +**Or via command line:** + +```bash +cd examples/offline_inference/text_to_image +python text_to_image.py \ + --model Your-org/your-model \ + --prompt "a cup of coffee on the table" \ + --negative_prompt "ugly, unclear" \ + --cfg_scale 4.0 \ + --num_inference_steps 50 \ + --output "tp_enabled.png" \ + --tensor_parallel_size 2 +``` + +**Verify:** + +1. Check the `e2e_time_ms` in the log for speedup +2. Compare generated image quality with TP disabled +3. Verify memory usage is reduced proportionally +4. Record comparison results in your PR + +--- + +## Troubleshooting + +### Issue: TP not activating + +**Symptoms:** Model runs on single GPU, no memory savings or speedup. + +**Causes & Solutions:** + +- **Still using `nn.Linear`:** + +**Problem:** Linear layers not replaced with parallel equivalents. + +**Solution:** Replace with parallel layers: +```python +# ❌ BAD +self.proj = nn.Linear(dim, dim) + +# ✅ GOOD +self.proj = RowParallelLinear(dim, dim, input_is_parallel=True) +``` + +### Issue: Dimension mismatch errors + +**Symptoms:** `RuntimeError: shape mismatch` during forward pass. + +**Causes & Solutions:** + +- **Missing `input_is_parallel=True`:** + +**Problem:** RowParallelLinear expects sharded input but receives full tensor. + +**Solution:** Set `input_is_parallel=True` when input comes from ColumnParallelLinear: +```python +# ✅ GOOD: Correct pairing +self.w1 = ColumnParallelLinear(dim, hidden_dim, return_bias=False,) +self.w2 = RowParallelLinear( + hidden_dim, + dim, + input_is_parallel=True, # Input sharded from w1 + return_bias=False, +) +``` + +- **Incorrect split dimensions:** + +**Problem:** QKV split sizes don't match sharded dimensions. + +**Solution:** Use `self.to_qkv.num_heads` (local heads per GPU): +```python +# ❌ BAD: Uses total heads +q_size = self.total_num_heads * self.head_dim + +# ✅ GOOD: Uses local heads +q_size = self.to_qkv.num_heads * self.head_dim +``` + +--- + +## Reference Implementations + +Complete examples in the codebase: + +| Model | Path | Pattern | Notes | +|-------|------|---------|-------| +| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation | +| **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams | +| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE | +| **TP Tests** | `tests/e2e/offline_inference/test_zimage_tensor_parallel.py` | E2E testing | TP correctness and performance | +| **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic | + +--- + +## Summary + +Adding Tensor Parallel support to a transformer: + +1. ✅ **Identify linear layers** - Which layers should be sharded? +2. ✅ **Replace with parallel layers** - Use QKVParallelLinear, ColumnParallelLinear, RowParallelLinear +3. ✅ **Validate TP constraints** - Ensure dimensions divisible by TP size +4. ✅ **Test** - Verify with `tensor_parallel_size=N`, check memory, speed, and quality diff --git a/docs/contributing/metrics.md b/docs/contributing/metrics.md new file mode 100644 index 000000000..7a46a2861 --- /dev/null +++ b/docs/contributing/metrics.md @@ -0,0 +1,165 @@ + +# Metrics vLLM-Omni: + +You can use these metrics in production to monitor the health and performance of the vLLM-omni system. Typical scenarios include: +- **Performance Monitoring**: Track throughput (e.g., `e2e_avg_tokens_per_s`), latency (e.g., `e2e_total_ms`), and resource utilization to verify that the system meets expected standards. +- **Debugging and Troubleshooting**: Use detailed per-request metrics to diagnose issues, such as high transfer times or unexpected token counts. + +## How to Enable and View Metrics + +### 1. Start the Service with Metrics Logging + +```bash +vllm serve /workspace/models/Qwen3-Omni-30B-A3B-Instruct --omni --port 8014 --log-stats +``` + +### 2. Send a Request + +```bash +python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image +``` + +### 3. What You Will See + +With `--log-stats` enabled, the server will output detailed metrics logs after each request. Example output: + + +#### Overall Summary + +| Field | Value | +|-----------------------------|--------------| +| e2e_requests | 1 | +| e2e_wall_time_ms | 41,299.190 | +| e2e_total_tokens | 5,202 | +| e2e_avg_time_per_request_ms | 41,299.190 | +| e2e_avg_tokens_per_s | 125.959 | +| e2e_stage_0_wall_time_ms | 10,192.289 | +| e2e_stage_1_wall_time_ms | 30,541.409 | +| e2e_stage_2_wall_time_ms | 207.496 | + +#### RequestE2EStats + +| Field | Value | +|-------------------------|------------| +| e2e_total_ms | 41,299.133 | +| e2e_total_tokens | 5,202 | +| transfers_total_time_ms | 245.895 | +| transfers_total_kbytes | 138,089.939| + +#### StageRequestStats + +| Field | 0 | 1 | 2 | +|------------------------|--------|--------|--------| +| audio_generated_frames | 0 | 0 | 525,525| +| batch_id | 38 | 274 | 0 | +| batch_size | 1 | 1 | 1 | +| num_tokens_in | 4,860 | 4,826 | 4,384 | +| num_tokens_out | 67 | 275 | 0 | +| postprocess_time_ms | 256.158| 0.491 | 0.000 | +| stage_gen_time_ms | 9,910.007|30,379.198|160.745| + +#### TransferEdgeStats + +| Field | 0->1 | 1->2 | +|---------------------|-------------|------------| +| size_kbytes | 109,277.349 | 28,812.591 | +| tx_time_ms | 78.701 | 18.790 | +| rx_decode_time_ms | 111.865 | 31.706 | +| in_flight_time_ms | 2.015 | 2.819 | + + +These logs include: +- **Overall summary**: total requests, wall time, average tokens/sec, etc. +- **E2E table**: per-request latency and token counts. +- **Stage table**: per-stage batch and timing details. +- **Transfer table**: data transfer and timing for each edge. + +You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above. + + +## Metrics Scope: Offline vs Online Inference + +For **offline inference** (batch mode), the summary includes both system-level metrics (aggregated across all requests) and per-request metrics. In this case, `e2e_requests` can be greater than 1, reflecting multiple completed requests in a batch. + +For **online inference** (serving mode), the summary is always per-request. `e2e_requests` is always 1, and only request-level metrics are reported for each completion. + +--- + +## Parameter Details + +| Field | Meaning | +|---------------------------|----------------------------------------------------------------------------------------------| +| `e2e_requests` | Number of completed requests. | +| `e2e_wall_time_ms` | Wall-clock time span from run start to last completion, in ms. | +| `e2e_total_tokens` | Total tokens counted across all completed requests (stage0 input + all stage outputs). | +| `e2e_avg_time_per_request_ms` | Average wall time per request: `e2e_wall_time_ms / e2e_requests`. | +| `e2e_avg_tokens_per_s` | Average token throughput over wall time: `e2e_total_tokens * 1000 / e2e_wall_time_ms`. | +| `e2e_stage_{i}_wall_time_ms` | Wall-clock time span for stage i, in ms. Each stage's wall time is reported as a separate field, e.g., `e2e_stage_0_wall_time_ms`, `e2e_stage_1_wall_time_ms`, etc. | + +--- + +## E2E Table (per request) + +| Field | Meaning | +|---------------------------|-----------------------------------------------------------------------| +| `e2e_total_ms` | End-to-end latency in ms. | +| `e2e_total_tokens` | Total tokens for the request (stage0 input + all stage outputs). | +| `transfers_total_time_ms` | Sum of transfer edge `total_time_ms` for this request. | +| `transfers_total_kbytes` | Sum of transfer kbytes for this request. | + + +--- + +## Stage Table (per stage event / request) + +| Field | Meaning | +|---------------------------|-------------------------------------------------------------------------------------------------| +| `batch_id` | Batch index. | +| `batch_size` | Batch size. | +| `num_tokens_in` | Input tokens to the stage. | +| `num_tokens_out` | Output tokens from the stage. | +| `stage_gen_time_ms` | Stage compute time in ms, excluding postprocessing time (reported separately as `postprocess_time_ms`). | +| `image_num` | Number of images generated (for diffusion/image stages). | +| `resolution` | Image resolution (for diffusion/image stages). | +| `postprocess_time_ms` | Diffusion/image: post-processing time in ms. | + +--- + +## Transfer Table (per edge / request) + +| Field | Meaning | +|----------------------|---------------------------------------------------------------------------| +| `size_kbytes` | Total kbytes transferred. | +| `tx_time_ms` | Sender transfer time in ms. | +| `rx_decode_time_ms` | Receiver decode time in ms. | +| `in_flight_time_ms` | In-flight time in ms. | + + +## Expectation of the Numbers (Verification) + +**Formulas:** +- `e2e_total_tokens = Stage0's num_tokens_in + sum(all stages' num_tokens_out)` +- `transfers_total_time_ms = sum(tx_time_ms + rx_decode_time_ms + in_flight_time_ms)` for every edge + +**Using the example above:** + +### e2e_total_tokens +- Stage0's `num_tokens_in`: **4,860** +- Stage0's `num_tokens_out`: **67** +- Stage1's `num_tokens_out`: **275** +- Stage2's `num_tokens_out`: **0** + +So, +``` +e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202 +``` +This matches the table value: `e2e_total_tokens = 5,202`. + +### transfers_total_time_ms +For each edge: +- 0->1: tx_time_ms (**78.701**) + rx_decode_time_ms (**111.865**) + in_flight_time_ms (**2.015**) = **192.581** +- 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315** + +Sum: 192.581 + 53.315 = **245.896** + +The table shows `transfers_total_time_ms = 245.895`, which matches the calculation (difference is due to rounding). diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 80212360f..529cab752 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -1,22 +1,27 @@ -# Adding a Diffusion Model -This guide walks through the process of adding a new Diffusion model to vLLM-Omni, using Qwen/Qwen-Image-Edit as a comprehensive example. +# Adding a Diffusion Model to vLLM-Omni + +This guide walks you through adding a new diffusion model to vLLM-Omni. We use **Qwen-Image** as the primary example, with references to other models (LongCat, Flux, Wan2.2) to illustrate different patterns. + + +--- + +## Table of Contents -# Table of Contents 1. [Overview](#overview) 2. [Directory Structure](#directory-structure) -3. [Step-by-Step Implementation](#step-by-step-implementation) -4. [Testing](#testing) -5. [Adding a Model Recipe](#adding-a-model-recipe) +3. [Basic Implementation](#basic-implementation) +4. [Advanced Features](#advanced-features) +5. [Troubleshooting](#troubleshooting) +6. [Pull Request Checklist](#pull-request-checklist) +7. [Reference Implementations](#reference-implementations) +8. [Summary](#summary) +--- -# Overview -When add a new diffusion model into vLLM-Omni, additional adaptation work is required due to the following reasons: +## Overview -+ New model must follow the framework’s parameter passing mechanisms and inference flow. +vLLM-Omni's diffusion inference follows this architecture: -+ Replacing the model’s default implementations with optimized modules, which is necessary to achieve the better performance. - -The diffusion execution flow as follow:

@@ -24,348 +29,892 @@ The diffusion execution flow as follow:

+**Key Components:** -# Directory Structure -File Structure for Adding a New Diffusion Model +1. **Request Handling:** User prompts → `OmniDiffusionRequest` +2. **Diffusion Engine:** Request → Preprocessing (Optional) → Pipeline execution -> Post-processing +3. **Pipeline Execution:** Request → Encode prompt → Diffusion steps → Vae decode + + +## Directory Structure + +Organize your model files following this structure: ``` vllm_omni/ -└── examples/ - └──offline_inference - └── example script # reuse existing if possible (e.g., image_edit.py) - └──online_serving - └── example script └── diffusion/ - └── registry.py # Registry work - ├── request.py # Request Info - └── models/your_model_name/ # Model directory (e.g., qwen_image) - └── pipeline_xxx.py # Model implementation (e.g., pipeline_qwen_image_edit.py) + ├── registry.py # ← Register your model here + ├── request.py # Request data structures + └── models/ + └── your_model_name/ # ← Create this directory + ├── __init__.py # Export pipeline and transformer + ├── pipeline_xxx.py # Pipeline implementation + └── xxx_transformer.py # Transformer implementation ``` -# Step-by-step-implementation -## Step 1: Model Implementation -The diffusion pipeline’s implementation follows **HuggingFace Diffusers**. -### 1.1 Define the Pipeline Class -Define the pipeline class, e.g., `QwenImageEditPipeline`, and initialize all required submodules, either from HuggingFace `diffusers` or custom implementations. In `QwenImageEditPipeline`, only `QwenImageTransformer2DModel` is re-implemented to support optimizations such as Ulysses-SP. When adding new models in the future, you can either reuse this re-implemented `QwenImageTransformer2DModel` or extend it as needed. +**Naming Conventions:** -### 1.2 Pre-Processing and Post-Processing Extraction -Extract the pre-processing and post-processing logic from the pipeline class to follow vLLM-Omni’s execution flow. For Qwen-Image-Edit: +- **Model directory:** `your_model_name` (lowercase, underscores), e.g., `qwen_image`, `flux`, `longcat_image`, `wan2_2` +- **Pipeline file:** `pipeline_xxx.py` where `xxx` describes the task, e.g., `pipeline_qwen_image.py`, `pipeline_qwen_image_edit.py` +- **Transformer file:** `xxx_transformer.py` matching transformer class name, e.g., `qwen_image_transformer.py`, `flux_transformer.py` + +--- + +## Basic Implementation + +This section covers the minimal steps to get a model working in vLLM-Omni with basic features (online/offline serving, batch requests). + +### Step 1: Adapt Transformer Model + +The transformer is the core denoising network. Start by copying the transformer implementation from Diffusers and making these adaptations. + + +#### 1.1: Remove Diffusers Mixins + +Diffusers' `Mixin` classes are not needed in vLLM-Omni. Remove them: + +```diff +# Before (Diffusers) +- from diffusers.models.modeling_utils import ModelMixin +- from diffusers.models.attention_processor import AttentionModuleMixin + +- class YourModelTransformer2DModel(ModelMixin, AttentionModuleMixin): ++ class YourModelTransformer2DModel(nn.Module): + """Your transformer model.""" +``` + +**Example mixins to remove:** + +- `ModelMixin` - Weight loading utilities (vLLM-Omni has its own weight loader) +- `AttentionModuleMixin` - Attention processors (using vLLM-Omni's Attention layer instead) +- `ConfigMixin` - Config management (not needed) +- `PeftAdapterMixin` - Parameter efficient finetune utilities (not needed) + +#### 1.2: Replace Attention Implementation + +**The most important adaptation:** Replace Diffusers' attention with vLLM-Omni's optimized `Attention` layer. + +**Before (Diffusers):** ```python -def get_qwen_image_edit_pre_process_func( - od_config: OmniDiffusionConfig, -): - """ - Define a pre-processing function that resizes input images and - pre-process for subsequent inference. - """ +from diffusers.models.attention_processor import dispatch_attention_fn + +class YourAttentionBlock(nn.Module): + def forward(self, hidden_states, encoder_hidden_states=None, ...): + ... + hidden_states = dispatch_attention_fn( + query, key, value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + ) ``` -```python -def get_qwen_image_edit_post_process_func( - od_config: OmniDiffusionConfig, -): - """ - Defines a post-processing function that post-process images. - """ -``` - -### 1.3 Define the forward function -The forward function of `QwenImageEditPipeline` follows the HuggingFace `diffusers` design for the most part. The key differences are: -+ As described in the overview, arguments are passed through `OnniDiffusionRequest`, so we need to get user parameters from it accordingly. -```python -prompt = req.prompt -``` -+ pre/post-processing are handled by the framework elsewhere, so skip them. - -### 1.4 Replace some ops or layers in DiT component - -vLLM-Omni provides a set of optimized operators with better performance and built-in support for parallelism, including attention, rotary embeddings (RoPE), and linear layers. - -Below is an example showing how to replace standard Transformer attention and FFN layers with vLLM-Omni implementations: - +**After (vLLM-Omni):** ```python from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + +class YourAttentionBlock(nn.Module): + def __init__(self, ...): + super().__init__() + + # Initialize vLLM-Omni's Attention layer + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim ** 0.5), + causal=False, # Diffusion models typically use bidirectional attention + num_kv_heads=self.num_kv_heads, + ) + + def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, ...): + ... + # Create attention metadata + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata) + +``` + +**Key Points:** + +- **Attention layer initialization:** Done in `__init__`, not per-forward +- **Tensor shapes:** vLLM-Omni `Attention` expects QKV to have `[B, seq, num_heads, head_dim]` shape +- **AttentionMetadata:** Wraps attention mask and other metadata + +**Attention backends:** vLLM-Omni automatically selects the attention backend given the environmental variable `DIFFUSION_ATTENTION_BACKEND`. The default attention backend is `FLASH_ATTN` for diffusion models. + +#### 1.3: Replace Imports and Utilities + +**Logger:** +```diff +- from diffusers.utils import logging +- logger = logging.get_logger(__name__) + ++ from vllm.logger import init_logger ++ logger = init_logger(__name__) +``` + +**Custom layers from vLLM and vLLM-Omni (if needed):** + +```python +from vllm.model_executor.layers.layernorm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm +``` -class MyAttention(nn.Module): - def __init__(self): +#### 1.4: Remove Training-Only Code + +Remove code that's only needed for training: + +```diff +# Remove gradient checkpointing +- if torch.is_grad_enabled() and self.gradient_checkpointing: +- hidden_states = torch.utils.checkpoint.checkpoint( +- self._forward_block, hidden_states, ... +- ) +- else: +- hidden_states = self._forward_block(hidden_states, ...) ++ hidden_states = self._forward_block(hidden_states, ...) + +# Remove training-specific attributes +- self.gradient_checkpointing = False + +# Remove dropout (set to 0 or remove) +- self.dropout = nn.Dropout(dropout_prob) ++ # Removed dropout for inference +``` + +#### 1.5: Add Configuration Support + +Add support for vLLM-Omni's `OmniDiffusionConfig`: + +```python +from vllm_omni.diffusion.data import OmniDiffusionConfig + +class YourModelTransformer2DModel(nn.Module): + def __init__( + self, + *, + od_config: OmniDiffusionConfig | None = None, # ← Add vLLM-Omni config + # ... other model-specific parameters + num_layers: int = 28, + hidden_size: int = 3072, + num_heads: int = 24, + **kwargs, + ): super().__init__() - self.attn = Attention() - self.to_qkv = QKVParallelLinear() - self.to_out = RowParallelLinear() - self.rope = RotaryEmbedding(is_neox_style=False) - def forward(self, hidden_states): - qkv, _ = self.to_qkv(hidden_states) - q, k, v = qkv.split(...) - q, k = self.rope(...) - attn_output = self.attn(q, k, v) - output = self.to_out(attn_output) + # Store config + self.od_config = od_config + self.parallel_config = od_config.parallel_config if od_config else None -class MyFFN(nn.Module): - def __init__(self): + # Model architecture + self.num_layers = num_layers + self.hidden_size = hidden_size + # ... initialize layers +``` + +### Step 2: Adapt Pipeline + +The pipeline orchestrates the full generation process (text encoding, denoising loop, VAE decoding). Adapt it from Diffusers format to vLLM-Omni's interface. + +#### 2.1: Remove Diffusers Inheritance + +**Remove Diffusers base classes:** +```diff +- from diffusers import DiffusionPipeline +- from diffusers.loaders import LoraLoaderMixin + +- class YourModelPipeline(DiffusionPipeline, LoraLoaderMixin): ++ class YourModelPipeline(nn.Module): + """Your model pipeline for vLLM-Omni.""" +``` + +#### 2.2: Adapt `__init__` Method + +**Before (Diffusers):** +```python +class YourModelPipeline(DiffusionPipeline): + def __init__( + self, + vae: AutoencoderKL, + text_encoder: CLIPTextModel, + tokenizer: CLIPTokenizer, + transformer: YourTransformer, + scheduler: FlowMatchScheduler, + ): super().__init__() - self.fc1 = ColumnParallelLinear() - self.fc2 = RowParallelLinear() - self.act = F.gelu - - def forward(self, hidden_states): - hidden, _ = self.fc1(hidden_states) - hidden = self.act(hidden) - output = self.fc2(hidden) - return output + self.register_modules( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + scheduler=scheduler, + ) ``` -In this example: - -+ Attention uses vLLM-Omni’s optimized attention kernel together with parallel QKV projection and RoPE. - -+ Linear layers are replaced with column- and row-parallel variants to enable tensor parallelism. - -+ The FFN follows a standard two-layer structure and can be further optimized (e.g., using fused or merged projections) if needed. - - -### 1.5 Provide a `_repeated_blocks` in DiT model -`_repeated_blocks` is the small and frequently-repeated block(s) of a model -- typically a transformer layer. - -It's used for torch compile optimizations. +**After (vLLM-Omni):** ```python -_repeated_blocks = ["QwenImageTransformerBlock"] +import os +from diffusers import AutoencoderKL +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from transformers import CLIPTextModel, CLIPTokenizer + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.models.your_model_name.your_model_transformer import ( + YourModelTransformer2DModel, +) + + +class YourModelPipeline(nn.Module): + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.parallel_config = od_config.parallel_config + self.device = get_local_device() + model = od_config.model + local_files_only = os.path.exists(model) + + # Load components from checkpoint + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only) + self.text_encoder = CLIPTextModel.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only).to(self.device) + self.tokenizer = CLIPTokenizer.from_pretrained( + model, subfolder="tokenizer", local_files_only=local_files_only) + self.vae = AutoencoderKL.from_pretrained( + model, subfolder="vae", local_files_only=local_files_only).to(self.device) + + # Initialize transformer with vLLM-Omni config + transformer_kwargs = get_transformer_config_kwargs( + od_config.tf_model_config, YourModelTransformer2DModel) + self.transformer = YourModelTransformer2DModel( + od_config=od_config, **transformer_kwargs) + + # Store VAE scale factor for latent space conversions + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.default_sample_size = 128 # Default latent size ``` +**Key Changes:** -### 1.6 (Optional) implement sequence parallelism -vLLM-Omni has a non-intrusive `_sp_plan` that enable sequence parallel without modifying `forward()` logic. -You can refer to [How to parallelize a new model](../../user_guide/diffusion/parallelism_acceleration.md) +1. **`od_config` parameter:** All configuration through `OmniDiffusionConfig` +2. **Manual component loading:** No `register_modules()`, load each component explicitly +3. **Local files support:** Check `os.path.exists(model)` for local checkpoints +4. **Transformer with config:** Pass `od_config` to transformer constructor +#### 2.3: Adapt `__call__` → `forward` Method -### 1.7 (Optional) integrate with Cache-Dit -vLLM-Omni supports acceleration via [Cache-Dit](../../user_guide/diffusion/cache_dit_acceleration.md). Most models compatible with Diffusers can use Cache-Dit seamlessly. For new models, you can extend support by modifying`cache_dit_backend.py` +**Change signature:** -## Step 2: Extend OmniDiffusionRequest Fields -User-provided inputs are ultimately passed to the model’s forward method through OmniDiffusionRequest, so we add the required fields here to support the new model. -```python -prompt: str | list[str] | None = None -negative_prompt: str | list[str] | None = None -... +```diff +- @torch.no_grad() +- def __call__( ++ def forward( + self, ++ req: OmniDiffusionRequest, # ← Add request parameter here +- ): ++ ) -> DiffusionOutput: # ← Add return type ``` -## Step 3: Registry -+ registry diffusion model in registry.py +[`OmniDiffusionRequest`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/request/#vllm_omni.diffusion.request.OmniDiffusionRequest) is a dataclass that contains the **prompts** and **sampling parameters** [`OmniDiffusionSamplingParams`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/inputs/data/#vllm_omni.inputs.data.OmniDiffusionSamplingParams) for the diffusion pipeline execution. It also contains a request_id for other components to trace this request and its outputs. + +See some parameters in `OmniDiffusionSamplingParams` as follows: + +| parameters | type |value | function | +|:---:|:---:|:---:|:---:| +| `num_inference_steps` | `int` | 50 | The number of diffusion steps during inference| +| `guidance_scale` | `float` | 0.0 | The classifier free guidance scale | +| `width` and `height` | `int` | None | The width and height of the generated image | + +**Extract parameters from request:** + ```python +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.data import DiffusionOutput + +def forward( + self, + req: OmniDiffusionRequest, +) -> DiffusionOutput: + # Extract prompts from request + if req.prompts is not None: + prompt = [ + p if isinstance(p, str) else (p.get("prompt") or "") + for p in req.prompts + ] + + # Extract sampling parameters + sampling_params = req.sampling_params + num_inference_steps = sampling_params.num_inference_steps or 50 + guidance_scale = sampling_params.guidance_scale or 7.5 + height = sampling_params.height or (self.default_sample_size * self.vae_scale_factor) + width = sampling_params.width or (self.default_sample_size * self.vae_scale_factor) + + # For image editing pipelines, extract images from multi_modal_data + if hasattr(req, 'multi_modal_data') and req.multi_modal_data: + input_images = req.multi_modal_data.get('image', []) + + # ... rest of generation logic +``` + +For an image editing model, an example `OmniDiffusionRequest` is like: +```python +{ + "prompt": "turn this cat to a dog", + "multi_modal_data": {"image": input_image} +}, +``` + +**Wrap output:** + +```diff + # Generate images + images = self.vae.decode(latents)[0] + +- return {"images": images} ++ return DiffusionOutput(output=images) +``` + +#### 2.4: Extract Pre/Post-Processing Functions + +vLLM-Omni separates image processing from the main pipeline for better modularity. + +**Post-processing function (required):** +```python +def get_your_model_post_process_func( + od_config: OmniDiffusionConfig, +): + """ + Create post-processing function for your model. + + Returns a function that converts latents to images. + """ + from diffusers.image_processor import VaeImageProcessor + import json + + # Load VAE config to get scale factor + model_path = od_config.model + if not os.path.exists(model_path): + from vllm_omni.diffusion.model_loader.utils import download_weights_from_hf_specific + model_path = download_weights_from_hf_specific(model_path, 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) + + # Create image processor + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) + + def post_process_func(images: torch.Tensor): + return image_processor.postprocess(images, output_type="pil") + + return post_process_func +``` + +**Pre-processing function (for image editing pipelines):** + +```python +def get_your_model_pre_process_func( + od_config: OmniDiffusionConfig, +): + """ + Create pre-processing function for image editing. + + Returns a function that prepares input images. + """ + from PIL import Image + from diffusers.image_processor import VaeImageProcessor + + # Load VAE config + # ... (similar to post_process_func) + + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) + + def pre_process_func( + request: OmniDiffusionRequest, + ): + for i, prompt in enumerate(request.prompts): + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + # image pre-processing + # after pre-processing, update the request attributes + ... + return request + + return pre_process_func +``` + +#### 2.5: Add Weight Loading Support + +Add methods for automatic weight downloading and loading: + +```python +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm.model_executor.models.utils import AutoWeightsLoader + +class YourModelPipeline(nn.Module): + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + # ... initialization code + + # Define weight sources for automatic loading + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """ + Customize the weight loading behavior, such as filter weights name. + + Args: + weights: Iterable of (param_name, param_tensor) tuples + + Returns: + Set of loaded parameter names + """ + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) +``` + +### Step 3: Register Model + +Register your model in `vllm_omni/diffusion/registry.py` so vLLM-Omni can discover and load it. + +#### 3.1: Register Pipeline Class + +```python +# vllm_omni/diffusion/registry.py + _DIFFUSION_MODELS = { - # arch:(mod_folder, mod_relname, cls_name) - ... - "QwenImageEditPipeline": ( - "qwen_image", - "pipeline_qwen_image_edit", - "QwenImageEditPipeline", + # Format: "PipelineClassName": (module_folder, module_file, class_name) + + # Existing models + "QwenImagePipeline": ("qwen_image", "pipeline_qwen_image", "QwenImagePipeline"), + "FluxPipeline": ("flux", "pipeline_flux", "FluxPipeline"), + + # Add your model + "YourModelPipeline": ( + "your_model_name", # Module folder name + "pipeline_your_model", # Python file name (without .py) + "YourModelPipeline", # Pipeline class name ), - ... } ``` -+ registry pre-process get function + +#### 3.2: Register Pre/Post-Processing Function + ```python +# vllm_omni/diffusion/registry.py _DIFFUSION_PRE_PROCESS_FUNCS = { # arch: pre_process_func - ... + # `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py, + # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key + "GlmImagePipeline": "get_glm_image_pre_process_func", "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func", - ... -} -``` -+ registry post-process get function -```python + # Add your model + "YourModelPipeline": "get_your_model_pre_process_func", # Optional +} _DIFFUSION_POST_PROCESS_FUNCS = { - # arch: post_process_func - ... - "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func", - ... + # Format: "PipelineClassName": "function_name" + + # Existing models + "QwenImagePipeline": "get_qwen_image_post_process_func", + "FluxPipeline": "get_flux_post_process_func", + + # Add your model + "YourModelPipeline": "get_your_model_post_process_func", } ``` -## Step 4: Add an Example Script -For each newly integrated model, we need to provide examples script under the examples/ to demonstrate how to initialize the pipeline with Omni, pass in user inputs, and generate outputs. -Key point for writing the example: -+ Use the Omni entrypoint to load the model and construct the pipeline. +#### 3.3: Export from Module -+ Show how to format user inputs and pass them via omni.generate(...). - -+ Demonstrate the common runtime arguments, such as: - - + model path or model name - - + input image(s) or prompt text - - + key diffusion parameters (e.g., inference steps, guidance scale) - - + optional acceleration backends (e.g., Cache-DiT, TeaCache) - -+ Save or display the generated results so users can validate the integration. - -## Step 5: TeaCache Coefficient Estimation (Optional) - -If your model supports TeaCache acceleration, you need to estimate the polynomial coefficients for optimal caching performance. - -### 5.1 Add Extractor Function - -First, implement an extractor function in `vllm_omni/diffusion/cache/teacache/extractors.py`. The extractor extracts the modulated input and defines how to run transformer blocks: +Create/update `__init__.py` to export your classes: ```python -def extract_your_model_context( - module: nn.Module, - hidden_states: torch.Tensor, - timestep: torch.Tensor, - **kwargs: Any, -) -> CacheContext: - # 1. Preprocessing - temb = module.time_embed(timestep) +# vllm_omni/diffusion/models/your_model_name/__init__.py - # 2. Extract modulated input (for cache decision) - modulated_input = module.transformer_blocks[0].norm1(hidden_states, temb) - - # 3. Define transformer execution - def run_transformer_blocks(): - h = hidden_states - for block in module.transformer_blocks: - h = block(h, temb=temb) - return (h,) - - # 4. Define postprocessing - def postprocess(h): - return module.proj_out(module.norm_out(h, temb)) - - return CacheContext( - modulated_input=modulated_input, - hidden_states=hidden_states, - encoder_hidden_states=None, - temb=temb, - run_transformer_blocks=run_transformer_blocks, - postprocess=postprocess, - ) -``` - -Register it in `EXTRACTOR_REGISTRY`: -```python -EXTRACTOR_REGISTRY = { - ... - "YourTransformer2DModel": extract_your_model_context, -} -``` - -### 5.2 Add Adapter for Coefficient Estimation - -Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`: - -```python -class YourModelAdapter: - @staticmethod - def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any: - # Load your pipeline - ... - - @staticmethod - def get_transformer(pipeline: Any) -> tuple[Any, str]: - return pipeline.transformer, "YourTransformer2DModel" - - @staticmethod - def install_hook(transformer: Any, hook: DataCollectionHook) -> None: - registry = HookRegistry.get_or_create(transformer) - registry.register_hook(hook._HOOK_NAME, hook) - -_MODEL_ADAPTERS["YourModel"] = YourModelAdapter -``` - -### 5.3 Run Coefficient Estimation - -Use the provided script to estimate coefficients: - -```python -from vllm_omni.diffusion.cache.teacache.coefficient_estimator import ( - TeaCacheCoefficientEstimator, +from .pipeline_your_model import ( + YourModelPipeline, + get_your_model_post_process_func, ) -from datasets import load_dataset -from tqdm import tqdm +from .your_model_transformer import YourModelTransformer2DModel -# Load model -estimator = TeaCacheCoefficientEstimator( - model_path="/path/to/model", - model_type="Bagel", # Your model type - device="cuda", -) - -# Load prompts (paper suggests ~70 prompts) -dataset = load_dataset("nateraw/parti-prompts", split="train") -prompts = dataset["Prompt"][:70] - -# Collect data -for prompt in tqdm(prompts): - estimator.collect_from_prompt(prompt, num_inference_steps=50) - -# Estimate coefficients -coeffs = estimator.estimate(poly_order=4) -print(f"Coefficients: {coeffs}") +__all__ = [ + "YourModelPipeline", + "YourModelTransformer2DModel", + "get_your_model_post_process_func", +] ``` -### 5.4 Interpreting Coefficient Estimation Results +--- -The estimator outputs statistics and polynomial coefficients. Here's how to interpret them: +### Step 4: Add Example Script -**Example Output:** -``` -Data statistics: -Count: 48 -Input Diffs (x): min=1.1089e-02, max=5.2555e-02, mean=2.8435e-02 -Output Diffs (y): min=2.8242e-02, max=2.9792e-01, mean=7.0312e-02 -Coefficients: [1333131.29, -168644.23, 7950.51, -163.75, 1.26] + +If your model is one of Text-to-Image, Text-to-Audio, Text-to-Video, Image-to-Image, Image-to-Video models, you can simply try one of the following offline inference scripts to run your model: + +| Model Category | Offline Inference Script | +|---|---| +| Image-to-Image | `examples/offline_inference/image_to_image/image_edit.py` | +| Image-to-Video | `examples/offline_inference/image_to_video/image_to_video.py` | +| Text-to-Image | `examples/offline_inference/text_to_image/text_to_image.py` | +| Text-to-Audio | `examples/offline_inference/text_to_audio/text_to_audio.py` | +| Text-to-Video | `examples/offline_inference/text_to_video/text_to_video.py` | + + +If new CLI arguments need to be added, please edit the offline inference script corresponding to your model category from the table above, and update the example inference script in its corresponding document file (e.g., `examples/offline_inference/text_to_video/text_to_video.md`). + +For online inference, all the supported tasks are listed in `docs/user_guide/examples/online_serving/`. If your model falls into these categories, please check the corresponding documentation in this folder and the example at `examples/online_serving/TASK_NAME`. Update them accordingly if needed. + +--- + +If your model is an Omni (understanding and generation) model, please follow the steps below. + +#### 4.1: Create Example File + +Taking **BAGEL** model as examples for both offline and online: + +- Offline: `examples/offline_inference/bagel/` +- Online: `examples/online_serving/bagel/` + +Add **two example folders** for your model: + +```bash +mkdir -p examples/offline_inference/your_model_name +mkdir -p examples/online_serving/your_model_name ``` -**What to Check:** -- **Count**: Number of timestep pairs analyzed. Should be at least 30-50 for reliable estimation. Low count suggests insufficient prompts or inference steps. -- **Input/Output Ranges**: Verify output differences correlate with input differences. If ranges seem unusual, check your prompt diversity. -- **Coefficient Magnitude**: Extremely large values (>1e8) may indicate numerical instability - try collecting more diverse data. +**Offline (recommended minimum):** create `examples/offline_inference/your_model_name/end2end.py` and a README. -**Troubleshooting:** -- If results seem unreliable, try: - - Increasing number of prompts (100+ recommended) - - Using more diverse prompts from multiple datasets - - Adjusting `num_inference_steps` (try 20, 50, 100) +- Script: `examples/offline_inference/your_model_name/end2end.py` + - Parse args like BAGEL (`--model`, `--modality`, optional `--image-path`, `--steps`, etc.) + - Use `from vllm_omni.entrypoints.omni import Omni` (or `OmniDiffusion` if your model is diffusion-only) + - Save outputs (images/audio/video/text) with deterministic filenames (e.g., `output_0_0.png`) +- Doc: `examples/offline_inference/your_model_name/README.md` + - Include at least one runnable command, e.g.: -### 5.5 Add Coefficients to Config - -Add the estimated coefficients to `vllm_omni/diffusion/cache/teacache/config.py`: - -```python -_MODEL_COEFFICIENTS = { - ... - "YourTransformer2DModel": [ - 1.04730573e+06, # a4 - -1.34150749e+05, # a3 - 6.51517806e+03, # a2 - -1.41209108e+02, # a1 - 1.17241808e+00, # a0 - ], -} +```bash +cd examples/offline_inference/your_model_name +python end2end.py --model your-org/your-model-name --modality text2img --prompts "A cute cat" ``` -## Step 6: Open a Pull Request -When submitting a pull request to add support for a new model, please include the following information in the PR description: +#### 4.2: Add Online Serving Example (OpenAI-Compatible) -+ Output verification: provide generation outputs to verify correctness and model behavior. +Mirror BAGEL’s online serving setup: -+ Inference speed: provide a comparison with the corresponding implementation in Diffusers. +- Server launcher: `examples/online_serving/your_model_name/run_server.sh` + - Wrap `vllm serve ... --omni --port ...` (and `--stage-configs-path ...` if needed) +- Client: `examples/online_serving/your_model_name/openai_chat_client.py` + - Send requests to `POST /v1/chat/completions` + - Support multimodal inputs (e.g., base64 image) if your model needs it +- Doc: `examples/online_serving/your_model_name/README.md` + - Include both “launch server” and “send request”: -+ Parallelism support: specify the supported parallel sizes and any relevant limitations. +```bash +# Terminal 1: launch server +cd examples/online_serving/your_model_name +bash run_server.sh -+ Cache acceleration: check whether the model can be accelerated using Cache-Dit or not. +# Terminal 2: send request +python openai_chat_client.py --prompt "A cute cat" --modality text2img +``` -Providing these details helps reviewers evaluate correctness, performance improvements, and parallel scalability of the new model integration. +### Step 5: Test Your Implementation -# Testing +Before submitting, thoroughly test your implementation. + +#### 5.1: Performance/Speed Check + +Manually compare **latency/throughput** and **output quality** against a Diffusers baseline. + +For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num_inference_steps**, and **guidance settings**, and run multiple trials to reduce randomness. Record the results (and your hardware / driver / CUDA versions) in your PR description. + + +#### 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). +--- + +## Advanced Features + +Once basic implementation works, add advanced features for better performance. + +### torch.compile Support + +Enable automatic compilation for repeated blocks: + +```python +# In your_model_transformer.py + +class YourModelTransformer2DModel(nn.Module): + # Specify which blocks can be compiled + _repeated_blocks = ["YourTransformerBlock"] # List of block class names + + def __init__(self, ...): + super().__init__() + # ... initialization +``` + +vLLM-Omni automatically compiles blocks in `_repeated_blocks` when `torch.compile` is available. + +### Tensor Parallelism + +See detailed guide: [How to add Tensor Parallel support](../features/tensor_parallel.md) + +**Quick setup:** + +1. Replace Linear layers by various parallel linear layers (e.g., `ColumnParallelLinear`) in vLLM +2. Check `tp_size` validity: `hidden_dim`, `num_heads`, and `num_kv_heads` must be divisible by `tp_size` + +**Usage:** Set `tensor_parallel_size` when initializing: +```python +omni = Omni(model="your-model", tensor_parallel_size=2) +``` + +### CFG Parallelism + +See detailed guide: [How to add CFG-Parallel support](../features/cfg_parallel.md) + +**Quick setup:** + +1. Implement `diffuse()` method +2. Inherit `CFGParallelMixin` in your pipeline class + +**Usage:** Set `cfg_parallel_size` when initializing: +```python +omni = Omni(model="your-model", cfg_parallel_size=2) +``` + +### Sequence Parallelism + +See detailed guide: [How to add Sequence Parallel support](../features/sequence_parallel.md) + +**Quick setup:** + +1. Add `_sp_plan` class attribute to transformer +2. Specify where to shard/gather tensors + +**Usage:** Set `ulysses_degree` and `ring_degree` when initializing: +```python +omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2) +``` + +### Cache Acceleration + +#### TeaCache + +See detailed guide: [How to add TeaCache support](../features/teacache.md) + +**Quick setup:** + +1. Write extractor function +2. Register in `EXTRACTOR_REGISTRY` +3. Add polynomial coefficients + +**Usage:** Set `cache_backend` and `cache_config` when initializing: +```python +omni = Omni(model="your-model", + cache_backend="tea_cache", + cache_config={"rel_l1_thresh": 0.2} +) + +``` + + +#### Cache-DiT + +See detailed guide: [How to add Cache-DiT support](../features/cache_dit.md) + +**Quick setup:** + +- For standard models: Works automatically +- For complex architectures: Write custom cache config + +**Usage:** Set `cache_backend` and `cache_config` when initializing: +```python +omni = Omni(model="your-model", + cache_backend="cache_dit", + cache_config={ + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + } +) +``` + +--- + +## Troubleshooting + + +**Issue: ImportError when loading model** + +**Symptoms:** `ModuleNotFoundError` or `ImportError` when calling `Omni(model="your-model")` + +**Causes:** + +1. Model not registered in `registry.py` +2. Wrong class name in registry +3. Missing `__init__.py` exports + + +**Issue: Shape mismatch in attention** + +**Symptoms:** `RuntimeError: shape mismatch` in attention forward + +**Cause:** Incorrect tensor reshaping for vLLM-Omni's attention interface + +**Solution:** Ensure correct shapes: + +```python +# vLLM-Omni expects: [batch, seq_len, num_heads, head_dim] +query = query.view(batch_size, seq_len, self.num_heads, self.head_dim) +key = key.view(batch_size, kv_seq_len, self.num_kv_heads, self.head_dim) +value = value.view(batch_size, kv_seq_len, self.num_kv_heads, self.head_dim) + +hidden_states = self.attn(query, key, value, attn_metadata=attn_metadata) + +# Reshape back: [batch, seq_len, num_heads, head_dim] → [batch, seq_len, hidden_size] +hidden_states = hidden_states.reshape(batch_size, seq_len, -1) +``` + +**Issue: Different outputs compared to Diffusers** + +**Symptoms:** Generated images look different from Diffusers + +**Causes:** + +1. Attention backend differences (FlashAttention vs PyTorch SDPA) +2. Missing normalization or scaling + +**4. Issue: Out of memory (OOM)** + +**Symptoms:** CUDA out of memory errors + +**Solutions:** + +1. **Reduce batch size:** + ```python + omni.generate(prompts=[...], max_batch_size=2) + ``` + +2. **Use smaller image size:** + ```python + sampling_params = OmniDiffusionSamplingParams(height=512, width=512) + ``` + +3. **Enable model offloading:** + ```python + omni = Omni(model="...", enable_cpu_offload=True) + ``` + +4. **Apply vae tiling and slicing** + ```python + omni = Omni(model="...", vae_use_slicing=True, vae_use_tiling=True,) + ``` + +--- + +## Pull Request Checklist + +When submitting a PR to add your model, include: + +**1. Implementation Files** + +- ✅ Transformer model (`xxx_transformer.py`) +- ✅ Pipeline (`pipeline_xxx.py`) +- ✅ Registry entries in `registry.py` +- ✅ `__init__.py` with proper exports + +**2. Example and Tests** + +- ✅ Example script in `examples/` +- ✅ Test file in `tests/e2e/` +- ✅ Documentation (`docs/`) creation or updates + +_Note: End-to-end test files in `tests/e2e/` are optional but strongly recommended. README updates are required for all new models._ + +**3. Documentation Updates** + +- ✅ Add model to supported models table in `docs/models/supported_models.md` +- ✅ If supporting acceleration features (e.g., sequence parallelism, CFG parallel), update acceleration feature tables in: + - `docs/user_guide/diffusion_acceleration.md` + - `docs/user_guide/diffusion/parallelism_acceleration.md` + +--- + +## Model Recipe -## Adding a Model Recipe After implementing and testing your model, please add a model recipe to the [vllm-project/recipes](https://github.com/vllm-project/recipes) repository. This helps other users understand how to use your model with vLLM-Omni. + +**What to Include** + +Your recipe should include: + +1. **Model Overview**: Brief description of the model and its capabilities +2. **Installation Instructions**: Step-by-step setup instructions including: + - Installing vllm-omni and dependencies + - Installing any additional required packages (e.g., xformers, diffusers) + - Any version requirements +3. **Usage Examples**: Command-line examples demonstrating how to run the model +4. **Configuration Details**: Important configuration parameters and their meanings + +**Example** + +For reference, see the [LongCat recipe example](https://github.com/vllm-project/recipes/pull/179) which demonstrates the expected format and structure. + +**Recipe Location** + +Create your recipe file in the appropriate directory structure: +- For organization-specific models: `OrganizationName/ModelName.md` +- For general models: `ModelName.md` + +The recipe should be a Markdown file that provides clear, reproducible instructions for users to get started with your model. + +--- + +## Reference Implementations + +Study these complete examples: + +| Model | Architecture | Key Features | Files | +|-------|--------------|--------------|-------| +| **Qwen-Image** | Dual-stream transformer | CFG-Parallel, SP, TP, Cache | `vllm_omni/diffusion/models/qwen_image/` | +| **Wan2.2** | Video transformer | Dual transformers, SP, CFG-Parallel | `vllm_omni/diffusion/models/wan2_2/` | + +--- + +## Summary + +Adding a diffusion model to vLLM-Omni involves: + +1. ✅ **Adapt transformer** - Replace attention, remove mixins, add config support +2. ✅ **Adapt pipeline** - Change interface, add request handling, extract processing +3. ✅ **Register model** - Add entries to `registry.py` +4. ✅ **Add examples** - Provide runnable scripts +5. ✅ **Test thoroughly** - Verify correctness and performance +6. ✅ **Add advanced features** - Enable parallelism and acceleration (optional) +7. ✅ **Submit PR** - Include verification results and documentation + +**Need help?** Check reference implementations or ask in [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai). diff --git a/docs/features/comfyui.md b/docs/features/comfyui.md new file mode 100644 index 000000000..3b7a9551e --- /dev/null +++ b/docs/features/comfyui.md @@ -0,0 +1,70 @@ +# vLLM-Omni ComfyUI Integration + +vLLM-Omni offers a ComfyUI integration on top of its online serving API. +It can send model inference requests to either a locally running vLLM-Omni service or a remote one. + +## Requirement + +- Python 3.12 or above +- [ComfyUI installed](https://docs.comfy.org/installation/system_requirements) +- [vLLM-Omni installed](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/) on either the same device or another device discoverable via the internet. +- No need to install additional packages apart from those already required by ComfyUI. + +!!! tip + If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them. + + +## Installation + +Copy the `apps/ComfyUI-vLLM-Omni` folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`. + +If you are running ComfyUI during copying, you should restart ComfyUI to load this extension. + +!!! tip + You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info. + +On the device and virtual environment you run ComfyUI, launch ComfyUI with +```bash +cd ComfyUI + +# The regular way +python main.py + +# If you are mainly using this node, launch it faster with +python main.py --cpu +``` + +On the device and virtual environment you run vLLM-Omni, start a model service with +```bash +vllm serve The_Model_ID_to_Serve --omni --port 8000 +``` + +Check **ComfyUI's sidebar -> Node Library**. There should be a new folder named **vLLM-Omni**. +If no, check your shell running the ComfyUI process. There may be some error messages before the line `Import times for custom nodes:` and the line `To see the GUI go to: http://127.0.0.1:8188`. + +## Quickstart + +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 +- **TTS** and **TTS Voice Clone** for TTS tasks + +This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**) + +!!! info + The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support. + +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. +- 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. + - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages. + +## Examples & Screenshots + +Please read the [ComfyUI integration's README](https://github.com/vllm-project/vllm-omni/tree/main/apps/ComfyUI-vLLM-Omni) for more info. diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0ba642c92..0fbbe6eb7 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -20,6 +20,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` | | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | +| `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/user_guide/diffusion/parallelism_acceleration.md b/docs/user_guide/diffusion/parallelism_acceleration.md index f91a626e0..d536f4bc7 100644 --- a/docs/user_guide/diffusion/parallelism_acceleration.md +++ b/docs/user_guide/diffusion/parallelism_acceleration.md @@ -233,143 +233,6 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m | Hybrid Ulysses + Ring | 2 | 2 | 24.3s | 1.87x | -##### How to parallelize a new model - -NOTE: "Terminology: SP vs CP" - Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py). - We use "Sequence Parallelism" to align with vLLM-Omni's terminology. - ---- - -###### Non-intrusive `_sp_plan` (Recommended) - -The `_sp_plan` mechanism allows SP without modifying `forward()` logic. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries. - -**Requirements for `forward()` function:** - -- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** (not inline Python operations) -- If your `forward()` contains inline tensor operations (e.g., `torch.cat`, `pad_sequence`) that need sharding, **extract them into a submodule** - -**When to create a submodule:** - -```python -# ❌ BAD: Inline operations - hooks cannot intercept -def forward(self, x, cap_feats): - unified = torch.cat([x, cap_feats], dim=1) # Cannot be sharded via _sp_plan - ... - -# ✅ GOOD: Extract into a submodule -class UnifiedPrepare(nn.Module): - def forward(self, x, cap_feats): - return torch.cat([x, cap_feats], dim=1) # Now can be sharded via _sp_plan - -class MyModel(nn.Module): - def __init__(self): - self.unified_prepare = UnifiedPrepare() # Submodule - - def forward(self, x, cap_feats): - unified = self.unified_prepare(x, cap_feats) # Hook can intercept here -``` - ---- - -###### Defining `_sp_plan` - -**Type definitions** (see [diffusers `_modeling_parallel.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py) for reference): - -```python -from vllm_omni.diffusion.distributed.sp_plan import ( - SequenceParallelInput, # Corresponds to diffusers' ContextParallelInput - SequenceParallelOutput, # Corresponds to diffusers' ContextParallelOutput -) -``` - -| Parameter | Description | -|-----------|-------------| -| `split_dim` | Dimension to split/gather (usually `1` for sequence) | -| `expected_dims` | Expected tensor rank for validation (optional) | -| `split_output` | `False`: shard **input** parameters; `True`: shard **output** tensors | -| `auto_pad` | Auto-pad if sequence not divisible by world_size (Ulysses only) | - -**Key naming convention:** - -| Key | Meaning | Python equivalent | -|-----|---------|-------------------| -| `""` | Root model | `model` | -| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` | -| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` | -| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` | - -**Dictionary key types:** - -| Key type | `split_output` | Description | -|----------|----------------|-------------| -| `"param_name"` (str) | `False` | Shard **input parameter** by name | -| `0`, `1` (int) | `True` | Shard **output tuple** by index | - -**Example** (similar to [diffusers `transformer_wan.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py)): - -```python -class MyTransformer(nn.Module): - _sp_plan = { - # Shard rope module OUTPUTS (returns tuple) - "rope": { - 0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # cos - 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True), # sin - }, - # Shard transformer block INPUT parameter - "blocks.0": { - "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), - }, - # Gather at final projection - "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), - } -``` - ---- - -###### Hook flow - -``` -Input → [SequenceParallelSplitHook: pre_forward] → Module.forward() → [post_forward] → ... - ↓ -... → [SequenceParallelGatherHook: post_forward] → Output -``` - -1. **SplitHook** shards tensors before/after the target module -2. **Attention layers** handle Ulysses/Ring communication internally -3. **GatherHook** collects sharded outputs - -The framework automatically applies these hooks when `sequence_parallel_size > 1`. - ---- - -###### Method 2: Intrusive modification (For complex cases) - -For models with dynamic sharding logic that cannot be expressed via `_sp_plan`: - -```python -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 ... - output = sp_gather(output, dim=1) - return output -``` - ---- - -###### Choosing the right approach - -| Scenario | Approach | -|----------|----------| -| Standard transformer | `_sp_plan` | -| Inline tensor ops need sharding | Extract to submodule + `_sp_plan` | -| Dynamic/conditional sharding | Intrusive modification | - - ### CFG-Parallel #### Offline Inference @@ -426,142 +289,3 @@ You can enable CFG-Parallel in online serving for diffusion models via `--cfg-pa ```bash vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2 ``` - -#### How to parallelize a pipeline - -This section describes how to add CFG-Parallel to a diffusion **pipeline**. We use the Qwen-Image pipeline (`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`) as the reference implementation. - -In `QwenImagePipeline`, each diffusion step runs two denoiser forward passes sequentially: - -- positive (prompt-conditioned) -- negative (negative-prompt-conditioned) - -CFG-Parallel assigns these two branches to different ranks in the **CFG group** and synchronizes the results. - -vLLM-omni provides `CFGParallelMixin` base class that encapsulates the CFG parallel logic. By inheriting from this mixin and calling its methods, pipelines can easily implement CFG parallel without writing repetitive code. - -**Key Methods in CFGParallelMixin:** -- `predict_noise_maybe_with_cfg()`: Automatically handles CFG parallel noise prediction -- `scheduler_step_maybe_with_cfg()`: Scheduler step with automatic CFG rank synchronization - -**Example Implementation:** - -```python -class QwenImageCFGParallelMixin(CFGParallelMixin): - """ - Base Mixin class for Qwen Image pipelines providing shared CFG methods. - """ - - def diffuse( - self, - prompt_embeds: torch.Tensor, - prompt_embeds_mask: torch.Tensor, - negative_prompt_embeds: torch.Tensor, - negative_prompt_embeds_mask: torch.Tensor, - latents: torch.Tensor, - img_shapes: torch.Tensor, - txt_seq_lens: torch.Tensor, - negative_txt_seq_lens: torch.Tensor, - timesteps: torch.Tensor, - do_true_cfg: bool, - guidance: torch.Tensor, - true_cfg_scale: float, - image_latents: torch.Tensor | None = None, - cfg_normalize: bool = True, - additional_transformer_kwargs: dict[str, Any] | None = None, - ) -> torch.Tensor: - self.transformer.do_true_cfg = do_true_cfg - - for i, t in enumerate(timesteps): - timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) - - # Prepare kwargs for positive (conditional) prediction - positive_kwargs = { - "hidden_states": latents, - "timestep": timestep / 1000, - "guidance": guidance, - "encoder_hidden_states_mask": prompt_embeds_mask, - "encoder_hidden_states": prompt_embeds, - "img_shapes": img_shapes, - "txt_seq_lens": txt_seq_lens, - } - - # Prepare kwargs for negative (unconditional) prediction - if do_true_cfg: - negative_kwargs = { - "hidden_states": latents, - "timestep": timestep / 1000, - "guidance": guidance, - "encoder_hidden_states_mask": negative_prompt_embeds_mask, - "encoder_hidden_states": negative_prompt_embeds, - "img_shapes": img_shapes, - "txt_seq_lens": negative_txt_seq_lens, - } - else: - negative_kwargs = None - - # Predict noise with automatic CFG parallel handling - # - In CFG parallel mode: rank0 computes positive, rank1 computes negative - # - Automatically gathers results and combines them on rank0 - noise_pred = self.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=true_cfg_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=cfg_normalize, - ) - - # Step scheduler with automatic CFG synchronization - # - Only rank0 computes the scheduler step - # - Automatically broadcasts updated latents to all ranks - latents = self.scheduler_step_maybe_with_cfg( - noise_pred, t, latents, do_true_cfg - ) - - return latents -``` - -**How it works:** -1. Prepare separate `positive_kwargs` and `negative_kwargs` for conditional and unconditional predictions -2. Call `predict_noise_maybe_with_cfg()` which: - - Detects if CFG parallel is enabled (`get_classifier_free_guidance_world_size() > 1`) - - Distributes computation: rank0 processes positive, rank1 processes negative - - Gathers predictions and combines them using `combine_cfg_noise()` on rank0 - - Returns combined noise prediction (only valid on rank0) -3. Call `scheduler_step_maybe_with_cfg()` which: - - Only rank0 computes the scheduler step - - Broadcasts the updated latents to all ranks for synchronization - -**How to customize** - -Some pipelines may need to customize the following functions in `CFGParallelMixin`: -1. You may need to edit `predict_noise` function for custom behaviors. -```python -def predict_noise(self, *args, **kwargs): - """ - Forward pass through transformer to predict noise. - - Subclasses should override this if they need custom behavior, - but the default implementation calls self.transformer. - """ - return self.transformer(*args, **kwargs)[0] - -``` -2. The default normalization function after combining the noise predictions from both branches is as follows. You may need to customize it. -```python -def cfg_normalize_function(self, noise_pred, comb_pred): - """ - Normalize the combined noise prediction. - - Args: - noise_pred: positive noise prediction - comb_pred: combined noise prediction after CFG - - Returns: - Normalized noise prediction tensor - """ - cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True) - noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True) - noise_pred = comb_pred * (cond_norm / noise_norm) - return noise_pred -``` diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index 397fd333e..8fa412065 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -36,7 +36,7 @@ def parse_args(): ) # OmniLLM init args - parser.add_argument("--enable-stats", action="store_true", default=False) + parser.add_argument("--log-stats", action="store_true", default=False) parser.add_argument("--init-sleep-seconds", type=int, default=20) parser.add_argument("--batch-timeout", type=int, default=5) parser.add_argument("--init-timeout", type=int, default=300) @@ -120,7 +120,7 @@ def main(): omni_kwargs.update( { - "log_stats": args.enable_stats, + "log_stats": args.log_stats, "init_sleep_seconds": args.init_sleep_seconds, "batch_timeout": args.batch_timeout, "init_timeout": args.init_timeout, diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 63b08202a..7cd8e1873 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -322,7 +322,7 @@ def main(args): query_result = query_func() omni_llm = Omni( model=model_name, - log_stats=args.enable_stats, + log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, batch_timeout=args.batch_timeout, init_timeout=args.init_timeout, @@ -439,7 +439,7 @@ def parse_args(): help="Query type.", ) parser.add_argument( - "--enable-stats", + "--log-stats", action="store_true", default=False, help="Enable writing detailed statistics (default: disabled)", diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 9e8e8e6c3..509886a02 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -297,7 +297,7 @@ def main(args): omni_llm = Omni( model=model_name, stage_configs_path=args.stage_configs_path, - log_stats=args.enable_stats, + log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, ) @@ -424,7 +424,7 @@ def parse_args(): help="Query type.", ) parser.add_argument( - "--enable-stats", + "--log-stats", action="store_true", default=False, help="Enable writing detailed statistics (default: disabled)", diff --git a/examples/offline_inference/qwen3_tts/end2end.py b/examples/offline_inference/qwen3_tts/end2end.py index 18bb9dedc..93aeba3ca 100644 --- a/examples/offline_inference/qwen3_tts/end2end.py +++ b/examples/offline_inference/qwen3_tts/end2end.py @@ -219,7 +219,7 @@ def main(args): omni = Omni( model=model_name, stage_configs_path=args.stage_configs_path, - log_stats=args.enable_stats, + log_stats=args.log_stats, stage_init_timeout=args.stage_init_timeout, ) @@ -275,7 +275,7 @@ def parse_args(): help="Query type.", ) parser.add_argument( - "--enable-stats", + "--log-stats", action="store_true", default=False, help="Enable writing detailed statistics (default: disabled)", diff --git a/pyproject.toml b/pyproject.toml index 706d0152e..0748b7eab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -175,6 +175,10 @@ markers = [ "slow: Slow tests (may skip in quick CI)", "benchmark: Benchmark tests", ] +filterwarnings = [ + "ignore:.*does not have '__test__' attribute.*:UserWarning", + "ignore:.*does not have '__bases__' attribute.*:UserWarning", +] [tool.typos.default] extend-ignore-identifiers-re = [ @@ -184,6 +188,7 @@ extend-ignore-identifiers-re = [ ".*arange", ".*MoBA", ".*temperal_downsample", + ".*_nd", ".*nothink.*", ".*NOTHINK.*", ".*nin.*", diff --git a/requirements/common.txt b/requirements/common.txt index 19ebca949..2331e0f5d 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -12,3 +12,4 @@ torchsde>=0.2.6 openai-whisper>=20250625 imageio[ffmpeg]>=2.37.2 sox>=1.5.0 +prettytable>=3.8.0 diff --git a/tests/benchmarks/test_serve_cli.py b/tests/benchmarks/test_serve_cli.py index 2c624b8e7..40244eb57 100644 --- a/tests/benchmarks/test_serve_cli.py +++ b/tests/benchmarks/test_serve_cli.py @@ -4,6 +4,7 @@ from pathlib import Path import pytest from tests.conftest import OmniServer +from tests.utils import hardware_test models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen3_omni_ci.yaml")] @@ -29,6 +30,9 @@ def omni_server(request): print("OmniServer stopped") +@pytest.mark.core_model +@pytest.mark.benchmark +@hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_bench_serve_chat(omni_server): command = [ diff --git a/tests/conftest.py b/tests/conftest.py index 84fbb91b7..98a21d43c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -588,8 +588,6 @@ def convert_audio_to_text(audio_data): """ Convert base64 encoded audio data to text using speech recognition. """ - import whisper - audio_data = base64.b64decode(audio_data) output_path = f"./test_{int(time.time())}" with open(output_path, "wb") as audio_file: @@ -597,6 +595,13 @@ def convert_audio_to_text(audio_data): print(f"audio data is saved: {output_path}") + text = convert_audio_file_to_text(output_path=output_path) + return text + + +def convert_audio_file_to_text(output_path): + import whisper + model = whisper.load_model("base") text = model.transcribe( output_path, @@ -604,6 +609,10 @@ def convert_audio_to_text(audio_data): word_timestamps=True, condition_on_previous_text=False, )["text"] + del model + if torch.cuda.is_available(): + gc.collect() + torch.cuda.empty_cache() if text: return text else: @@ -614,7 +623,6 @@ def merge_base64_and_convert_to_text(base64_list): """ Merge a list of base64 encoded audio data and convert to text. """ - import whisper from pydub import AudioSegment merged_audio = None @@ -627,17 +635,8 @@ def merge_base64_and_convert_to_text(base64_list): merged_audio += seg output_path = f"./test_{int(time.time())}" merged_audio.export(output_path, format="wav") - model = whisper.load_model("base") - text = model.transcribe( - output_path, - temperature=0.0, - word_timestamps=True, - condition_on_previous_text=False, - )["text"] - if text: - return text - else: - return "" + text = convert_audio_file_to_text(output_path) + return text def modify_stage_config( @@ -886,6 +885,7 @@ class OmniServer: model: str, serve_args: list[str], *, + port: int | None = None, env_dict: dict[str, str] | None = None, ) -> None: _run_pre_test_cleanup(enable_force=True) @@ -896,7 +896,10 @@ class OmniServer: self.env_dict = env_dict self.proc: subprocess.Popen | None = None self.host = "127.0.0.1" - self.port = get_open_port() + if port is None: + self.port = get_open_port() + else: + self.port = port def _start_server(self) -> None: """Start the vLLM-Omni server subprocess.""" diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index ed9301410..a9312f4b1 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/tests/diffusion/cache/test_cache_backends.py @@ -22,6 +22,8 @@ from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend from vllm_omni.diffusion.data import DiffusionCacheConfig +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class TestCacheDiTBackend: """Test CacheDiTBackend implementation.""" diff --git a/tests/diffusion/lora/test_base_linear.py b/tests/diffusion/lora/test_base_linear.py index 42bdf6526..9386d0909 100644 --- a/tests/diffusion/lora/test_base_linear.py +++ b/tests/diffusion/lora/test_base_linear.py @@ -5,10 +5,13 @@ from __future__ import annotations from dataclasses import dataclass +import pytest import torch from vllm_omni.diffusion.lora.layers.base_linear import DiffusionBaseLinearLayerWithLoRA +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + @dataclass class _DummyLoRAConfig: diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 84fafe3bc..e5f1c4700 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -3,6 +3,7 @@ from __future__ import annotations +import pytest import torch from vllm.lora.lora_weights import LoRALayerWeights from vllm.lora.utils import get_supported_lora_modules @@ -11,6 +12,8 @@ from vllm.model_executor.layers.linear import LinearBase from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.lora.request import LoRARequest +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class _DummyLoRALayer: def __init__(self, n_slices: int, output_slices: tuple[int, ...]): diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 220f210a3..6e57355a2 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -17,6 +17,8 @@ import torch from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + @pytest.fixture def mock_od_config(): diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py index 8c7ff79ca..2b35718e7 100644 --- a/tests/distributed/omni_connectors/test_kv_flow.py +++ b/tests/distributed/omni_connectors/test_kv_flow.py @@ -1,7 +1,6 @@ import pytest import torch -from tests.utils import hardware_test from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( OmniKVCacheConfig, @@ -9,6 +8,8 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.cache] + class MockConnector: def __init__(self): @@ -58,11 +59,6 @@ def common_constants(): } -@pytest.mark.cache -@hardware_test( - res={"cuda": "L4"}, - num_cards=2, -) def test_manager_extraction(kv_config, mock_connector, common_constants): """Test extraction and sending logic in OmniKVTransferManager.""" num_layers = common_constants["num_layers"] @@ -109,11 +105,6 @@ def test_manager_extraction(kv_config, mock_connector, common_constants): assert data["layer_blocks"]["key_cache"][0].shape == expected_shape -@pytest.mark.cache -@hardware_test( - res={"cuda": "L4"}, - num_cards=2, -) def test_manager_reception(kv_config, mock_connector, common_constants): """Test reception and injection logic in OmniKVTransferManager.""" num_layers = common_constants["num_layers"] @@ -171,11 +162,6 @@ def test_manager_reception(kv_config, mock_connector, common_constants): assert req.kv_metadata["seq_len"] == seq_len -@pytest.mark.cache -@hardware_test( - res={"cuda": "L4"}, - num_cards=2, -) def test_integration_flow(common_constants): """Simulate extraction -> connector -> reception.""" num_layers = common_constants["num_layers"] @@ -211,7 +197,8 @@ def test_integration_flow(common_constants): recv_timeout=1.0, ) receiver_manager = OmniKVTransferManager(receiver_config) - receiver_manager._connector = connector # Share the same mock connector instance + # Share the same mock connector instance + receiver_manager._connector = connector req = OmniDiffusionRequest( prompts=["test_integ"], @@ -228,11 +215,6 @@ def test_integration_flow(common_constants): assert req.kv_metadata["seq_len"] == 10 -@pytest.mark.cache -@hardware_test( - res={"cuda": "L4"}, - num_cards=2, -) def test_manager_extraction_no_connector(kv_config, common_constants): """Test extraction when connector is unavailable (should still return IDs).""" block_size = common_constants["block_size"] diff --git a/tests/e2e/offline_inference/test_cache_dit.py b/tests/e2e/offline_inference/test_cache_dit.py index 281ac48dc..65c59c6f7 100644 --- a/tests/e2e/offline_inference/test_cache_dit.py +++ b/tests/e2e/offline_inference/test_cache_dit.py @@ -15,6 +15,7 @@ from pathlib import Path import pytest import torch +from tests.utils import hardware_test from vllm_omni.inputs.data import OmniDiffusionSamplingParams # ruff: noqa: E402 @@ -32,6 +33,10 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" models = ["riverclouds/qwen_image_random"] +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cache +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) @pytest.mark.parametrize("model_name", models) def test_cache_dit(model_name: str): """Test cache-dit backend with diffusion model.""" diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index 35e106df8..e46d7ec96 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_cpu_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 GPUMemoryMonitor, hardware_test from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -45,6 +45,9 @@ def inference(model_name: str, offload: bool = True): return peak +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) @pytest.mark.skipif(current_omni_platform.is_npu() or current_omni_platform.is_rocm(), reason="Hardware not supported") @pytest.mark.parametrize("model_name", models) def test_cpu_offload_diffusion_model(model_name: str): diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni.py b/tests/e2e/offline_inference/test_qwen2_5_omni.py index af9d793c1..eda2f28b5 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni.py @@ -13,10 +13,10 @@ from vllm.assets.video import VideoAsset from vllm.envs import VLLM_USE_MODELSCOPE from vllm.multimodal.image import convert_image_mode +from tests.utils import create_new_process_for_each_test, hardware_test from vllm_omni.platforms import current_omni_platform from .conftest import OmniRunner -from .utils import create_new_process_for_each_test models = ["Qwen/Qwen2.5-Omni-3B"] @@ -34,8 +34,10 @@ test_params = [(model, stage_config) for model in models] @pytest.mark.core_model -@pytest.mark.parametrize("test_config", test_params) +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @create_new_process_for_each_test("spawn") +@pytest.mark.parametrize("test_config", test_params) def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None: """Test processing audio, image, and video together, generating audio output.""" model, stage_config_path = test_config @@ -94,8 +96,10 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t @pytest.mark.core_model -@pytest.mark.parametrize("test_config", test_params) +@pytest.mark.omni +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @create_new_process_for_each_test("spawn") +@pytest.mark.parametrize("test_config", test_params) def test_mixed_modalities_to_text_only(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None: """Test processing audio, image, and video together, generating audio output.""" model, stage_config_path = test_config diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index 2c5c66fe3..ecd09fdd3 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -14,6 +14,7 @@ from pathlib import Path import pytest from vllm.assets.video import VideoAsset +from tests.utils import hardware_test from vllm_omni.platforms import current_omni_platform from .conftest import OmniRunner @@ -31,6 +32,9 @@ else: test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("test_config", test_params) def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None: """Test processing video, generating audio output.""" diff --git a/tests/e2e/offline_inference/test_sequence_parallel.py b/tests/e2e/offline_inference/test_sequence_parallel.py index 3e7bb5617..d66104aa1 100644 --- a/tests/e2e/offline_inference/test_sequence_parallel.py +++ b/tests/e2e/offline_inference/test_sequence_parallel.py @@ -20,16 +20,17 @@ import torch import torch.distributed as dist from PIL import Image +from tests.utils import hardware_test +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform # ruff: noqa: E402 REPO_ROOT = Path(__file__).resolve().parents[3] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from vllm_omni import Omni -from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.platforms import current_omni_platform # Test configuration MODELS = ["riverclouds/qwen_image_random"] @@ -144,10 +145,20 @@ def _run_inference( # SP configurations: (ulysses_degree, ring_degree, height, width, warmup, is_perf_test) # - warmup: whether to run warmup for this SP config # - is_perf_test: whether this is a performance test (show speedup metrics) -SP_CONFIGS = [ - (2, 1, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True), # Ulysses-2 - performance test +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 - (2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False), # Hybrid - correctness only + # Hybrid - correctness only + (2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False), +] + +SP_CONFIGS_L3 = [ + # 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), (4, 1, 272, 272, False, False), # Ulysses-4 - shape and correctness ] @@ -162,6 +173,10 @@ def _get_sp_mode(ulysses_degree: int, ring_degree: int) -> str: return f"hybrid-{ulysses_degree}x{ring_degree}" +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 2, "rocm": 2}) @pytest.mark.parametrize("model_name", MODELS) def test_sp_correctness(model_name: str): """Test that SP inference produces correct outputs and measure performance. @@ -184,7 +199,130 @@ def test_sp_correctness(model_name: str): print(f"Available GPUs: {device_count}") print("=" * 70) - for ulysses_degree, ring_degree, height, width, sp_warmup, is_perf_test in SP_CONFIGS: + for ulysses_degree, ring_degree, height, width, sp_warmup, is_perf_test in SP_CONFIGS_L2: + sp_size = ulysses_degree * ring_degree + sp_mode = _get_sp_mode(ulysses_degree, ring_degree) + + if device_count < sp_size: + print(f"\n[{sp_mode}] SKIPPED (requires {sp_size} GPUs)") + continue + + # Determine baseline warmup: only for default size (performance tests) + cache_key = (height, width) + baseline_warmup = height == DEFAULT_HEIGHT and width == DEFAULT_WIDTH + + # Get or compute baseline for this (height, width) + if cache_key not in baseline_cache: + print(f"\n--- Running baseline {height}x{width} (warmup={baseline_warmup}) ---") + baseline = _run_inference( + model_name, + torch.bfloat16, + "sdpa", + height=height, + width=width, + warmup=baseline_warmup, + ) + assert len(baseline.images) == 1 + baseline_cache[cache_key] = baseline + print(f"[baseline] {height}x{width}: {baseline.elapsed_ms:.0f}ms") + else: + baseline = baseline_cache[cache_key] + + # Run SP + print(f"\n--- Running {sp_mode} (warmup={sp_warmup}) ---") + sp_result = _run_inference( + model_name, + torch.bfloat16, + "sdpa", + ulysses_degree=ulysses_degree, + ring_degree=ring_degree, + height=height, + width=width, + warmup=sp_warmup, + ) + assert len(sp_result.images) == 1 + + # Compare outputs (correctness) + mean_diff, max_diff = _diff_metrics(baseline.images[0], sp_result.images[0]) + + # Build result entry + result = { + "mode": sp_mode, + "sp_size": sp_size, + "height": height, + "width": width, + "baseline_ms": baseline.elapsed_ms, + "sp_ms": sp_result.elapsed_ms, + "mean_diff": mean_diff, + "max_diff": max_diff, + "is_perf_test": is_perf_test, + } + results.append(result) + + # Output based on test type + if is_perf_test: + speedup = baseline.elapsed_ms / sp_result.elapsed_ms if sp_result.elapsed_ms > 0 else 0 + result["speedup"] = speedup + print( + f"[{sp_mode}] {sp_size} GPUs | " + f"baseline: {baseline.elapsed_ms:.0f}ms, sp: {sp_result.elapsed_ms:.0f}ms, " + f"speedup: {speedup:.2f}x" + ) + else: + print(f"[{sp_mode}] {sp_size} GPUs | sp: {sp_result.elapsed_ms:.0f}ms (correctness only)") + + print(f"[{sp_mode}] diff: mean={mean_diff:.6e}, max={max_diff:.6e}") + + # Assert correctness + assert mean_diff <= DIFF_MEAN_THRESHOLD and max_diff <= DIFF_MAX_THRESHOLD, ( + f"[{sp_mode}] SP output differs from baseline: mean={mean_diff:.6e}, max={max_diff:.6e}" + ) + + # Summary + print("\n" + "=" * 70) + print("SUMMARY") + print("=" * 70) + print(f"{'Mode':<15} {'GPUs':<6} {'Size':<10} {'Baseline':<12} {'SP':<12} {'Speedup':<10} {'Status'}") + print("-" * 70) + for r in results: + speedup_str = f"{r['speedup']:.2f}x" if r.get("speedup") else "N/A" + baseline_str = f"{r['baseline_ms']:.0f}ms" if r["is_perf_test"] else "N/A" + status = "PASS" if r["mean_diff"] <= DIFF_MEAN_THRESHOLD else "FAIL" + print( + f"{r['mode']:<15} {r['sp_size']:<6} {r['height']}x{r['width']:<5} " + f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}" + ) + print("=" * 70) + + +# TODO: After PR#1272 is merged, add markers +# @pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +@pytest.mark.parametrize("model_name", MODELS) +def test_sp_correctness_advanced(model_name: str): + """Test that SP inference produces correct outputs and measure performance. + + Runs baseline once per unique (height, width), then tests all SP configs. + + Note: Run with `pytest -v -s` to see detailed output. + """ + device_count = current_omni_platform.get_device_count() + + # Cache baseline results by (height, width) + # Key: (height, width), Value: (result, warmup_used) + baseline_cache: dict[tuple[int, int], InferenceResult] = {} + + # Collect results for summary + results: list[dict] = [] + + print("\n" + "=" * 70) + print(f"Sequence Parallel Test - Model: {model_name}") + print(f"Available GPUs: {device_count}") + print("=" * 70) + + for ulysses_degree, ring_degree, height, width, sp_warmup, is_perf_test in SP_CONFIGS_L3: sp_size = ulysses_degree * ring_degree sp_mode = _get_sp_mode(ulysses_degree, ring_degree) diff --git a/tests/e2e/offline_inference/test_stable_audio_model.py b/tests/e2e/offline_inference/test_stable_audio_model.py index df2ca5e42..ad879d751 100644 --- a/tests/e2e/offline_inference/test_stable_audio_model.py +++ b/tests/e2e/offline_inference/test_stable_audio_model.py @@ -5,6 +5,7 @@ import numpy as np import pytest import torch +from tests.utils import hardware_test from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -19,6 +20,9 @@ from vllm_omni import Omni models = ["linyueqian/stable_audio_random"] +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}) @pytest.mark.parametrize("model_name", models) def test_stable_audio_model(model_name: str): m = Omni(model=model_name) diff --git a/tests/e2e/offline_inference/test_t2i_model.py b/tests/e2e/offline_inference/test_t2i_model.py index 74c1bd1ce..e2cb09f30 100644 --- a/tests/e2e/offline_inference/test_t2i_model.py +++ b/tests/e2e/offline_inference/test_t2i_model.py @@ -5,6 +5,7 @@ from pathlib import Path import pytest import torch +from tests.utils import hardware_test from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -32,6 +33,9 @@ elif current_omni_platform.is_rocm(): models = ["Tongyi-MAI/Z-Image-Turbo"] +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 2}) @pytest.mark.parametrize("model_name", models) def test_diffusion_model(model_name: str): m = None diff --git a/tests/e2e/offline_inference/test_teacache.py b/tests/e2e/offline_inference/test_teacache.py index 7d6261388..d97991779 100644 --- a/tests/e2e/offline_inference/test_teacache.py +++ b/tests/e2e/offline_inference/test_teacache.py @@ -15,6 +15,7 @@ from pathlib import Path import pytest import torch +from tests.utils import hardware_test from vllm_omni.inputs.data import OmniDiffusionSamplingParams # ruff: noqa: E402 @@ -31,6 +32,10 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" models = ["riverclouds/qwen_image_random"] +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.cache +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) @pytest.mark.parametrize("model_name", models) def test_teacache(model_name: str): """Test TeaCache backend with diffusion model.""" diff --git a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py index 0e3b97ec3..2d051e5aa 100644 --- a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py +++ b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py @@ -12,18 +12,18 @@ 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 vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform # ruff: noqa: E402 REPO_ROOT = Path(__file__).resolve().parents[2] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from tests.utils import GPUMemoryMonitor -from vllm_omni import Omni -from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform # os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -127,7 +127,10 @@ def _run_zimage_generate( cleanup_dist_env_and_memory() -@pytest.mark.integration +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) 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.") diff --git a/tests/e2e/online_serving/test_async_omni.py b/tests/e2e/online_serving/test_async_omni.py index d90727e1b..cab3e6e22 100644 --- a/tests/e2e/online_serving/test_async_omni.py +++ b/tests/e2e/online_serving/test_async_omni.py @@ -8,7 +8,8 @@ import pytest from vllm import SamplingParams from vllm.inputs import PromptType -from vllm_omni.entrypoints.async_omni import AsyncOmni, ClientRequestState +from tests.utils import hardware_test +from vllm_omni.entrypoints.async_omni import AsyncOmni os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -57,6 +58,9 @@ async def generate( return count, request_id +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.asyncio async def test_abort(): with ExitStack() as after: @@ -111,51 +115,3 @@ async def test_abort(): num_generated_tokens, request_id = await task assert num_generated_tokens == NUM_EXPECTED_TOKENS await asyncio.sleep(5) - - -@pytest.mark.asyncio -async def test_build_and_log_summary(monkeypatch): - from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e - - RealCRS = ClientRequestState - capture_metrics = {} - - class MockCRS(RealCRS): - def __init__(self, request_id: str): - super().__init__(request_id) - capture_metrics[request_id] = self - - monkeypatch.setattr("vllm_omni.entrypoints.async_omni.ClientRequestState", MockCRS) - monkeypatch.setattr("vllm_omni.entrypoints.client_request_state.ClientRequestState", MockCRS) - - with ExitStack() as after: - # Avoid SHM IPC in tests to prevent /dev/shm exhaustion and SIGBUS. - engine = AsyncOmni( - model=model, - stage_configs_path=stage_config, - shm_threshold_bytes=sys.maxsize, - ) - after.callback(engine.shutdown) - prompt = "Hello my name is Robert and " - NUM_EXPECTED_TOKENS = 64 - NUM_REQUESTS = 3 - request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)] - - # Create concurrent requests. - tasks: list[asyncio.Task] = [] - for idx, request_id in enumerate(request_ids): - tasks.append(asyncio.create_task(generate(engine, request_id, prompt, NUM_EXPECTED_TOKENS))) - - # Confirm the requests are okay. - for idx, task in enumerate(tasks): - await task - output_modalities = ["text"] - final_stage_id_for_e2e = get_final_stage_id_for_e2e( - output_modalities, engine.output_modalities, engine.stage_list - ) - summary = capture_metrics[request_ids[idx]].metrics.build_and_log_summary(final_stage_id_for_e2e) - - # Check that total tokens matches sum of stage tokens. - assert summary["e2e_total_tokens"] == sum(stage["tokens"] for stage in summary["stages"]) - # Check that total time matches sum of stage times. - assert summary["e2e_total_time_ms"] >= sum(stage["total_time_ms"] for stage in summary["stages"]) diff --git a/tests/e2e/online_serving/test_image_gen_edit.py b/tests/e2e/online_serving/test_image_gen_edit.py index 8db0d50fb..7db740f20 100644 --- a/tests/e2e/online_serving/test_image_gen_edit.py +++ b/tests/e2e/online_serving/test_image_gen_edit.py @@ -22,6 +22,8 @@ from PIL import Image from vllm.assets.image import ImageAsset from vllm.utils.network_utils import get_open_port +from tests.utils import hardware_test + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" # Increase timeout for downloading assets from S3 (default 5s is too short for CI) os.environ.setdefault("VLLM_IMAGE_FETCH_TIMEOUT", "60") @@ -178,6 +180,9 @@ def _decode_data_url_to_image_bytes(data_url: str) -> bytes: return base64.b64decode(b64_data) +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_i2i_multi_image_input_qwen_image_edit_2509( omni_server, diff --git a/tests/e2e/online_serving/test_images_generations_lora.py b/tests/e2e/online_serving/test_images_generations_lora.py index e912c420d..85b80cafb 100644 --- a/tests/e2e/online_serving/test_images_generations_lora.py +++ b/tests/e2e/online_serving/test_images_generations_lora.py @@ -23,6 +23,7 @@ from PIL import Image from safetensors.torch import save_file from tests.conftest import OmniServer +from tests.utils import hardware_test os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -144,6 +145,9 @@ def _basic_payload() -> dict: } +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) def test_images_generations_per_request_lora_switching(omni_server: OmniServer, tmp_path: Path) -> None: # Base generation. base_img = _post_images(omni_server, _basic_payload()) diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index 073419fb8..7f879939f 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -28,6 +28,7 @@ from tests.conftest import ( merge_base64_and_convert_to_text, modify_stage_config, ) +from tests.utils import hardware_test from vllm_omni.platforms import current_omni_platform models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] @@ -145,6 +146,9 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> None: """ @@ -215,6 +219,9 @@ def test_mix_to_text_audio_001(client: openai.OpenAI, omni_server, request) -> N assert similarity > 0.9, "The audio content is not same as the text" +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_text_audio_001(client: openai.OpenAI, omni_server) -> None: """ diff --git a/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml new file mode 100644 index 000000000..aaed4de92 --- /dev/null +++ b/tests/e2e/stage_configs/qwen2_5_omni_ci.yaml @@ -0,0 +1,106 @@ +# stage config for running qwen2.5-omni with architecture of OmniLLM. + +# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090). +# This config is optimized for CI e2e tests. +stage_args: + - stage_id: 0 + runtime: + process: true # Run this stage in a separate process + devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device) + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: Qwen2_5OmniForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + max_model_len: 32768 + max_num_batched_tokens: 32768 + max_num_seqs: 1 + gpu_memory_utilization: 0.8 + skip_mm_profiling: true + enforce_eager: true # Now we only support eager mode + trust_remote_code: true + engine_output_type: latent + enable_prefix_caching: false + is_comprehension: true + final_output: true + final_output_type: text + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 128 + seed: 42 + detokenize: True + repetition_penalty: 1.1 + - stage_id: 1 + runtime: + process: true + devices: "1" + max_batch_size: 1 + engine_args: + model_stage: talker + model_arch: Qwen2_5OmniForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + max_model_len: 32768 + max_num_batched_tokens: 32768 + max_num_seqs: 1 + gpu_memory_utilization: 0.8 + skip_mm_profiling: true + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + engine_output_type: latent + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker + default_sampling_params: + temperature: 0.9 + top_p: 0.8 + top_k: 40 + max_tokens: 128 + seed: 42 + detokenize: True + repetition_penalty: 1.05 + stop_token_ids: [8294] + - stage_id: 2 + runtime: + process: true + devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU + max_batch_size: 1 + engine_args: + model_stage: code2wav + model_arch: Qwen2_5OmniForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + gpu_memory_utilization: 0.15 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + engine_output_type: audio + max_num_batched_tokens: 4069 + engine_input_source: [1] + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 128 + seed: 42 + detokenize: True + repetition_penalty: 1.1 + +# Top-level runtime config (concise): default windows and stage edges +runtime: + enabled: true + defaults: + window_size: -1 # Simplified: trigger downstream only after full upstream completion + max_inflight: 1 # Simplified: process serially within each stage + edges: + - from: 0 # thinker → talker: trigger only after receiving full input (-1) + to: 1 + window_size: -1 + - from: 1 # talker → code2wav: trigger only after receiving full input (-1) + to: 2 + window_size: -1 diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 6ff6782ae..0c6479cce 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -23,6 +23,8 @@ from vllm_omni.entrypoints.openai.image_api_utils import ( ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + # Unit Tests diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index 240fd2051..58d4d253f 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -11,6 +11,8 @@ from unittest.mock import MagicMock import pytest from vllm.sampling_params import SamplingParams +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + @pytest.fixture def mock_comprehension_stage(): diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index da73acb5a..2db98c068 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -16,6 +16,8 @@ from vllm_omni.entrypoints.openai.serving_speech import ( ) from vllm_omni.outputs import OmniRequestOutput +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + logger = logging.getLogger(__name__) diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 6b49eba2c..ba0390040 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + from vllm_omni.entrypoints import omni as omni_module from vllm_omni.entrypoints.async_omni import AsyncOmni +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_default_stage_config_includes_cache_backend(monkeypatch): """Ensure cache_backend/cache_config are preserved in default diffusion stage.""" diff --git a/tests/entrypoints/test_omni_diffusion.py b/tests/entrypoints/test_omni_diffusion.py index c4884e3ab..ab1194a25 100644 --- a/tests/entrypoints/test_omni_diffusion.py +++ b/tests/entrypoints/test_omni_diffusion.py @@ -9,6 +9,8 @@ import pytest from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK from vllm_omni.inputs.data import OmniDiffusionSamplingParams +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + # Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies. warnings.filterwarnings( "ignore", diff --git a/tests/entrypoints/test_omni_input_preprocessor.py b/tests/entrypoints/test_omni_input_preprocessor.py index 77c84f06b..422154bf9 100644 --- a/tests/entrypoints/test_omni_input_preprocessor.py +++ b/tests/entrypoints/test_omni_input_preprocessor.py @@ -1,5 +1,9 @@ +import pytest + from vllm_omni.inputs.preprocess import OmniInputPreprocessor +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def _make_preprocessor(monkeypatch): preprocessor = object.__new__(OmniInputPreprocessor) diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index f99c6d833..4f05575ca 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/tests/entrypoints/test_omni_llm.py @@ -9,6 +9,8 @@ from vllm import SamplingParams from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + # Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies. warnings.filterwarnings( "ignore", @@ -296,9 +298,9 @@ def _setup_log_mocks(monkeypatch): # Mock OrchestratorMetrics to be a simple class that doesn't require file operations class _FakeOrchestratorMetrics: - def __init__(self, num_stages, enable_stats, wall_start_ts): + def __init__(self, num_stages, log_stats, wall_start_ts): self.num_stages = num_stages - self.enable_stats = enable_stats + self.log_stats = log_stats self.stage_first_ts = [None] * num_stages self.stage_last_ts = [None] * num_stages self.e2e_done = set() @@ -316,7 +318,7 @@ def _setup_log_mocks(monkeypatch): return "Fake summary" monkeypatch.setattr( - "vllm_omni.entrypoints.omni.OrchestratorMetrics", + "vllm_omni.entrypoints.omni.OrchestratorAggregator", _FakeOrchestratorMetrics, raising=False, ) diff --git a/tests/entrypoints/test_omni_new_request_data.py b/tests/entrypoints/test_omni_new_request_data.py index 776509d5b..b1ad56cdd 100644 --- a/tests/entrypoints/test_omni_new_request_data.py +++ b/tests/entrypoints/test_omni_new_request_data.py @@ -1,9 +1,12 @@ from types import SimpleNamespace +import pytest import torch from vllm_omni.core.sched.output import OmniNewRequestData +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_omni_new_request_data_copies_payloads(): prompt_embeds = torch.randn(2, 3) diff --git a/tests/entrypoints/test_omni_stage_diffusion_config.py b/tests/entrypoints/test_omni_stage_diffusion_config.py index 5fe04cbbd..f464c55fd 100644 --- a/tests/entrypoints/test_omni_stage_diffusion_config.py +++ b/tests/entrypoints/test_omni_stage_diffusion_config.py @@ -1,8 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest + from vllm_omni.entrypoints.omni_stage import _build_od_config +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_build_od_config_includes_diffusion_fields(): engine_args = { diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index ac503639b..ab7358d9f 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -6,6 +6,8 @@ import pytest from vllm_omni.entrypoints.stage_utils import set_stage_devices +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def _make_dummy_torch(call_log): class _Props: diff --git a/tests/examples/online_serving/test_qwen2_5_omni.py b/tests/examples/online_serving/test_qwen2_5_omni.py new file mode 100644 index 000000000..5a529662b --- /dev/null +++ b/tests/examples/online_serving/test_qwen2_5_omni.py @@ -0,0 +1,271 @@ +""" +Example online tests for Qwen2.5-Omni-7B model. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import re +import subprocess +from pathlib import Path + +import pytest + +from tests.conftest import OmniServer, convert_audio_file_to_text, cosine_similarity_text + +models = ["Qwen/Qwen2.5-Omni-7B"] + + +stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")] + +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 = [(model, stage_config) for model in models for stage_config in stage_configs] + + +@pytest.fixture(scope="module") +def omni_server(request): + """Start vLLM-Omni server as a subprocess with actual model weights. + Uses module scope so the server starts only once per test module. + Multi-stage initialization can take 10-20+ minutes. + """ + model, stage_config_path = request.param + + print(f"Starting OmniServer with model: {model}") + print("This may take 10-20+ minutes for initialization...") + + with OmniServer( + model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"], port=8091 + ) as server: + print("OmniServer started successfully") + yield server + print("OmniServer stopped") + + +def run_cmd(command): + result = subprocess.run( + command, + capture_output=True, + text=True, + ) + + if result.returncode != 0: + print(f"STDERR: {result.stderr}") + raise subprocess.CalledProcessError(result.returncode, command) + + all_output = result.stdout + print(f"All output:\n{all_output}") + return all_output + + +def extract_content_after_keyword(keywords, text): + matches = re.findall(rf"{keywords}\s*(.+)", text, re.DOTALL) + + if not matches: + raise AssertionError(f"Keywords {keywords} not found in provided text output") + return matches[0] + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_send_multimodal_request_001(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "mixed_modalities", + ] + + result = run_cmd(command) + + text_content = extract_content_after_keyword("Chat completion output from text:", result) + + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"text content is: {text_content}") + print(f"audio content is: {audio_content}") + + assert all(keyword in text_content for keyword in ["baby", "book"]), ( + "The output does not contain any of the keywords in video description." + ) + # There is currently an issue with incorrect image descriptions. + # assert "cherry blossom" in text_content, "The output does not contain any of the keywords in image description." + assert "lamb" in text_content, "The output does not contain any of the keywords in audio description." + + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + print(f"similarity is: {similarity}") + assert similarity > 0.9, "The audio content is not same as the text" + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_send_multimodal_request_002(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "mixed_modalities", + "--prompt", + "Analyze all the media content and provide a comprehensive summary.", + ] + result = run_cmd(command) + + text_content = extract_content_after_keyword("Chat completion output from text:", result) + + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"text content is: {text_content}") + assert all(keyword in text_content for keyword in ["baby", "book"]), ( + "The output does not contain any of the keywords in video description." + ) + # There is currently an issue with incorrect image descriptions. + # assert "cherry blossom" in text_content, "The output does not contain any of the keywords in image description." + assert "lamb" in text_content, "The output does not contain any of the keywords in audio description." + + print(f"audio content is: {audio_content}") + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + print(f"similarity is: {similarity}") + assert similarity > 0.9, "The audio content is not same as the text" + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_send_multimodal_request_003(omni_server) -> None: + command = ["bash", os.path.join(example_dir, "run_curl_multimodal_generation.sh"), "mixed_modalities"] + + result = run_cmd(command) + + text_content = extract_content_after_keyword("Output of request:", result) + + # Verify text output same as audio output + print(f"text content is: {text_content}") + assert all(keyword in text_content for keyword in ["baby", "book"]), ( + "The output does not contain any of the keywords in video description." + ) + # There is currently an issue with incorrect image descriptions. + # assert "cherry blossom" in text_content, "The output does not contain any of the keywords in image description." + assert "lamb" in text_content, "The output does not contain any of the keywords in audio description." + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_modality_control_001(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "mixed_modalities", + "--modalities", + "text", + ] + + result = run_cmd(command) + + text_content = extract_content_after_keyword("Chat completion output from text:", result) + + # Verify text output + print(f"text content is: {text_content}") + assert all(keyword in text_content for keyword in ["baby", "book"]), ( + "The output does not contain any of the keywords in video description." + ) + # There is currently an issue with incorrect image descriptions. + # assert "cherry blossom" in text_content, "The output does not contain any of the keywords in image description." + assert "lamb" in text_content, "The output does not contain any of the keywords in audio description." + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_modality_control_002(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "mixed_modalities", + "--modalities", + "audio", + ] + + run_cmd(command) + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"audio content is: {audio_content}") + assert all(keyword in audio_content for keyword in ["baby", "book"]), ( + "The output does not contain any of the keywords in video description." + ) + # There is currently an issue with incorrect image descriptions. + # assert "cherry blossom" in audio_content, "The output does not contain any of the keywords in image description." + assert "lamb" in audio_content, "The output does not contain any of the keywords in audio description." + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_modality_control_003(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "mixed_modalities", + "--modalities", + "audio,text", + ] + + result = run_cmd(command) + + text_content = extract_content_after_keyword("Chat completion output from text:", result) + + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"text content is: {text_content}") + assert all(keyword in text_content for keyword in ["baby", "book"]), ( + "The output does not contain any of the keywords in video description." + ) + # There is currently an issue with incorrect image descriptions. + # assert "cherry blossom" in text_content, "The output does not contain any of the keywords in image description." + assert "lamb" in text_content, "The output does not contain any of the keywords in audio description." + + print(f"audio content is: {audio_content}") + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + print(f"similarity is: {similarity}") + assert similarity > 0.9, "The audio content is not same as the text" + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.skip(reason="There is a known issue with stream error.") +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_stream_001(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "mixed_modalities", + "--stream", + ] + + result = run_cmd(command) + + text_content = extract_content_after_keyword("content:", result) + + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"text content is: {text_content}") + assert all(keyword in text_content for keyword in ["baby", "book"]), ( + "The output does not contain any of the keywords in video description." + ) + # There is currently an issue with incorrect image descriptions. + # assert "cherry blossom" in text_content, "The output does not contain any of the keywords in image description." + assert "lamb" in text_content, "The output does not contain any of the keywords in audio description." + + print(f"audio content is: {audio_content}") + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + print(f"similarity is: {similarity}") + assert similarity > 0.9, "The audio content is not same as the text" + # TODO: Verify the E2E latency after confirmation baseline. + + +# TODO: test local web ui diff --git a/tests/examples/online_serving/test_qwen3_omni.py b/tests/examples/online_serving/test_qwen3_omni.py new file mode 100644 index 000000000..675146ff5 --- /dev/null +++ b/tests/examples/online_serving/test_qwen3_omni.py @@ -0,0 +1,230 @@ +""" +Example Online tests for Qwen3-Omni model. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import re +import subprocess +from pathlib import Path + +import pytest + +from tests.conftest import OmniServer, convert_audio_file_to_text, cosine_similarity_text + +models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] + + +stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "qwen3_omni_ci.yaml")] + +example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen3_omni") +# Create parameter combinations for model and stage config +test_params = [(model, stage_config) for model in models for stage_config in stage_configs] + + +@pytest.fixture(scope="module") +def omni_server(request): + """Start vLLM-Omni server as a subprocess with actual model weights. + Uses module scope so the server starts only once per test module.. + Multi-stage initialization can take 10-20+ minutes. + """ + model, stage_config_path = request.param + + print(f"Starting OmniServer with model: {model}") + print("This may take 10-20+ minutes for initialization...") + + with OmniServer( + model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"], port=8091 + ) as server: + print("OmniServer started successfully") + yield server + print("OmniServer stopped") + + +def run_cmd(command): + result = subprocess.run( + command, + capture_output=True, + text=True, + ) + + if result.returncode != 0: + print(f"STDERR: {result.stderr}") + raise subprocess.CalledProcessError(result.returncode, command) + + all_output = result.stdout + print(f"All output:\n{all_output}") + return all_output + + +def extract_content_after_keyword(keywords, text): + matches = re.findall(rf"{keywords}\s*(.+)", text, re.DOTALL) + + if not matches: + raise AssertionError(f"Keywords {keywords} not found in provided text output") + return matches[0] + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_send_multimodal_request_001(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "use_image", + ] + + result = run_cmd(command) + text_content = extract_content_after_keyword("Chat completion output from text:", result) + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"text content is: {text_content}") + print(f"audio content is: {audio_content}") + + assert "cherry blossom" in text_content, "The output does not contain any of the keywords." + + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + print(f"similarity is: {similarity}") + assert similarity > 0.9, "The audio content is not same as the text" + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_send_multimodal_request_002(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "use_video", + "--prompt", + "What are the main activities shown in this video?", + ] + result = run_cmd(command) + + text_content = extract_content_after_keyword("Chat completion output from text:", result) + + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"text content is: {text_content}") + print(f"audio content is: {audio_content}") + assert all(keyword in text_content for keyword in ["baby", "book"]), ( + "The output does not contain any of the keywords." + ) + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + print(f"similarity is: {similarity}") + assert similarity > 0.9, "The audio content is not same as the text" + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_send_multimodal_request_003(omni_server) -> None: + command = ["bash", os.path.join(example_dir, "run_curl_multimodal_generation.sh"), "use_image"] + + result = run_cmd(command) + + text_content = extract_content_after_keyword("Output of request:", result) + + # Verify text output same as audio output + print(f"text content is: {text_content}") + assert "cherry blossom" in text_content, "The output does not contain any of the keywords." + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_modality_control_001(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "use_image", + "--modalities", + "text", + ] + + result = run_cmd(command) + + text_content = extract_content_after_keyword("Chat completion output from text:", result) + + # Verify text output + print(f"text content is: {text_content}") + assert "cherry blossom" in text_content, "The output does not contain any of the keywords." + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_modality_control_002(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "use_image", + "--modalities", + "audio", + ] + + run_cmd(command) + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"audio content is: {audio_content}") + assert "cherry blossom" in audio_content, "The output does not contain any of the keywords." + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_modality_control_003(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "use_image", + "--modalities", + "audio,text", + ] + + result = run_cmd(command) + + text_content = extract_content_after_keyword("Chat completion output from text:", result) + + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"text content is: {text_content}") + assert "cherry blossom" in audio_content, "The output does not contain any of the keywords." + print(f"audio content is: {audio_content}") + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + print(f"similarity is: {similarity}") + assert similarity > 0.9, "The audio content is not same as the text" + + # TODO: Verify the E2E latency after confirmation baseline. + + +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_stream_001(omni_server) -> None: + command = [ + "python", + os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py"), + "--query-type", + "use_image", + "--stream", + ] + + result = run_cmd(command) + + text_content = extract_content_after_keyword("content:", result) + + # Verify text output same as audio output + audio_content = convert_audio_file_to_text(output_path="./audio_0.wav") + print(f"text content is: {text_content}") + assert "cherry blossom" in audio_content, "The output does not contain any of the keywords." + print(f"audio content is: {audio_content}") + similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) + print(f"similarity is: {similarity}") + assert similarity > 0.9, "The audio content is not same as the text" + # TODO: Verify the E2E latency after confirmation baseline. + + +# TODO: test local web ui diff --git a/tests/metrics/test_stats.py b/tests/metrics/test_stats.py new file mode 100644 index 000000000..8b18c5dff --- /dev/null +++ b/tests/metrics/test_stats.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +import pytest + +from vllm_omni.metrics import OrchestratorAggregator +from vllm_omni.metrics.stats import RequestE2EStats, StageRequestStats, StageStats + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _get_request_entry(table: list[dict], request_id: str) -> dict: + for entry in table: + if entry.get("request_id") == request_id: + return entry + raise AssertionError(f"request_id={request_id} not found") + + +def test_orchestrator_aggregator_builds_summary() -> None: + agg = OrchestratorAggregator(num_stages=2, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=1) + agg.stage_first_ts[0] = 0.0 + agg.stage_last_ts[0] = 0.03 + agg.stage_first_ts[1] = 0.05 + agg.stage_last_ts[1] = 0.07 + + agg.on_forward(0, 1, "r1", size_bytes=1024, tx_ms=5.0, used_shm=False) + + agg.on_stage_metrics( + 0, + "r1", + StageRequestStats( + batch_id=1, + batch_size=1, + num_tokens_in=3, + num_tokens_out=3, + stage_gen_time_ms=30.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats(), + ), + ) + agg.on_stage_metrics( + 1, + "r1", + StageRequestStats( + batch_id=1, + batch_size=1, + num_tokens_in=0, + num_tokens_out=4, + stage_gen_time_ms=20.0, + rx_transfer_bytes=1024, + rx_decode_time_ms=5.0, + rx_in_flight_time_ms=2.0, + stage_stats=StageStats(), + ), + ) + agg.on_finalize_request(1, "r1", req_start_ts=0.0) + + summary = agg.build_and_log_summary() + overall = summary["overall_summary"] + assert overall["e2e_requests"] == 1 + + stage_entry = _get_request_entry(summary["stage_table"], "r1") + stage_ids = [row["stage_id"] for row in stage_entry["stages"]] + assert stage_ids == [0, 1] + + transfer_entry = _get_request_entry(summary["trans_table"], "r1") + assert transfer_entry["transfers"][0]["edge"] == "0->1" + assert transfer_entry["transfers"][0]["size_kbytes"] == 1.0 + + e2e_entry = _get_request_entry(summary["e2e_table"], "r1") + assert e2e_entry["e2e_total_tokens"] == 10 + + +def test_build_and_log_summary_e2e_only() -> None: + agg = OrchestratorAggregator(num_stages=1, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e=0) + agg.e2e_events.append( + RequestE2EStats( + request_id="r", + e2e_total_ms=10.0, + e2e_total_tokens=5, + transfers_total_time_ms=0.0, + transfers_total_bytes=0, + ) + ) + + summary = agg.build_and_log_summary() + e2e_entry = _get_request_entry(summary["e2e_table"], "r") + assert e2e_entry["e2e_total_tokens"] == 5 + stage_entry = _get_request_entry(summary["stage_table"], "r") + assert stage_entry["stages"] == [] + + +def test_build_and_log_summary_multiple_requests() -> None: + agg = OrchestratorAggregator( + num_stages=2, log_stats=True, wall_start_ts=0.0, final_stage_id_for_e2e={"r1": 1, "r2": 0} + ) + + # Request r1 goes through both stages + agg.on_stage_metrics( + 0, + "r1", + StageRequestStats( + batch_id=1, + batch_size=1, + num_tokens_in=2, + num_tokens_out=4, + stage_gen_time_ms=10.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats(), + ), + ) + agg.on_stage_metrics( + 1, + "r1", + StageRequestStats( + batch_id=1, + batch_size=1, + num_tokens_in=4, + num_tokens_out=5, + stage_gen_time_ms=8.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats(), + ), + ) + agg.on_finalize_request(1, "r1", req_start_ts=0.0) + + # Request r2 only goes through stage 0 + agg.on_stage_metrics( + 0, + "r2", + StageRequestStats( + batch_id=2, + batch_size=1, + num_tokens_in=1, + num_tokens_out=2, + stage_gen_time_ms=12.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats(), + ), + ) + agg.on_finalize_request(0, "r2", req_start_ts=0.0) + + summary = agg.build_and_log_summary() + assert len(summary["stage_table"]) == 2 + assert {entry["request_id"] for entry in summary["e2e_table"]} == {"r1", "r2"} + # Check that r1 has two stages and r2 has one + r1_stage_entry = next(e for e in summary["stage_table"] if e["request_id"] == "r1") + r2_stage_entry = next(e for e in summary["stage_table"] if e["request_id"] == "r2") + assert len(r1_stage_entry["stages"]) == 2 + assert len(r2_stage_entry["stages"]) == 1 diff --git a/tests/model_executor/models/qwen2_5_omni/test_audio_length.py b/tests/model_executor/models/qwen2_5_omni/test_audio_length.py index dd5f09817..6156ed3a9 100644 --- a/tests/model_executor/models/qwen2_5_omni/test_audio_length.py +++ b/tests/model_executor/models/qwen2_5_omni/test_audio_length.py @@ -3,6 +3,8 @@ import pytest +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + def test_resolve_max_mel_frames_default(): from vllm_omni.model_executor.models.qwen2_5_omni.audio_length import resolve_max_mel_frames diff --git a/tests/utils.py b/tests/utils.py index f5c513a4d..3219821b6 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -495,18 +495,8 @@ def hardware_test(*, res: dict[str, str], num_cards: int | dict[str, int] = 1): raise ValueError(f"Unsupported platform: {platform}") all_marks.extend(marks) - create_new_process_flag = False - for cards in num_cards_dict.values(): - if cards > 1: - create_new_process_flag = True - break - def wrapper(f: Callable[_P, None]) -> Callable[_P, None]: - if create_new_process_flag: - # only for distributed tests - func = create_new_process_for_each_test()(f) - else: - func = f + func = f for mark in reversed(all_marks): func = mark(func) return func diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index f08df8754..8ecb4b0f8 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -1,7 +1,10 @@ +import pytest import torch from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class _DummyInputBatch: def __init__(self): diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index eb1adf227..c7836123a 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -1,10 +1,13 @@ from contextlib import contextmanager from types import SimpleNamespace +import pytest import torch from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + class DummyBuffer: """A minimal buffer wrapper that exposes the `.gpu` attribute.""" diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 83cfe84d7..f83bb294d 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -14,10 +14,11 @@ from vllm.logger import init_logger from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy +from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention from vllm_omni.diffusion.attention.selector import get_attn_backend from vllm_omni.diffusion.distributed.parallel_state import get_sp_group -from vllm_omni.diffusion.forward_context import get_forward_context +from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available logger = init_logger(__name__) @@ -87,6 +88,21 @@ class Attention(nn.Module): gather_idx=gather_idx, use_sync=use_sync, ) + # Fallback strategy when SP is not active (outside sharded regions) + self._no_parallel_strategy = NoParallelAttention() + + def _get_active_parallel_strategy(self): + """Get the parallel strategy based on current SP active state. + + Returns NoParallelAttention if we're outside an SP sharded region + (e.g., in noise_refiner/context_refiner before unified_prepare in Z-Image). + This avoids unnecessary SP communication for layers not covered by _sp_plan. + """ + if is_forward_context_available(): + ctx = get_forward_context() + if not ctx.sp_active: + return self._no_parallel_strategy + return self.parallel_strategy def forward( self, @@ -95,20 +111,23 @@ class Attention(nn.Module): value: torch.Tensor, attn_metadata: AttentionMetadata = None, ) -> torch.Tensor: + # Get the appropriate parallel strategy based on SP active state + strategy = self._get_active_parallel_strategy() + # 1. Prepare inputs (Communication / Resharding) # For Ulysses: AllToAll Q/K/V; Slicing joint_q/k/v # For Ring: Concat joint_q - query, key, value, attn_metadata, ctx = self.parallel_strategy.pre_attention(query, key, value, attn_metadata) + query, key, value, attn_metadata, ctx = strategy.pre_attention(query, key, value, attn_metadata) # 2. Kernel Execution (Computation) - if self.use_ring: + if self.use_ring and strategy is not self._no_parallel_strategy: out = self._run_ring_attention(query, key, value, attn_metadata) else: out = self._run_local_attention(query, key, value, attn_metadata) # 3. Post-processing (Reverse Communication) # For Ulysses: AllToAll Output, and AllGather Joint Output - out = self.parallel_strategy.post_attention(out, ctx) + out = strategy.post_attention(out, ctx) return out diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 8e19f1242..2ab94d5b7 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -102,16 +102,20 @@ class DiffusionEngine: if not isinstance(outputs, list): outputs = [outputs] if outputs is not None else [] + metrics = { + "image_num": int(request.sampling_params.num_outputs_per_prompt), + "resolution": int(request.sampling_params.resolution), + "postprocess_time_ms": postprocess_time * 1000, + } + if self.pre_process_func is not None: + metrics["preprocessing_time_ms"] = preprocess_time * 1000 + # Handle single request or multiple requests if len(request.prompts) == 1: # Single request: return single OmniRequestOutput prompt = request.prompts[0] request_id = request.request_ids[0] if request.request_ids else "" - metrics = {} - if output.trajectory_timesteps is not None: - metrics["trajectory_timesteps"] = output.trajectory_timesteps - if supports_audio_output(self.od_config.model_class_name): audio_payload = outputs[0] if len(outputs) == 1 else outputs return [ @@ -149,10 +153,6 @@ class DiffusionEngine: request_outputs = outputs[output_idx : output_idx + num_outputs] if output_idx < len(outputs) else [] output_idx += num_outputs - metrics = {} - if output.trajectory_timesteps is not None: - metrics["trajectory_timesteps"] = output.trajectory_timesteps - if supports_audio_output(self.od_config.model_class_name): audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs results.append( diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 9fa0a5bfc..12da9c43b 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -757,6 +757,7 @@ def initialize_model_parallel( backend=backend, parallel_mode="data", ) + vllm_parallel_state._DP = _DP global _CFG assert _CFG is None, "classifier_free_guidance group is already initialized" @@ -774,6 +775,7 @@ def initialize_model_parallel( backend=backend, parallel_mode="pipeline", ) + vllm_parallel_state._PP = _PP global _SP assert _SP is None, "sequence parallel group is already initialized" diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index 02082e7b1..d2e68c8be 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -31,6 +31,16 @@ class ForwardContext: # Original sequence length before padding (for removing padding in gather) sp_original_seq_len: int | None = None + # SP active scope tracking + # Tracks the depth of SP sharding - incremented on shard, decremented on gather + # Used by attention layers to determine if SP communication should be enabled + _sp_shard_depth: int = 0 + + @property + def sp_active(self) -> bool: + """Returns True when inside an SP sharded region (between shard and gather).""" + return self._sp_shard_depth > 0 + def __post_init__(self): pass diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index 8f2c2b755..7044a3aa6 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/vllm_omni/diffusion/hooks/sequence_parallel.py @@ -235,6 +235,8 @@ class SequenceParallelSplitHook(ModelHook): def post_forward(self, module: nn.Module, output: Any) -> Any: """Shard outputs for split_output=True entries.""" + from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + is_tensor = isinstance(output, torch.Tensor) is_tensor_list = isinstance(output, (list, tuple)) and all(isinstance(x, torch.Tensor) for x in output) @@ -243,6 +245,7 @@ class SequenceParallelSplitHook(ModelHook): return output output_list = [output] if is_tensor else list(output) + actually_sharded = False for index, spm in self.metadata.items(): if not isinstance(index, int): @@ -252,7 +255,14 @@ class SequenceParallelSplitHook(ModelHook): if index >= len(output_list): raise ValueError(f"Index {index} out of bounds for output of length {len(output_list)}.") - output_list[index] = self._prepare_sp_input(output_list[index], spm, self._last_args, self._last_kwargs) + original = output_list[index] + output_list[index] = self._prepare_sp_input(original, spm, self._last_args, self._last_kwargs) + if output_list[index] is not original: + actually_sharded = True + + # Mark SP as active only if at least one tensor was actually sharded + if actually_sharded and is_forward_context_available(): + get_forward_context()._sp_shard_depth += 1 return output_list[0] if is_tensor else type(output)(output_list) @@ -445,6 +455,8 @@ class SequenceParallelGatherHook(ModelHook): ctx = get_forward_context() original_seq_len = ctx.sp_original_seq_len + actually_gathered = False + for i, spm in enumerate(self.metadata): if spm is None: continue @@ -465,6 +477,12 @@ class SequenceParallelGatherHook(ModelHook): logger.debug(f"Removed padding: gathered shape {gathered.shape} (original_seq_len={original_seq_len})") output[i] = gathered + actually_gathered = True + + # Mark SP as inactive only if at least one tensor was actually gathered + if actually_gathered and is_forward_context_available(): + ctx = get_forward_context() + ctx._sp_shard_depth = max(0, ctx._sp_shard_depth - 1) return output[0] if is_tensor else type(output)(output) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py b/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py new file mode 100644 index 000000000..98a3ac07b --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Hunyuan Image 3 diffusion model components.""" + +from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_image_3_transformer import ( + HunyuanImage3Model, + HunyuanImage3Text2ImagePipeline, +) +from vllm_omni.diffusion.models.hunyuan_image_3.pipeline_hunyuan_image_3 import ( + HunyuanImage3Pipeline, +) + +__all__ = [ + "HunyuanImage3Pipeline", + "HunyuanImage3Model", + "HunyuanImage3Text2ImagePipeline", +] diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/autoencoder.py b/vllm_omni/diffusion/models/hunyuan_image_3/autoencoder.py new file mode 100644 index 000000000..efba2f274 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/autoencoder.py @@ -0,0 +1,668 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils import BaseOutput +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange +from torch import Tensor, nn + + +class DiagonalGaussianDistribution: + def __init__(self, parameters: torch.Tensor, deterministic: bool = False): + if parameters.ndim == 3: + dim = 2 # (B, L, C) + elif parameters.ndim == 5 or parameters.ndim == 4: + dim = 1 # (B, C, T, H ,W) / (B, C, H, W) + else: + raise NotImplementedError + self.parameters = parameters + self.mean, self.logvar = torch.chunk(parameters, 2, dim=dim) + self.logvar = torch.clamp(self.logvar, -30.0, 20.0) + self.deterministic = deterministic + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + if self.deterministic: + self.var = self.std = torch.zeros_like( + self.mean, device=self.parameters.device, dtype=self.parameters.dtype + ) + + def sample(self, generator: torch.Generator | None = None) -> torch.FloatTensor: + # make sure sample is on the same device as the parameters and has same dtype + sample = randn_tensor( + self.mean.shape, + generator=generator, + device=self.parameters.device, + dtype=self.parameters.dtype, + ) + x = self.mean + self.std * sample + return x + + +@dataclass +class DecoderOutput(BaseOutput): + sample: torch.FloatTensor + posterior: DiagonalGaussianDistribution | None = None + + +def swish(x: Tensor) -> Tensor: + return x * torch.sigmoid(x) + + +def forward_with_checkpointing(module, *inputs, use_checkpointing=False): + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if use_checkpointing: + return torch.utils.checkpoint.checkpoint(create_custom_forward(module), *inputs, use_reentrant=False) + else: + return module(*inputs) + + +class Conv3d(nn.Conv3d): + """ + Perform Conv3d on patches with numerical differences from nn.Conv3d within 1e-5. + Only symmetric padding is supported. + """ + + def forward(self, input): + B, C, T, H, W = input.shape + memory_count = (C * T * H * W) * 2 / 1024**3 + if memory_count > 2: + n_split = math.ceil(memory_count / 2) + assert n_split >= 2 + chunks = torch.chunk(input, chunks=n_split, dim=-3) + padded_chunks = [] + for i in range(len(chunks)): + if self.padding[0] > 0: + padded_chunk = F.pad( + chunks[i], + (0, 0, 0, 0, self.padding[0], self.padding[0]), + mode="constant" if self.padding_mode == "zeros" else self.padding_mode, + value=0, + ) + if i > 0: + padded_chunk[:, :, : self.padding[0]] = chunks[i - 1][:, :, -self.padding[0] :] + if i < len(chunks) - 1: + padded_chunk[:, :, -self.padding[0] :] = chunks[i + 1][:, :, : self.padding[0]] + else: + padded_chunk = chunks[i] + padded_chunks.append(padded_chunk) + padding_bak = self.padding + self.padding = (0, self.padding[1], self.padding[2]) + outputs = [] + for i in range(len(padded_chunks)): + outputs.append(super().forward(padded_chunks[i])) + self.padding = padding_bak + return torch.cat(outputs, dim=-3) + else: + return super().forward(input) + + +class AttnBlock(nn.Module): + """Attention with torch sdpa implementation.""" + + def __init__(self, in_channels: int): + super().__init__() + self.in_channels = in_channels + + self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + + self.q = Conv3d(in_channels, in_channels, kernel_size=1) + self.k = Conv3d(in_channels, in_channels, kernel_size=1) + self.v = Conv3d(in_channels, in_channels, kernel_size=1) + self.proj_out = Conv3d(in_channels, in_channels, kernel_size=1) + + def attention(self, h_: Tensor) -> Tensor: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, c, f, h, w = q.shape + q = rearrange(q, "b c f h w -> b 1 (f h w) c").contiguous() + k = rearrange(k, "b c f h w -> b 1 (f h w) c").contiguous() + v = rearrange(v, "b c f h w -> b 1 (f h w) c").contiguous() + h_ = nn.functional.scaled_dot_product_attention(q, k, v) + + return rearrange(h_, "b 1 (f h w) c -> b c f h w", f=f, h=h, w=w, c=c, b=b) + + def forward(self, x: Tensor) -> Tensor: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True) + self.conv1 = Conv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True) + self.conv2 = Conv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + self.nin_shortcut = Conv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + h = x + h = self.norm1(h) + h = swish(h) + h = self.conv1(h) + + h = self.norm2(h) + h = swish(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + return x + h + + +class DownsampleDCAE(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_downsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_downsample else 1 * 2 * 2 + assert out_channels % factor == 0 + self.conv = Conv3d(in_channels, out_channels // factor, kernel_size=3, stride=1, padding=1) + + self.add_temporal_downsample = add_temporal_downsample + self.group_size = factor * in_channels // out_channels + + def forward(self, x: Tensor): + r1 = 2 if self.add_temporal_downsample else 1 + h = self.conv(x) + h = rearrange(h, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + shortcut = rearrange(x, "b c (f r1) (h r2) (w r3) -> b (r1 r2 r3 c) f h w", r1=r1, r2=2, r3=2) + + B, C, T, H, W = shortcut.shape + shortcut = shortcut.view(B, h.shape[1], self.group_size, T, H, W).mean(dim=2) + return h + shortcut + + +class UpsampleDCAE(nn.Module): + def __init__(self, in_channels: int, out_channels: int, add_temporal_upsample: bool = True): + super().__init__() + factor = 2 * 2 * 2 if add_temporal_upsample else 1 * 2 * 2 + self.conv = Conv3d(in_channels, out_channels * factor, kernel_size=3, stride=1, padding=1) + + self.add_temporal_upsample = add_temporal_upsample + self.repeats = factor * out_channels // in_channels + + def forward(self, x: Tensor): + r1 = 2 if self.add_temporal_upsample else 1 + h = self.conv(x) + h = rearrange(h, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + shortcut = x.repeat_interleave(repeats=self.repeats, dim=1) + shortcut = rearrange(shortcut, "b (r1 r2 r3 c) f h w -> b c (f r1) (h r2) (w r3)", r1=r1, r2=2, r3=2) + return h + shortcut + + +class Encoder(nn.Module): + """ + The encoder network of AutoencoderKLConv3D. + """ + + def __init__( + self, + in_channels: int, + z_channels: int, + block_out_channels: tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + downsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[-1] % (2 * z_channels) == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + # downsampling + self.conv_in = Conv3d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1) + + self.down = nn.ModuleList() + block_in = block_out_channels[0] + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + down = nn.Module() + down.block = block + + add_spatial_downsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_downsample = add_spatial_downsample and bool( + i_level >= np.log2(ffactor_spatial // ffactor_temporal) + ) + if add_spatial_downsample or add_temporal_downsample: + assert i_level < len(block_out_channels) - 1 + block_out = block_out_channels[i_level + 1] if downsample_match_channel else block_in + down.downsample = DownsampleDCAE(block_in, block_out, add_temporal_downsample) + block_in = block_out + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = Conv3d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1) + + self.gradient_checkpointing = False + + def forward(self, x: Tensor) -> Tensor: + use_checkpointing = bool(self.training and self.gradient_checkpointing) + + # downsampling + h = self.conv_in(x) + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks): + h = forward_with_checkpointing( + self.down[i_level].block[i_block], h, use_checkpointing=use_checkpointing + ) + if hasattr(self.down[i_level], "downsample"): + h = forward_with_checkpointing(self.down[i_level].downsample, h, use_checkpointing=use_checkpointing) + + # middle + h = forward_with_checkpointing(self.mid.block_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.block_2, h, use_checkpointing=use_checkpointing) + + # end + group_size = self.block_out_channels[-1] // (2 * self.z_channels) + shortcut = rearrange(h, "b (c r) f h w -> b c r f h w", r=group_size).mean(dim=2) + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + h += shortcut + return h + + +class Decoder(nn.Module): + """ + The decoder network of AutoencoderKLConv3D. + """ + + def __init__( + self, + z_channels: int, + out_channels: int, + block_out_channels: tuple[int, ...], + num_res_blocks: int, + ffactor_spatial: int, + ffactor_temporal: int, + upsample_match_channel: bool = True, + ): + super().__init__() + assert block_out_channels[0] % z_channels == 0 + + self.z_channels = z_channels + self.block_out_channels = block_out_channels + self.num_res_blocks = num_res_blocks + + # z to block_in + block_in = block_out_channels[0] + self.conv_in = Conv3d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in) + + # upsampling + self.up = nn.ModuleList() + for i_level, ch in enumerate(block_out_channels): + block = nn.ModuleList() + block_out = ch + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out)) + block_in = block_out + up = nn.Module() + up.block = block + + add_spatial_upsample = bool(i_level < np.log2(ffactor_spatial)) + add_temporal_upsample = bool(i_level < np.log2(ffactor_temporal)) + if add_spatial_upsample or add_temporal_upsample: + assert i_level < len(block_out_channels) - 1 + block_out = block_out_channels[i_level + 1] if upsample_match_channel else block_in + up.upsample = UpsampleDCAE(block_in, block_out, add_temporal_upsample) + block_in = block_out + self.up.append(up) + + # end + self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True) + self.conv_out = Conv3d(block_in, out_channels, kernel_size=3, stride=1, padding=1) + + self.gradient_checkpointing = False + + def forward(self, z: Tensor) -> Tensor: + use_checkpointing = bool(self.training and self.gradient_checkpointing) + + # z to block_in + repeats = self.block_out_channels[0] // (self.z_channels) + h = self.conv_in(z) + z.repeat_interleave(repeats=repeats, dim=1) + + # middle + h = forward_with_checkpointing(self.mid.block_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.attn_1, h, use_checkpointing=use_checkpointing) + h = forward_with_checkpointing(self.mid.block_2, h, use_checkpointing=use_checkpointing) + + # upsampling + for i_level in range(len(self.block_out_channels)): + for i_block in range(self.num_res_blocks + 1): + h = forward_with_checkpointing(self.up[i_level].block[i_block], h, use_checkpointing=use_checkpointing) + if hasattr(self.up[i_level], "upsample"): + h = forward_with_checkpointing(self.up[i_level].upsample, h, use_checkpointing=use_checkpointing) + + # end + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h + + +class AutoencoderKLConv3D(ModelMixin, ConfigMixin): + """ + Autoencoder model with KL-regularized latent space based on 3D convolutions. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int, + out_channels: int, + latent_channels: int, + block_out_channels: tuple[int, ...], + layers_per_block: int, + ffactor_spatial: int, + ffactor_temporal: int, + sample_size: int, + sample_tsize: int, + scaling_factor: float | None = None, + shift_factor: float | None = None, + downsample_match_channel: bool = True, + upsample_match_channel: bool = True, + only_encoder: bool = False, # only build encoder for saving memory + only_decoder: bool = False, # only build decoder for saving memory + ): + super().__init__() + self.ffactor_spatial = ffactor_spatial + self.ffactor_temporal = ffactor_temporal + self.scaling_factor = scaling_factor + self.shift_factor = shift_factor + + # build model + if not only_decoder: + self.encoder = Encoder( + in_channels=in_channels, + z_channels=latent_channels, + block_out_channels=block_out_channels, + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + downsample_match_channel=downsample_match_channel, + ) + if not only_encoder: + self.decoder = Decoder( + z_channels=latent_channels, + out_channels=out_channels, + block_out_channels=list(reversed(block_out_channels)), + num_res_blocks=layers_per_block, + ffactor_spatial=ffactor_spatial, + ffactor_temporal=ffactor_temporal, + upsample_match_channel=upsample_match_channel, + ) + + # slicing and tiling related + self.use_slicing = False + self.slicing_bsz = 1 + self.use_spatial_tiling = False + self.use_temporal_tiling = False + self.use_tiling_during_training = False + + # only relevant if vae tiling is enabled + self.tile_sample_min_size = sample_size + self.tile_latent_min_size = sample_size // ffactor_spatial + self.tile_sample_min_tsize = sample_tsize + self.tile_latent_min_tsize = sample_tsize // ffactor_temporal + self.tile_overlap_factor = 0.25 + + # use torch.compile for faster encode speed + self.use_compile = False + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, (Encoder, Decoder)): + module.gradient_checkpointing = value + + def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + blend_extent = min(a.shape[-1], b.shape[-1], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + blend_extent = min(a.shape[-2], b.shape[-2], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int): + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + def spatial_tiled_encode(self, x: torch.Tensor): + """spatial tailing for frames""" + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor)) # 256 * (1 - 0.25) = 192 + blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor) # 8 * 0.25 = 2 + row_limit = self.tile_latent_min_size - blend_extent # 8 - 2 = 6 + + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size] + tile = self.encoder(tile) + row.append(tile) + rows.append(row) + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + moments = torch.cat(result_rows, dim=-2) + return moments + + def temporal_tiled_encode(self, x: torch.Tensor): + """temporal tailing for frames""" + B, C, T, H, W = x.shape + overlap_size = int(self.tile_sample_min_tsize * (1 - self.tile_overlap_factor)) # 64 * (1 - 0.25) = 48 + blend_extent = int(self.tile_latent_min_tsize * self.tile_overlap_factor) # 8 * 0.25 = 2 + t_limit = self.tile_latent_min_tsize - blend_extent # 8 - 2 = 6 + + row = [] + for i in range(0, T, overlap_size): + tile = x[:, :, i : i + self.tile_sample_min_tsize, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_sample_min_size or tile.shape[-2] > self.tile_sample_min_size + ): + tile = self.spatial_tiled_encode(tile) + else: + tile = self.encoder(tile) + row.append(tile) + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + moments = torch.cat(result_row, dim=-3) + return moments + + def spatial_tiled_decode(self, z: torch.Tensor): + """spatial tailing for frames""" + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor)) # 8 * (1 - 0.25) = 6 + blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor) # 256 * 0.25 = 64 + row_limit = self.tile_sample_min_size - blend_extent # 256 - 64 = 192 + + rows = [] + for i in range(0, H, overlap_size): + row = [] + for j in range(0, W, overlap_size): + tile = z[:, :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size] + decoded = self.decoder(tile) + row.append(decoded) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = self.blend_v(rows[i - 1][j], tile, blend_extent) + if j > 0: + tile = self.blend_h(row[j - 1], tile, blend_extent) + result_row.append(tile[:, :, :, :row_limit, :row_limit]) + result_rows.append(torch.cat(result_row, dim=-1)) + dec = torch.cat(result_rows, dim=-2) + return dec + + def temporal_tiled_decode(self, z: torch.Tensor): + """temporal tailing for frames""" + B, C, T, H, W = z.shape + overlap_size = int(self.tile_latent_min_tsize * (1 - self.tile_overlap_factor)) # 8 * (1 - 0.25) = 6 + blend_extent = int(self.tile_sample_min_tsize * self.tile_overlap_factor) # 64 * 0.25 = 16 + t_limit = self.tile_sample_min_tsize - blend_extent # 64 - 16 = 48 + assert 0 < overlap_size < self.tile_latent_min_tsize + + row = [] + for i in range(0, T, overlap_size): + tile = z[:, :, i : i + self.tile_latent_min_tsize, :, :] + if self.use_spatial_tiling and ( + tile.shape[-1] > self.tile_latent_min_size or tile.shape[-2] > self.tile_latent_min_size + ): + decoded = self.spatial_tiled_decode(tile) + else: + decoded = self.decoder(tile) + row.append(decoded) + + result_row = [] + for i, tile in enumerate(row): + if i > 0: + tile = self.blend_t(row[i - 1], tile, blend_extent) + result_row.append(tile[:, :, :t_limit, :, :]) + dec = torch.cat(result_row, dim=-3) + return dec + + def encode(self, x: Tensor, return_dict: bool = True): + """ + Encodes the input by passing through the encoder network. + Support slicing and tiling for memory efficiency. + """ + + def _encode(x): + if self.use_temporal_tiling and x.shape[-3] > self.tile_sample_min_tsize: + return self.temporal_tiled_encode(x) + if self.use_spatial_tiling and ( + x.shape[-1] > self.tile_sample_min_size or x.shape[-2] > self.tile_sample_min_size + ): + return self.spatial_tiled_encode(x) + + if self.use_compile: + + @torch.compile + def encoder(x): + return self.encoder(x) + + return encoder(x) + return self.encoder(x) + + if len(x.shape) != 5: # (B, C, T, H, W) + x = x[:, :, None] + assert len(x.shape) == 5 # (B, C, T, H, W) + if x.shape[2] == 1: + x = x.expand(-1, -1, self.ffactor_temporal, -1, -1) + else: + assert x.shape[2] != self.ffactor_temporal and x.shape[2] % self.ffactor_temporal == 0 + + if self.use_slicing and x.shape[0] > 1: + if self.slicing_bsz == 1: + encoded_slices = [_encode(x_slice) for x_slice in x.split(1)] + else: + sections = [self.slicing_bsz] * (x.shape[0] // self.slicing_bsz) + if x.shape[0] % self.slicing_bsz != 0: + sections.append(x.shape[0] % self.slicing_bsz) + encoded_slices = [_encode(x_slice) for x_slice in x.split(sections)] + h = torch.cat(encoded_slices) + else: + h = _encode(x) + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + def decode(self, z: Tensor, return_dict: bool = True, generator=None): + """ + Decodes the input by passing through the decoder network. + Support slicing and tiling for memory efficiency. + """ + + def _decode(z): + if self.use_temporal_tiling and z.shape[-3] > self.tile_latent_min_tsize: + return self.temporal_tiled_decode(z) + if self.use_spatial_tiling and ( + z.shape[-1] > self.tile_latent_min_size or z.shape[-2] > self.tile_latent_min_size + ): + return self.spatial_tiled_decode(z) + return self.decoder(z) + + if self.use_slicing and z.shape[0] > 1: + decoded_slices = [_decode(z_slice) for z_slice in z.split(1)] + decoded = torch.cat(decoded_slices) + else: + decoded = _decode(z) + + if z.shape[-3] == 1: + decoded = decoded[:, :, -1:] + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py new file mode 100644 index 000000000..360904b5e --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py @@ -0,0 +1,1155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import random +import warnings +from collections import defaultdict +from copy import deepcopy +from typing import Any + +import torch +import torch.nn.functional as F +from diffusers.utils.outputs import BaseOutput +from transformers import AutoTokenizer + +from .hunyuan_image_3_transformer import ImageInfo, JointImageInfo, default + + +class TokenizerEncodeOutput(BaseOutput): + tokens: torch.Tensor | None = None + timestep_scatter_index: torch.Tensor | None = None + guidance_scatter_index: torch.Tensor | None = None + text_slices: list[slice] | None = None + gen_image_slices: list[slice] | None = None + joint_image_slices: list[slice] | None = None + cond_vae_image_slices: list[slice] | None = None + cond_vit_image_slices: list[slice] | None = None + text_mask: torch.Tensor | None = None + gen_image_mask: torch.Tensor | None = None + cond_vae_image_mask: torch.Tensor | None = None + cond_vit_image_mask: torch.Tensor | None = None + real_pos: torch.Tensor | None = None + all_image_slices: list[slice] | None = None + cond_timestep_scatter_index: torch.Tensor | None = None + gen_timestep_scatter_index: torch.Tensor | None = None + + +class Conversation: + roles: list[str] = ["User", "Assistant"] + sep: str = "\n\n" + + +class TokenizerWrapper: + def __init__(self, tokenizer): + if isinstance(tokenizer, str): + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer) + else: + self.tokenizer = tokenizer + + # Define short names + self.bos_token_id = self.tokenizer.bos_token_id + self.eos_token_id = self.tokenizer.eos_token_id + self.pad_token_id = self.tokenizer.pad_token_id + self.boi_token_id = self.tokenizer.convert_tokens_to_ids("") + self.eoi_token_id = self.tokenizer.convert_tokens_to_ids("") + self.img_token_id = self.tokenizer.convert_tokens_to_ids("") + self.cfg_token_id = self.tokenizer.convert_tokens_to_ids("") + self.end_answer_token_id = self.tokenizer.convert_tokens_to_ids("") + self.end_recaption_token_id = self.tokenizer.convert_tokens_to_ids("") + self.ratio_token_offset = self.tokenizer.convert_tokens_to_ids("") + self.special_token_map = self.tokenizer.added_tokens_encoder + + def pad(self, tensor_list, dim=0, pad_val=None): + if pad_val is None: + pad_val = self.pad_token_id + max_len = max([t.shape[dim] for t in tensor_list]) + padded_tensor_list = [] + for t in tensor_list: + if t.shape[dim] < max_len: + assert pad_val is not False, "Not allowed pad." + t = F.pad(t, (0, max_len - t.shape[dim]), value=pad_val) + padded_tensor_list.append(t) + return padded_tensor_list + + def encode(self, *args, **kwargs): + return self.tokenizer.encode(*args, **kwargs) + + def decode(self, *args, **kwargs): + return self.tokenizer.decode(*args, **kwargs) + + def encode_text( + self, + *texts, + uncond_enabled: bool | list[bool] | None = None, + uncond_p: float | None = None, + max_length: int | None = None, + pad: str | None = None, + return_lengths: bool = False, + ) -> list[torch.Tensor]: + r""" + Encode text and image for AR-like model training of the text-to-image/instruction tuning tasks. + Support encode multiple texts at once. Each text can be separately conditioned or unconditioned + based on the uncond_flags and a uniform uncond_p. + ** token is always prepended to the text tokens.** + + Args: + texts (`str` or `List[str]`): + List of texts to be encoded. + uncond_enabled (`bool` or `List[bool]`): + List of flags to indicate whether the text should be unconditioned. + If False, the text will never be unconditioned. + If True, the text will be unconditioned with uncond_p. + uncond_p (`float`): + Probability to the unconditional text. Only works when uncond_enabled is True. + max_length (`int`): + Maximum length of the encoded text. + pad (`str`, *optional*): + Padding method. Can be 'left' or 'right'. + return_lengths (`bool`): + Whether to return the length of each encoded text. + + Returns: + `tuple[torch.Tensor, List[int]]` or `torch.Tensor`: + If `return_lengths` is True, returns a tuple of (encoded_tokens, lengths). + If `return_lengths` is False, returns only the encoded_tokens. + """ + if pad is not None: + assert max_length is not None, "max_length should be provided when pad is not None." + + if uncond_enabled is None: + uncond_enabled = [True] * len(texts) + elif isinstance(uncond_enabled, bool): + uncond_enabled = [uncond_enabled] * len(texts) + if len(uncond_enabled) != len(texts): + print(uncond_enabled, texts) + assert len(uncond_enabled) == len(texts), ( + f"Length of uncond_flags should be equal to the number of texts, " + f"but got {len(uncond_enabled)} and {len(texts)}." + ) + + # Prepare text/uncond tokens + # TODO: If len(texts) > 1, such as instruction + prompt in inpainting, we need to determine how to do uncond. + # Now all texts will be cond or uncond at the same time. + do_uncond_drop = (uncond_p is not None) and (random.random() < uncond_p) + text_tokens, lengths = [], [] + cum_length = 0 + for text, uncond_flag in zip(texts, uncond_enabled): + # If reach the max_length and there still have unencoded texts, give a warning message and break the loop. + if max_length is not None and cum_length >= max_length: + warnings.warn( + f"Text length exceeds the max_length({max_length}). The remaining texts will be ignored: " + f"{text[:80]}..." + ) + break + # Set add_special_tokens=False to avoid adding token in some LLMs. + if isinstance(text, str): + text_token = self.tokenizer.encode(text, add_special_tokens=False) + else: + text_token = text + if uncond_flag and do_uncond_drop: + text_token = [self.cfg_token_id] * len(text_token) + # Cutoff the text by max_length if necessary + if max_length is not None and (cum_length + len(text_token)) > max_length: + text_token = text_token[: max_length - cum_length] + text_tokens.extend(text_token) + lengths.append(len(text_token)) + cum_length += len(text_token) + + # Prepend/Append tokens if applicable + if pad is not None and (pad_length := max_length - len(text_tokens)) > 0: + if pad == "left": + text_tokens = [self.pad_token_id] * pad_length + text_tokens + elif pad == "right": + text_tokens = text_tokens + [self.pad_token_id] * pad_length + else: + raise ValueError(f"Unsupported padding method: {pad}.") + + if return_lengths: + return text_tokens, lengths + return text_tokens + + @staticmethod + def _check_key_number_matched(keys, data): + # Assert keys and token_source are matched + assert set(keys) == set(data.keys()), ( + f"Keys in the template and token source should be matched, but got {set(keys)} and {list(data.keys())}." + ) + key_counts = {k: 0 for k in keys} + for key in keys: + key_counts[key] += 1 + for key, count in key_counts.items(): + assert len(data[key]) == count, ( + f"Number of `{key}` in the token source should be matched with the template, but got " + f"{data[key]}({len(data[key])}) and {count}." + ) + + def _add_image_meta_info_token( + self, + token_seq, + token_count, + extra_token_pos, + add_timestep_token=False, + add_image_shape_token=False, + base_size=None, + ratio_idx=None, + image_type=None, + add_guidance_token=False, + ): + if add_image_shape_token: + token_seq.extend( + [self.special_token_map[f""], self.special_token_map[f""]] + ) + token_count += 2 + if add_timestep_token: + token_seq.extend([self.special_token_map[""]]) + extra_token_pos["timestep"].append(token_count) + if image_type is not None: + if image_type == "gen_image": + extra_token_pos["gen_timestep"].append(token_count) + elif image_type in ["joint_image"]: + extra_token_pos["cond_timestep"].append(token_count) + else: + raise ValueError(f"Unsupported image type: {image_type}.") + token_count += 1 + if add_guidance_token: + token_seq.extend([self.special_token_map[""]]) + extra_token_pos["guidance"].append(token_count) + token_count += 1 + return token_count + + def encode_sequence( + self, + template: str, + token_source: dict[str, list], + total_length=None, + add_timestep_token=False, + add_guidance_token=False, + last_key_only_prefix=False, + add_eos=True, + use_front_boi_token=True, + add_pad=True, + add_bos=True, + drop_last: str | bool = "auto", + add_image_shape_token=False, + ) -> tuple[list, dict[str, list]]: + r""" + Encode a sequence based on the template (e.g., `text-image` for t2i, `text-image-image` for instruction tuning) + and token source. + + Args: + template (`str`): + Template of the sequence. E.g., "text-gen_image" means the sequence is composed of text and an image. + "text-text-gen_image" means the sequence is composed of two sections of text and an image. + token_source (`Dict[str, List]`): + Token source for each key in the template, in order. + - text: List[Dict]. + - gen_image: List[Dict]. + - joint_image: List[Dict]. + total_length (`int`): + Total length of the encoded sequence, include padding tokens. + add_timestep_token (`bool`): + Whether to add timestep token before the image tokens. + (Right after the tokens) + add_guidance_token (`bool`): + Whether to add guidance token before the image tokens. + last_key_only_prefix (`bool`): + Whether to only use the modal prefix in the last key. + add_eos (`bool` or `'auto'`): + Whether to add eos token at the end of the sequence. If True, always add eos token. If 'auto', + add eos token only when the total_length is not reached and the last token is not . + use_front_boi_token (`bool`): + Whether to put the token at the front of iw, ih and timestep tokens. + add_pad (`bool` or `'auto'`): + Whether to add padding tokens to the sequence. + add_bos (`bool`): + Whether to add bos token at the beginning of the sequence. + drop_last (`bool` or `'auto'`): + - If auto, drop last tokens exceeding the total_length if the total_length is provided. If cut point is + in the middle of the image tokens, an error will raised. + - If True, drop last tokens exceeding the total_length. + - If False, keep the last tokens exceeding the total_length, even if the total_length is reached. + add_image_shape_token (`bool`): + Whether to add image shape token before the image tokens. (Right before the token) + + Returns: + `tuple[list, dict]`: A tuple containing: + - token_seq (`list`): Encoded token sequence. + - extra_token_pos (`dict`): Positions of extra tokens. + """ + if last_key_only_prefix: + assert add_eos is not True, "add_eos should not be True when last_key_only_prefix is True." + if drop_last is True and total_length is None: + raise ValueError("total_length should be provided when drop_last is True.") + + keys = template.split("-") + modal_length = len(keys) + index_indicator = {k: 0 for k in token_source} + for k, v in token_source.items(): + assert isinstance(v, (list, tuple)), ( + f"Value of `{k}` in the token source should be a list or tuple, but got {type(v)}." + ) + self._check_key_number_matched(keys, token_source) + + token_seq = [] + token_count = 0 + extra_token_pos = defaultdict(list) + if add_bos: + token_seq.append(self.bos_token_id) + token_count += 1 + # If drop_last is True, we check the token_count on the fly and exit the loop if the total_length is reached. + # This check is only applied to the block tokens. Block tokens mean the tokens that are unsplittable, like + # image tokens. Text tokens are splittable, so we don't need to check the token_count for text. + # If the loop is broken by drop_last, we don't add the eos token at the end because the sequence is not + # complete. + drop_last_break = False + for i, key in enumerate(keys): + source = token_source[key][index_indicator[key]] + if key == "text": + token_seq.extend(source) # text token sequence + extra_token_pos["_start"].append(token_count) + token_count += len(source) + extra_token_pos["_end"].append(token_count - 1) + + elif key == "gen_image": + if isinstance(source, int): + source = {"length": source} + extra_count = ( + 2 + + (1 if source.get("timestep", add_timestep_token) else 0) + + (1 if source.get("guidance", add_guidance_token) else 0) + + (2 if source.get("image_shape", add_image_shape_token) else 0) + ) + if drop_last is True and token_count + extra_count + source["length"] > total_length: + drop_last_break = True + break + if source.get("front_boi", use_front_boi_token): + token_seq.append(self.boi_token_id) + extra_token_pos["boi"].append(token_count) + token_count += 1 + token_count = self._add_image_meta_info_token( + token_seq=token_seq, + token_count=token_count, + extra_token_pos=extra_token_pos, + add_timestep_token=source.get("timestep", add_timestep_token), + add_guidance_token=source.get("guidance", add_guidance_token), + add_image_shape_token=source.get("image_shape", add_image_shape_token), + base_size=source.get("base_size"), + ratio_idx=source.get("ratio_idx"), + image_type=key, + ) + if not source.get("front_boi", use_front_boi_token): + token_seq.append(self.boi_token_id) + extra_token_pos["boi"].append(token_count) + token_count += 1 + if last_key_only_prefix and i == modal_length - 1: + pass # for AR inference + else: + token_seq.extend( + [self.img_token_id] * source["length"] # token number + + [self.eoi_token_id] + ) + extra_token_pos["_start"].append(token_count) + extra_token_pos["_start"].append(token_count) + token_count += source["length"] + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["eoi"].append(token_count) + token_count += 1 # + + elif key == "joint_image": + assert isinstance(source["length"], list) and len(source["length"]) == 2, ( + "joint_image length should be a list of two integers" + ) + extra_count = ( + 2 + + 1 + + ( # boi, eoi, joint_img_sep + 1 if source.get("timestep", add_timestep_token) else 0 + ) + + (2 if source.get("image_shape", add_image_shape_token) else 0) + ) + if drop_last is True and token_count + extra_count + sum(source["length"]) > total_length: + drop_last_break = True + break + if source.get("front_boi", use_front_boi_token): + token_seq.append(self.boi_token_id) # Use patched boi for Janus, otherwise useing default + extra_token_pos["boi"].append(token_count) + token_count += 1 + token_count = self._add_image_meta_info_token( + token_seq=token_seq, + token_count=token_count, + extra_token_pos=extra_token_pos, + add_timestep_token=source.get("timestep", add_timestep_token), + add_image_shape_token=source.get("image_shape", add_image_shape_token), + base_size=source.get("base_size"), + ratio_idx=source.get("ratio_idx"), + image_type=key, + ) + if not source.get("front_boi", use_front_boi_token): + token_seq.append(self.boi_token_id) + extra_token_pos["boi"].append(token_count) + token_count += 1 + if last_key_only_prefix and i == modal_length - 1: + pass # for AR inference + else: + token_seq.extend([self.img_token_id] * source["length"][0]) + extra_token_pos["_start"].append(token_count) + extra_token_pos["_start"].append(token_count) + extra_token_pos["_start"].append(token_count) + token_count += source["length"][0] + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["_end"].append(token_count - 1) + + token_seq.extend([self.special_token_map[""]]) + extra_token_pos["joint_img_sep"].append(token_count) + token_count += 1 + + token_seq.extend([self.img_token_id] * source["length"][1]) + extra_token_pos["_start"].append(token_count) + extra_token_pos["_start"].append(token_count) + token_count += source["length"][1] + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["_end"].append(token_count - 1) + extra_token_pos["_end"].append(token_count - 1) + + token_seq.extend([self.eoi_token_id]) + extra_token_pos["eoi"].append(token_count) + token_count += 1 # + + else: + raise ValueError(f"Not supported key: {key}") + index_indicator[key] += 1 + + if add_eos is True and not drop_last_break: + # Typically used for t2i task. + token_seq.append(self.eos_token_id) + extra_token_pos["eos"].append(token_count) + token_count += 1 + elif add_eos == "auto" and not drop_last_break: + # Typically used for lm and mmu task. + if token_seq[-1] != self.eos_token_id and (total_length is None or token_count < total_length): + token_seq.append(self.eos_token_id) + extra_token_pos["eos"].append(token_count) + token_count += 1 + + if total_length: + # Check token count and clip sequence if necessary + if token_count > total_length and drop_last: + # Assert clip position is not in the middle of the block-wise tokens (gen_image, joint_image) + for start_key, end_key in [ + ("_start", "_end"), + ("_start", "_end"), + ("_start", "_end"), + ("_start", "_end"), + ]: + if start_key in extra_token_pos and end_key in extra_token_pos: + assert all( + (start > total_length or end + 1 < total_length) + for start, end in zip(extra_token_pos[start_key], extra_token_pos[end_key]) + ), ( + "Clip position should not be in the middle of the image tokens.\n" + f"Below is the text:\n{self._shorten_text(self.tokenizer.decode(token_seq))}" + ) + token_seq = token_seq[:total_length] + + # Pad the sequence if necessary + pad_num = max(0, total_length - len(token_seq)) + if add_pad and pad_num: + token_seq.extend([self.pad_token_id] * pad_num) + extra_token_pos["first_pad"].append(token_count) + + return token_seq, extra_token_pos + + def batch_gen_infer( + self, + infer_fn, + prompt_list: list, + negative_prompt_list: list | None = None, + infer_fn_kwargs_list: list[dict[str, int]] | None = None, + do_classifier_free_guidance=False, + condition_repeat_times: int = 1, + uncondition_repeat_times: int = 1, + ): + r""" + Batch inference for the AR-like model training of the text-to-image/instruction tuning tasks. + + Args: + infer_fn (`callable`): + Inference function to encode the prompt. + prompt_list (`list`): + List of prompts. Each element can be a single prompt or a list of prompts passed to the infer_fn. + negative_prompt_list (`list`, *optional*): + List of negative prompts. Only used when do_classifier_free_guidance is True. If None, will use + token sequence as negative prompt. + infer_fn_kwargs_list (`List[Dict[str, int]]`, *optional*): + List of keyword arguments for the infer_fn. + do_classifier_free_guidance (`bool`): + Whether to do classifier-free guidance. + condition_repeat_times (`int`): + Support multi-condition. + uncondition_repeat_times (`int`): + Support multi-uncondition. + """ + if infer_fn_kwargs_list is None: + infer_fn_kwargs_list = [{} for _ in prompt_list] + + # [n_output, bsz] + cond_results_list = None + uncond_results_list = None + output_type_list = [] + + for prompt_idx, (prompt, infer_fn_kwargs) in enumerate(zip(prompt_list, infer_fn_kwargs_list)): + if not isinstance(prompt, (list, tuple)): + prompt = [prompt] + cond_kwargs = {"uncond_p": 0.0} if do_classifier_free_guidance else {} + results = infer_fn( + *prompt, + **infer_fn_kwargs, + **cond_kwargs, + ) + output_type_list.append((type(results), len(results) if isinstance(results, (list, tuple)) else 1)) + if isinstance(results, dict): + raise ValueError("Make batch on dict is not supported. Please return list or tuple for infer_fn.") + if not isinstance(results, (list, tuple)): + results = (results,) + if cond_results_list is None: + cond_results_list = [[] for _ in results] + uncond_results_list = [[] for _ in results] + for i, result in enumerate(results): + cond_results_list[i].append(result) + + if do_classifier_free_guidance: + if negative_prompt_list is None: + uncond_kwargs = {"uncond_p": 1.0} + uncond_results = infer_fn( + *prompt, + **infer_fn_kwargs, + **uncond_kwargs, + ) + else: + negative_prompt = negative_prompt_list[prompt_idx] + if not isinstance(negative_prompt, (list, tuple)): + negative_prompt = [negative_prompt] + uncond_results = infer_fn( + *negative_prompt, + **infer_fn_kwargs, + ) + if isinstance(uncond_results, TokenizerEncodeOutput): + uncond_results_list.append(uncond_results) + else: + for i, result in enumerate(uncond_results): + uncond_results_list[i].append(result) + + assert all(output_type_list[0] == n for n in output_type_list), ( + f"Number of outputs should be equal for all samples, but got {output_type_list}." + ) + output_type, output_num = output_type_list[0] + + def make_batch(batch_cond_item, batch_uncond_item): + # Process each output item to make batch + first = batch_cond_item[0] # The first element in the batch + if isinstance(first, torch.Tensor): + stacked_item = torch.stack( + self.pad( + batch_cond_item * condition_repeat_times + batch_uncond_item * uncondition_repeat_times, + ) + ) + + elif first is None: + assert all(item is None for item in batch_cond_item + batch_uncond_item), ( + f"The first cond item is None, but some items are not None:\n\n" + f"condition: {batch_cond_item}\n\n" + f"uncondition: {batch_uncond_item}" + ) + stacked_item = None + + elif isinstance(first, (list, tuple)): + # If the output item is a list or tuple, we treat it as a whole, and won't make nested batch any more. + stacked_item = batch_cond_item * condition_repeat_times + batch_uncond_item * uncondition_repeat_times + + elif isinstance(first, TokenizerEncodeOutput): + stacked_item = {} + # Traverse not-None attributes + for key in list(first.keys()): + merged_list = [cond_item[key] for cond_item in batch_cond_item] * condition_repeat_times + [ + uncond_item[key] for uncond_item in batch_uncond_item + ] * uncondition_repeat_times + if isinstance(first[key], torch.Tensor): + if "mask" in key: + pad_val = 0.0 + elif key == "tokens": + pad_val = self.special_token_map[""] + else: + pad_val = False # Should not pad for other tensors + stacked_item[key] = torch.stack(self.pad(merged_list, pad_val=pad_val), dim=0) + elif isinstance(first[key], list): + stacked_item[key] = merged_list + elif first[key] is None: + pass + else: + raise ValueError(f"Unsupported type of {key}: {type(first[key])}.") + stacked_item = TokenizerEncodeOutput(stacked_item) + + else: + raise TypeError(f"Making batch on type {type(first)} is not supported.") + + return stacked_item + + stacked_outputs = [] + for cond_results, uncond_results in zip(cond_results_list, uncond_results_list): + stacked_outputs.append(make_batch(cond_results, uncond_results)) + + if output_type is list: + return stacked_outputs + elif output_type is tuple: + return tuple(stacked_outputs) + elif output_num == 1: + return stacked_outputs[0] + else: + raise ValueError(f"Unsupported output type: {output_type}.") + + @staticmethod + def parse_extra_token_pos(extra_token_pos, prefix, tokens, rng=None): + if rng is None: + rng = slice(None) + image_slices = ( + [ + slice(start, end + 1) + for start, end in zip( + extra_token_pos[f"<{prefix}>_start"][rng], extra_token_pos[f"<{prefix}>_end"][rng] + ) + ] + if f"<{prefix}>_start" in extra_token_pos and f"<{prefix}>_end" in extra_token_pos + else [] + ) + if image_slices: + image_mask = torch.zeros_like(tokens, dtype=torch.bool) + for image_slice in image_slices: + image_mask[image_slice] = True + else: + image_mask = None + return image_slices, image_mask + + def encode_general( + self, + sections: list[dict[str, Any]] | None = None, + max_token_length: int | None = None, + add_eos="auto", + use_text_mask=True, + add_pad="auto", + add_bos=True, + drop_last="auto", + ) -> TokenizerEncodeOutput: + r""" + General encode function to encode a sequence with multiple sections of text and images. + Each section is a dict with a `type` key and other keys depending on the type. + + Supported section types: + + - text: dict with keys: + - text (`str` or `List[int]`): Text to be encoded. Either `text` or `tokens` should be provided. + - tokens (`List[int]`): Pre-encoded text tokens. Either `text` or `tokens` should be provided. + - uncond_enabled (`bool`): Whether to enable uncondition for this text section. + - uncond_p (`float`): Probability to drop the text section for uncondition. + - max_length (`int`): Maximum length of the text section. + - ignore (`bool`): Whether to ignore this text section in the text mask. + - start_offset (`int`): Start offset of the text mask. + - end_offset (`int`): End offset of the text mask. + + - gen_image: dict with keys: + - token_length (`int`): Number of image tokens. + - add_timestep_token (`bool`): Whether to add timestep token before the image tokens. + - add_guidance_token (`bool`): Whether to add guidance token before the image tokens. + - use_front_boi_token (`bool`): Whether to put the token. + - add_image_shape_token (`bool`): Whether to add image shape token before the image tokens. + - base_size (`int`): Base size of the image. + - ratio_idx (`int`): Ratio index of the image. + + - joint_image: dict with keys: + - token_length (`List[int]`): Number of image tokens for the two images. + - add_timestep_token (`bool`): Whether to add timestep token before the image tokens. + - use_front_boi_token (`bool`): Whether to put the token. + - add_image_shape_token (`bool`): Whether to add image shape token before the image tokens. + - base_size (`int`): Base size of the image. + - ratio_idx (`int`): Ratio index of the image. + + Args: + sections (`List[Dict[str, Any]]`): + List of sections to be encoded. + max_token_length (`int`): + Maximum length of the encoded token sequence. + add_eos (`bool` or `'auto'`): + Whether to add eos token at the end of the sequence. + use_text_mask (`bool`): + Whether to generate text mask. + add_pad (`bool` or `'auto'`): + Whether to add padding tokens to the sequence. If True and total_length is not reached, + add padding tokens. + add_bos (`bool`): + Whether to add bos token at the beginning of the sequence. + drop_last (`bool` or `'auto'`): + - If auto, drop last tokens exceeding the total_length if the total_length is provided. + If cut point is in the middle of the image tokens, an error will raised. + - If True, drop last tokens exceeding the total_length. If cut point is in the + middle of the image tokens, all the successive image tokens will be dropped. + - If False, keep the last tokens exceeding the total_length, even if the total_length + is reached. + + Returns: + `TokenizerEncodeOutput`: Encoded token sequence and extra information. + """ + if sections is None: + raise ValueError("sections must be provided.") + template = "-".join([section["type"] for section in sections]) + + sections = deepcopy(sections) + token_source = defaultdict(list) + text_mask_specs = [] + for section in sections: + if section["type"] == "text": + text = self.encode_text( + section["text"] if "text" in section else section["tokens"], + uncond_enabled=section.get("uncond_enabled"), + uncond_p=section.get("uncond_p"), + max_length=section.get("max_length"), + ) + token_source["text"].append(text) + text_mask_specs.append( + dict( + ignore=section.get("ignore", False), + start_offset=section.get("start_offset", 0), + end_offset=section.get("end_offset", 0), + ) + ) + elif section["type"] == "gen_image": + token_source["gen_image"].append( + dict( + length=section["token_length"], + timestep=section.get("add_timestep_token", False), + guidance=section.get("add_guidance_token", False), + front_boi=section.get("use_front_boi_token", False), + image_shape=section.get("add_image_shape_token", False), + base_size=section.get("base_size"), + ratio_idx=section.get("ratio_idx"), + ) + ) + elif section["type"] == "joint_image": + token_source["joint_image"].append( + dict( + length=section["token_length"], + timestep=section.get("add_timestep_token", False), + front_boi=section.get("use_front_boi_token", False), + image_shape=section.get("add_image_shape_token", False), + base_size=section.get("base_size"), + ratio_idx=section.get("ratio_idx"), + ) + ) + else: + raise ValueError(f"Invalid section type: {section['type']}") + + # Combine text and image tokens + full_token_seq, extra_token_pos = self.encode_sequence( + template=template, + token_source=dict(token_source), + total_length=max_token_length, + add_eos=add_eos, + add_pad=add_pad, + add_bos=add_bos, + drop_last=drop_last, + ) + full_seq_token_tensor = torch.tensor(full_token_seq, dtype=torch.long) + + timestep_scatter_index = ( + torch.tensor(extra_token_pos["timestep"], dtype=torch.long) if "timestep" in extra_token_pos else None + ) + guidance_scatter_index = ( + torch.tensor(extra_token_pos["guidance"], dtype=torch.long) if "guidance" in extra_token_pos else None + ) + cond_timestep_scatter_index = ( + torch.tensor(extra_token_pos["cond_timestep"], dtype=torch.long) + if "cond_timestep" in extra_token_pos + else None + ) + gen_timestep_scatter_index = ( + torch.tensor(extra_token_pos["gen_timestep"], dtype=torch.long) + if "gen_timestep" in extra_token_pos + else None + ) + + # Gen image mask + gen_image_slices, gen_image_mask = self.parse_extra_token_pos(extra_token_pos, "img", full_seq_token_tensor) + # Joint image + joint_image_slices, _ = self.parse_extra_token_pos(extra_token_pos, "joint_img", full_seq_token_tensor) + # Conditional vae image + cond_vae_image_slices, cond_vae_image_mask = self.parse_extra_token_pos( + extra_token_pos, "vae_img", full_seq_token_tensor + ) + # Conditional vit image + cond_vit_image_slices, cond_vit_image_mask = self.parse_extra_token_pos( + extra_token_pos, "vit_img", full_seq_token_tensor + ) + # All image slices (gen_image, joint_image) + all_image_slices = ( + [ + slice(start, end + 1) + for start, end in zip(extra_token_pos["_start"], extra_token_pos["_end"]) + ] + if "_start" in extra_token_pos and "_end" in extra_token_pos + else [] + ) + + # Text mask + text_slices = ( + [ + slice(start, end + 1) + for start, end in zip(extra_token_pos["_start"], extra_token_pos["_end"]) + ] + if "_start" in extra_token_pos and "_end" in extra_token_pos + else [] + ) + assert len(text_slices) <= len(text_mask_specs), ( + f"Number of text slices ({len(text_slices)}) should be less than or equal to " + f"number of text mask specs ({len(text_mask_specs)})" + ) + if use_text_mask: + text_mask = torch.zeros_like(full_seq_token_tensor, dtype=torch.float32) + for text_slice, mask_spec in zip(text_slices, text_mask_specs): + if not mask_spec["ignore"]: + real_slice = slice( + text_slice.start + mask_spec["start_offset"], text_slice.stop + mask_spec["end_offset"] + ) + text_mask[real_slice] = 1.0 + else: + text_mask = None + + # real_pos is the first position of the token + real_pos = torch.tensor(extra_token_pos.get("first_pad", [full_seq_token_tensor.shape[0]]), dtype=torch.long) + + return TokenizerEncodeOutput( + tokens=full_seq_token_tensor, + timestep_scatter_index=timestep_scatter_index, + guidance_scatter_index=guidance_scatter_index, + text_slices=text_slices, + gen_image_slices=gen_image_slices, + joint_image_slices=joint_image_slices, + cond_vae_image_slices=cond_vae_image_slices, + cond_vit_image_slices=cond_vit_image_slices, + text_mask=text_mask, + gen_image_mask=gen_image_mask, + cond_vae_image_mask=cond_vae_image_mask, + cond_vit_image_mask=cond_vit_image_mask, + real_pos=real_pos, + all_image_slices=all_image_slices, + cond_timestep_scatter_index=cond_timestep_scatter_index, + gen_timestep_scatter_index=gen_timestep_scatter_index, + ) + + def get_cot_sections(self, cot_text, uncond_kwargs, cot_max_length=None, drop_think=False): + if not cot_text: # None or empty + return [] + if "" in cot_text and "" in cot_text: + before_think_sec = cot_text.split("")[0] + after_think_sec = cot_text.split("")[1] + think_sec = cot_text.split("")[1].split("")[0] + return ( + self.get_cot_sections(before_think_sec, uncond_kwargs, drop_think=drop_think) + + ( + [ + dict(type="text", text=""), + dict(type="text", text=think_sec, max_length=cot_max_length, **uncond_kwargs), + dict(type="text", text=""), + ] + if not drop_think + else [] + ) + + self.get_cot_sections(after_think_sec, uncond_kwargs, drop_think=drop_think) + ) + + if "" in cot_text and "" in cot_text: + before_recaption_sec = cot_text.split("")[0] + after_recaption_sec = cot_text.split("")[1] + recaption_sec = cot_text.split("")[1].split("")[0] + return ( + self.get_cot_sections(before_recaption_sec, uncond_kwargs, drop_think=drop_think) + + [ + dict(type="text", text=""), + dict(type="text", text=recaption_sec, max_length=cot_max_length, **uncond_kwargs), + dict(type="text", text=""), + ] + + self.get_cot_sections(after_recaption_sec, uncond_kwargs, drop_think=drop_think) + ) + + return [ + dict(type="text", text=cot_text, **uncond_kwargs), + ] + + def apply_general_template( + self, + message_list, + max_length=None, + add_assistant_prefix=False, + answer="auto", + bot_task="auto", + sequence_template="instruct", + uncond_p=0.0, + cfg_factor=1, + batchify=False, + image_base_size=1024, + drop_think=False, + ): + # If cfg_factor > 1, we need to repeat the unconditioned part + if batchify: + assert isinstance(message_list[0], list), ( + f"When batchify is True, message_list should be a list of list, but got [{type(message_list[0])}, ...]." + ) + return self.batch_gen_infer( + infer_fn=self.apply_general_template, + prompt_list=[[]], + infer_fn_kwargs_list=[ + dict( + message_list=message_list_i, + max_length=max_length, + add_assistant_prefix=add_assistant_prefix, + answer=answer, + bot_task=bot_task, + sequence_template=sequence_template, + image_base_size=image_base_size, + drop_think=drop_think, + ) + for message_list_i in message_list + ], + do_classifier_free_guidance=cfg_factor > 1, + condition_repeat_times=1, + uncondition_repeat_times=cfg_factor - 1, + ) + + conv = Conversation() + uncond_kwargs = dict(uncond_enabled=uncond_p == 1.0, uncond_p=uncond_p) + + def process_successive_message( + _message_list, _cur_message_idx, role, prefix, suffix, answer_prefix="", answer_suffix="" + ): + _sub_sections = [] + while _cur_message_idx < len(message_list) and _message_list[_cur_message_idx]["role"] == role: + message = _message_list[_cur_message_idx] + if message["type"] == "text": + text = message["content"] + if role == "system": + _sub_sections.append(dict(type="text", text=text)) + elif role == "assistant": + if ("" in text and "" in text) or ( + "" in text and "" in text + ): + _sub_sections.extend(self.get_cot_sections(text, uncond_kwargs, drop_think=drop_think)) + else: + _sub_sections.append(dict(type="text", text=text, **uncond_kwargs)) + else: + _sub_sections.append( + dict(type="text", text=f"{answer_prefix}{text}{answer_suffix}", **uncond_kwargs) + ) + elif message["type"] == "gen_image": + info = message["content"] + assert isinstance(info, ImageInfo), f"Expected ImageInfo, but got {type(info)}" + if role == "assistant": + _sub_sections.append(dict(type="text", text=answer_prefix)) + _sub_sections.append(dict(type=message["type"], **info.meta_info)) + if role == "assistant": + _sub_sections.append(dict(type="text", text=answer_suffix)) + elif message["type"] == "joint_image": + info = message["content"] + assert isinstance(info, JointImageInfo), f"Expected JointImageInfo, but got {type(info)}" + _sub_sections.append(dict(type=message["type"], **info.meta_info)) + else: + raise ValueError(f"Unknown message type: {message['type']}") + _cur_message_idx += 1 + if len(_sub_sections) > 0: + # Add role prefix and suffix + _sub_sections.insert(0, dict(type="text", text=prefix)) + _sub_sections.append(dict(type="text", text=suffix)) + return _sub_sections, _cur_message_idx + + # Define assistant prefix and suffix + if (answer == "auto" and sequence_template == "instruct") or answer is True: + answer_prefix, answer_suffix = "", "" + else: + answer_prefix, answer_suffix = "", "" + if sequence_template == "pretrain": + system_suffix = "" + user_prefix = "" + user_suffix = "" + bot_prefix = "" + bot_suffix = "" + else: + system_suffix = f"{conv.sep}" + user_prefix = f"{conv.roles[0]}: " + user_suffix = f"{conv.sep}" + bot_prefix = f"{conv.roles[1]}: " + bot_suffix = f"{conv.sep}" + + # Process successive user and assistant messages + sections = [] + cur_message_idx = 0 + final_role = None + while cur_message_idx < len(message_list): + # Process successive system messages + sub_sections, cur_message_idx = process_successive_message( + message_list, cur_message_idx, role="system", prefix="", suffix=system_suffix + ) + # Add to the template and sections + sections.extend(sub_sections) + if len(sub_sections) > 0: + final_role = "system" + + # Process successive user messages + sub_sections, cur_message_idx = process_successive_message( + message_list, cur_message_idx, role="user", prefix=user_prefix, suffix=user_suffix + ) + # Add to the template and sections + sections.extend(sub_sections) + if len(sub_sections) > 0: + final_role = "user" + + # Process successive assistant messages + sub_sections, cur_message_idx = process_successive_message( + message_list, + cur_message_idx, + role="assistant", + prefix=bot_prefix, + suffix=bot_suffix, + answer_prefix=answer_prefix, + answer_suffix=answer_suffix, + ) + # Add to the template and sections + sections.extend(sub_sections) + if len(sub_sections) > 0: + final_role = "assistant" + + if add_assistant_prefix: + if final_role == "assistant": + # Avoid adding prefix twice + _bot_prefix = "" + # Remove the final bot_suffix + if len(sections) > 0 and sections[-1]["type"] == "text" and sections[-1]["text"] == bot_suffix: + sections = sections[:-1] + else: + _bot_prefix = bot_prefix + # We can add special tokens for the bot latest message according to different tasks + bot_response_prefix = dict( + auto=_bot_prefix, + image="", + think=f"{_bot_prefix}", + recaption=f"{_bot_prefix}", + img_ratio=f"{_bot_prefix}{answer_prefix}", + )[bot_task] + sections.append(dict(type="text", text=bot_response_prefix)) + + output = self.encode_general( + sections=sections, + use_text_mask=False, + add_eos=False, + add_pad=False, + ) + + if max_length is not None: + if output.tokens.shape[-1] > max_length: + raise ValueError( + f"Encoded token length {output.tokens.shape[-1]} exceeds max_length {max_length}.\n" + f"Please set a larger max_length or check the input messages:\n{message_list}" + ) + + return output, sections + + def apply_chat_template( + self, + batch_prompt: list[str] | None = None, + batch_message_list: list[list[dict[str, Any]]] | None = None, + mode: str = "gen_text", + batch_gen_image_info: list[ImageInfo] | None = None, + batch_cond_image_info: list[JointImageInfo] | list[list[JointImageInfo]] | None = None, + batch_system_prompt: list[str] | None = None, + batch_cot_text: list[str] | None = None, + max_length: int | None = None, + bot_task: str = "auto", # auto/image/think/recaption/img_ratio + image_base_size: int = 1024, + sequence_template: str = "pretrain", + cfg_factor: int = 1, + add_assistant_prefix: bool | None = None, + drop_think: bool = False, + ) -> dict[str, Any]: + assert bot_task in ["image", "auto", "think", "recaption", "img_ratio"], ( + f"bot_task should be one of ['image', 'auto', 'think', 'recaption', 'img_ratio'], but got {bot_task}." + ) + + if batch_message_list is None: + # Simple text-to-image or text-cot-to-image task + batch_size = len(batch_prompt) + + # Batchify inputs + if not isinstance(batch_system_prompt, list): + batch_system_prompt = [batch_system_prompt] * batch_size + if not isinstance(batch_gen_image_info, list): + batch_gen_image_info = [batch_gen_image_info] * batch_size + if batch_cot_text is not None: + assert len(batch_cot_text) == batch_size, ( + f"batch_cot_text should have the same length as batch_size ({batch_size}), " + f"but got {len(batch_cot_text)}." + ) + else: + batch_cot_text = [None] * batch_size + if batch_cond_image_info is not None: + assert len(batch_cond_image_info) == batch_size, ( + f"batch_cond_image_info should have the same length as batch_size ({batch_size}), " + f"but got {len(batch_cond_image_info)}." + ) + batch_cond_image_info = [ + cond_image_info if isinstance(cond_image_info, list) else [cond_image_info] + for cond_image_info in batch_cond_image_info + ] + else: + batch_cond_image_info = [[] for _ in range(batch_size)] + + # Convert single round materials into standard message list + batch_message_list = [] + for prompt, system_prompt, cot_text, gen_image_info, cond_image_info_list in zip( + batch_prompt, + batch_system_prompt, + batch_cot_text, + batch_gen_image_info, + batch_cond_image_info, + ): + message_list = [] + # 1. system prompt section + if system_prompt: + message_list.append(dict(role="system", type="text", content=system_prompt, context_type="str")) + # 2. user inputs sections + # 2.1 image inputs + if len(cond_image_info_list) > 0: + message_list.extend( + [ + dict(role="user", type="joint_image", content=cond_image_info, context_type="image_info") + for cond_image_info in cond_image_info_list + ] + ) + # 2.2 text inputs + message_list.append(dict(role="user", type="text", content=prompt, context_type="str")) + # 3. assistant answer sections + if cot_text is not None: + message_list.append(dict(role="assistant", type="text", content=cot_text, context_type="str")) + if mode == "gen_image": + message_list.append( + dict(role="assistant", type="gen_image", content=gen_image_info, context_type="image_info") + ) + batch_message_list.append(message_list) + + output, sections = self.apply_general_template( + message_list=batch_message_list, + max_length=max_length, + add_assistant_prefix=default(add_assistant_prefix, mode != "gen_image"), + bot_task=bot_task, + sequence_template=sequence_template, + cfg_factor=cfg_factor, + batchify=True, + image_base_size=image_base_size, + drop_think=drop_think, + ) + return dict(output=output, sections=sections) 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 new file mode 100644 index 000000000..adb582741 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py @@ -0,0 +1,2796 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import glob +import inspect +import logging +import math +import os +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import Any, cast + +import numpy as np +import regex as re +import torch +import torch.nn.functional as F +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.image_processor import VaeImageProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils.outputs import BaseOutput +from diffusers.utils.torch_utils import randn_tensor +from einops import rearrange +from PIL import Image +from safetensors.torch import load_file +from torch import nn +from torchvision import transforms +from transformers import PretrainedConfig, Siglip2ImageProcessorFast +from transformers.cache_utils import Cache +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, +) +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, + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.utils import ( + PPMissingLayer, + is_pp_missing_parameter, + make_layers, +) +from vllm.v1.attention.backend import AttentionType + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.parallel_state import get_pp_group +from vllm_omni.diffusion.layers.rope import RotaryEmbedding + +logger = logging.getLogger(__name__) + + +def _is_moe(config: PretrainedConfig) -> bool: + num_experts = getattr(config, "num_experts", None) + if isinstance(num_experts, int): + return num_experts > 1 + if isinstance(num_experts, list) and num_experts: + # Ensure all elements are integers before calling max. + if all(isinstance(e, int) for e in num_experts): + return max(num_experts) > 1 + else: + return False + return False + + +def _get_cla_factor(config: PretrainedConfig) -> int: + if not getattr(config, "use_cla", False): + return 1 + return getattr(config, "cla_share_factor", 1) + + +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[list[int], int]: + """ + 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 + + +def real_batched_index_select(t, dim, idx): + """index_select for batched index and batched t""" + assert t.ndim >= 2 and idx.ndim >= 2, f"{t.ndim=} {idx.ndim=}" + assert len(t) == len(idx), f"{len(t)=} != {len(idx)=}" + return torch.stack([torch.index_select(t[i], dim - 1, idx[i]) for i in range(len(t))]) + + +def conv_nd(dims, *args, **kwargs): # noqa: N802 + """ + Create a 1D, 2D, or 3D convolution module. + """ + if dims == 1: + return nn.Conv1d(*args, **kwargs) + elif dims == 2: + return nn.Conv2d(*args, **kwargs) + elif dims == 3: + return nn.Conv3d(*args, **kwargs) + raise ValueError(f"unsupported dimensions: {dims}") + + +def normalization(channels, **kwargs): + """ + Make a standard normalization layer. + :param channels: number of input channels. + :return: a nn.Module for normalization. + """ + return nn.GroupNorm(32, channels, **kwargs) + + +def linear(*args, **kwargs): + """ + Create a linear module. + """ + return nn.Linear(*args, **kwargs) + + +def zero_module(module): + """ + Zero out the parameters of a module and return it. + """ + for p in module.parameters(): + p.detach().zero_() + return module + + +def _to_tuple(x, dim=2): + if isinstance(x, int): + return (x,) * dim + elif len(x) == dim: + return x + else: + raise ValueError(f"Expected length {dim} or int, but got {x}") + + +def get_meshgrid_nd(start, *args, dim=2): # noqa: N802 + if len(args) == 0: + # start is grid_size + num = _to_tuple(start, dim=dim) + start = (0,) * dim + stop = num + elif len(args) == 1: + # start is start, args[0] is stop, step is 1 + start = _to_tuple(start, dim=dim) + stop = _to_tuple(args[0], dim=dim) + num = [stop[i] - start[i] for i in range(dim)] + # assert num are all integers + num_int = [int(x) for x in num] + assert (torch.tensor(num) == torch.tensor(num_int)).all(), f"num should be int, but got {num}" + num = num_int + elif len(args) == 2: + # start is start, args[0] is stop, args[1] is num + start = _to_tuple(start, dim=dim) # Left-Top eg: 12,0 + stop = _to_tuple(args[0], dim=dim) # Right-Bottom eg: 20,32 + num = _to_tuple(args[1], dim=dim) # Target Size eg: 32,124 + else: + raise ValueError(f"len(args) should be 0, 1 or 2, but got {len(args)}") + + # PyTorch implement of np.linspace(start[i], stop[i], num[i], endpoint=False) + axis_grid = [] + for i in range(dim): + a, b, n = start[i], stop[i], num[i] + g = torch.linspace(a, b, n + 1, dtype=torch.float32)[:n] + axis_grid.append(g) + grid = torch.meshgrid(*axis_grid, indexing="ij") # dim x [H, W] + grid = torch.stack(grid, dim=0) # [dim, H, W] + + return grid + + +def build_2d_rope( + seq_len: int, + n_elem: int, + image_infos: list[tuple[slice, tuple[int, int]]] | None = None, + device: torch.device | None = None, + base: int = 10000, + base_rescale_factor: float = 1.0, + return_all_pos: bool = False, +): + assert n_elem % 4 == 0, f"n_elem must be divisible by 4, but got {n_elem}." + + # theta + if base_rescale_factor != 1.0: + base *= base_rescale_factor ** (n_elem / (n_elem - 2)) + theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device).float() / n_elem)) + theta = theta.reshape(1, n_elem // 4, 2) # [1, half_d, 2] + + # position indices + if image_infos is None: + image_infos = [] + + image_infos_list = [image_infos] + sample_seq_lens = [seq_len] + + # Prepare position indices for each sample + x_sections = [] + y_sections = [] + for sample_id, sample_image_infos in enumerate(image_infos_list): + last_pos = 0 + for sec_slice, (h, w) in sample_image_infos: + L = sec_slice.start # start from 0, so image_slice.start is just L + # previous text + if last_pos < L: + y_sections.append(torch.arange(last_pos, L)) + x_sections.append(torch.arange(last_pos, L)) + elif h is None: + # Interleave data has overlapped positions for tokens. + y_sections.append(torch.arange(sec_slice.start, sec_slice.stop)) + x_sections.append(torch.arange(sec_slice.start, sec_slice.stop)) + continue + else: + # Interleave data has overlapped positions for noised image and the successive clean image, + # leading to last_pos (= last text end L + noise w * h) > L (last text end L). + pass + # current image + beta_y = L + (w * h - h) / 2 + beta_x = L + (w * h - w) / 2 + grid = get_meshgrid_nd((beta_y, beta_x), (beta_y + h, beta_x + w)) # [2, h, w] # noqa: N802 + grid = grid.reshape(2, -1) # (y, x) + y_sections.append(grid[0]) + x_sections.append(grid[1]) + # step + last_pos = L + w * h + # final text + y_sections.append(torch.arange(last_pos, sample_seq_lens[sample_id])) + x_sections.append(torch.arange(last_pos, sample_seq_lens[sample_id])) + + x_pos = torch.cat(x_sections).long() + y_pos = torch.cat(y_sections).long() + # If there are overlap positions, we need to remove them. + x_pos = x_pos[:seq_len] + y_pos = y_pos[:seq_len] + all_pos = torch.stack((y_pos, x_pos), dim=1).unsqueeze(1).to(device) # [seq_len, 1, 2] + + # calc rope + idx_theta = (all_pos * theta).reshape(all_pos.shape[0], n_elem // 2).repeat(1, 1) + + cos = torch.cos(idx_theta) + sin = torch.sin(idx_theta) + + if return_all_pos: + return cos, sin, all_pos + + return cos, sin + + +def build_batch_2d_rope( + seq_len: int, + n_elem: int, + image_infos: list[list[tuple[slice, tuple[int, int]]]] | None = None, + device: torch.device | None = None, + base: int = 10000, + base_rescale_factor: float = 1.0, + return_all_pos: bool = False, +): + cos_list, sin_list, all_pos_list = [], [], [] + if image_infos is None: + image_infos = [None] + for i, image_info in enumerate(image_infos): + res = build_2d_rope( + seq_len, + n_elem, + image_infos=image_info, + device=device, + base=base, + base_rescale_factor=base_rescale_factor, + return_all_pos=return_all_pos, + ) + if isinstance(res, tuple) and len(res) == 3: + cos, sin, all_pos = res + elif isinstance(res, tuple) and len(res) == 2: + cos, sin = res + all_pos = None + else: + raise ValueError( + "build_2d_rope must return a tuple of length 2 or 3 " + f"when return_all_pos={return_all_pos}, got: {type(res)} with length " + f"{len(res) if isinstance(res, tuple) else 'N/A'}" + ) + cos_list.append(cos) + sin_list.append(sin) + all_pos_list.append(all_pos) + stacked_cos = torch.stack(cos_list, dim=0) + stacked_sin = torch.stack(sin_list, dim=0) + if return_all_pos: + return stacked_cos, stacked_sin, all_pos_list + + return stacked_cos, stacked_sin + + +def get_full_state_dict(model_path): + files = glob.glob(os.path.join(model_path, "*.safetensors")) + full_sd = {} + for f in files: + full_sd.update(load_file(f)) + return full_sd + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb( + q, k, cos, sin, position_ids=None, unsqueeze_dim=1, mla=False +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): + The query tensor. + k (`torch.Tensor`): + The key tensor. + cos (`torch.Tensor`): + The cosine part of the rotary embedding. + sin (`torch.Tensor`): + The sine part of the rotary embedding. + position_ids (`torch.Tensor`): + The position indices of the tokens corresponding to the query and key tensors. For example, this can be + used to pass offsetted position ids when working with a KV-cache. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + + Returns: + `Tuple[torch.Tensor, torch.Tensor]`: + A tuple comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + if position_ids is not None: + cos = cos[position_ids] + sin = sin[position_ids] + + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + if mla: + b, h, s, d = q.shape + q = q.reshape(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + b, h, s, d = k.shape + k = k.reshape(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def default(value, default_value): + return value if value is not None else default_value + + +class Resolution: + def __init__(self, size, *args): + if isinstance(size, str): + if "x" in size: + size = size.split("x") + size = (int(size[0]), int(size[1])) + else: + size = int(size) + if len(args) > 0: + size = (size, args[0]) + if isinstance(size, int): + size = (size, size) + + self.h = self.height = size[0] + self.w = self.width = size[1] + self.r = self.ratio = self.height / self.width + + def __getitem__(self, idx): + if idx == 0: + return self.h + elif idx == 1: + return self.w + else: + raise IndexError(f"Index {idx} out of range") + + def __str__(self): + return f"{self.h}x{self.w}" + + +class ResolutionGroup: + def __init__(self, base_size=None, step=None, align=1): + self.align = align + self.base_size = base_size + assert base_size % align == 0, f"base_size {base_size} is not divisible by align {align}" + if base_size is not None and not isinstance(base_size, int): + raise ValueError(f"base_size must be None or int, but got {type(base_size)}") + if step is None: + step = base_size // 16 + if step is not None and step > base_size // 2: + raise ValueError(f"step must be smaller than base_size // 2, but got {step} > {base_size // 2}") + + self.step = step + self.data = self._calc_by_step() + + self.ratio = np.array([x.ratio for x in self.data]) + self.attr = ["" for _ in range(len(self.data))] + self.prefix_space = 0 + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx] + + def __repr__(self): + prefix = self.prefix_space * " " + prefix_close = (self.prefix_space - 4) * " " + res_str = f"ResolutionGroup(base_size={self.base_size}, step={self.step}, data=" + attr_maxlen = max([len(x) for x in self.attr] + [5]) + res_str += ( + f"\n{prefix}ID: height width ratio {' ' * max(0, attr_maxlen - 4)}count h/16 w/16 tokens\n{prefix}" + ) + res_str += ("\n" + prefix).join( + [ + f"{i:2d}: ({x.h:4d}, {x.w:4d}) {self.ratio[i]:.4f} {self.attr[i]:>{attr_maxlen}s} " + f"({x.h // 16:3d}, {x.w // 16:3d}) {x.h // 16 * x.w // 16:6d}" + for i, x in enumerate(self.data) + ] + ) + res_str += f"\n{prefix_close})" + return res_str + + def _calc_by_step(self): + assert self.align <= self.step, f"align {self.align} must be smaller than step {self.step}" + + min_height = self.base_size // 2 + min_width = self.base_size // 2 + max_height = self.base_size * 2 + max_width = self.base_size * 2 + + resolutions = [Resolution(self.base_size, self.base_size)] + + cur_height, cur_width = self.base_size, self.base_size + while True: + if cur_height >= max_height and cur_width <= min_width: + break + + cur_height = min(cur_height + self.step, max_height) + cur_width = max(cur_width - self.step, min_width) + resolutions.append(Resolution(cur_height // self.align * self.align, cur_width // self.align * self.align)) + + cur_height, cur_width = self.base_size, self.base_size + while True: + if cur_height <= min_height and cur_width >= max_width: + break + + cur_height = max(cur_height - self.step, min_height) + cur_width = min(cur_width + self.step, max_width) + resolutions.append(Resolution(cur_height // self.align * self.align, cur_width // self.align * self.align)) + + resolutions = sorted(resolutions, key=lambda x: x.ratio) + + return resolutions + + def get_target_size(self, width, height): + ratio = height / width + idx = np.argmin(np.abs(self.ratio - ratio)) + reso = self.data[idx] + return reso.w, reso.h + + def get_base_size_and_ratio_index(self, width, height): + ratio = height / width + idx = np.argmin(np.abs(self.ratio - ratio)) + return self.base_size, idx + + +class ImageInfo: + """Class to store image information for processing and generation.""" + + def __init__( + self, + image_type: str = None, + image_tensor: torch.Tensor = None, + image_width: int = None, + image_height: int = None, + token_width: int = None, + token_height: int = None, + image_token_length: int = None, + base_size: int = None, + ratio_index: int = None, + **kwargs, + ): + self.image_type = image_type + self.image_tensor = image_tensor + self.image_width = image_width + self.w = image_width + self.image_height = image_height + self.h = image_height + self.token_width = token_width + self.tk_w = token_width + self.token_height = token_height + self.tk_h = token_height + self.image_token_length = default( + image_token_length, + token_width * token_height if token_width is not None and token_height is not None else None, + ) + self.base_size = base_size + self.ratio_index = ratio_index + + self.add_timestep_token = kwargs.get("add_timestep_token", True) + self.add_guidance_token = kwargs.get("add_guidance_token", False) + self.use_front_boi_token = kwargs.get("use_front_boi_token", True) + self.add_image_shape_token = kwargs.get("add_image_shape_token", True) + + def __getitem__(self, key: str) -> Any: + """Allow dictionary-like access to attributes.""" + if hasattr(self, key): + return getattr(self, key) + raise KeyError(f"Key '{key}' not found in ImageInfo") + + def __setitem__(self, key: str, value: Any) -> None: + """Allow dictionary-like assignment to attributes.""" + if hasattr(self, key): + setattr(self, key, value) + else: + raise KeyError(f"Key '{key}' not found in ImageInfo") + + def __contains__(self, key: str) -> bool: + """Check if the key exists in the ImageInfo object.""" + return hasattr(self, key) + + def __repr__(self): + return ( + f"ImageInfo(image_type={self.image_type}, image_tensor={self.image_tensor}, " + f"image_width={self.image_width}, image_height={self.image_height}, " + f"token_width={self.token_width}, token_height={self.token_height}, " + f"image_token_length={self.image_token_length}, " + f"base_size={self.base_size}, ratio_index={self.ratio_index}" + ) + + @property + def meta_info(self): + # Used for image sections of tkwrapper.encode_general() + if self.image_type in ["vae", "gen_image"]: + return dict( + token_length=self.image_token_length, + add_timestep_token=self.add_timestep_token, + add_guidance_token=self.add_guidance_token, + use_front_boi_token=self.use_front_boi_token, + add_image_shape_token=self.add_image_shape_token, + base_size=self.base_size, + ratio_idx=self.ratio_index, + # for rope 2d + token_height=self.token_height, + token_width=self.token_width, + # for bc + image_height=self.image_height, + image_width=self.image_width, + ) + else: + raise ValueError(f"Unknown image type '{self.image_type}'") + + +class ImageTensor(torch.Tensor): + # This class is just for type hinting purposes. Attribute `i` should be defined + # as an instance attribute of the torch.Tensor instance, like: tensor.i = ImageInfo(...) + i: ImageInfo + vision_encoder_kwargs: dict + + +class JointImageInfo: + def __init__(self, vae_image_info: ImageInfo, vision_image_info: ImageInfo, vision_encoder_kwargs: dict = None): + self.vae_image_info = vae_image_info + self.vision_image_info = vision_image_info + self.vision_encoder_kwargs = vision_encoder_kwargs + + # Define key attributes to align with ImageInfo for uniformity + self.image_type = "joint_image" + self.image_token_length = vae_image_info.image_token_length + vision_image_info.image_token_length + + self.add_timestep_token = vae_image_info.add_timestep_token + self.use_front_boi_token = vae_image_info.use_front_boi_token + self.add_image_shape_token = vae_image_info.add_image_shape_token + + def __repr__(self): + return f"JointImageInfo(vae_image={self.vae_image_info}, vision_image={self.vision_image_info})" + + @property + def meta_info(self): + # Used for image sections of tkwrapper.encode_general() + return dict( + token_length=[self.vae_image_info.image_token_length, self.vision_image_info.image_token_length], + add_timestep_token=self.add_timestep_token, + use_front_boi_token=self.use_front_boi_token, + add_image_shape_token=self.add_image_shape_token, + base_size=self.vae_image_info.base_size, + ratio_idx=self.vae_image_info.ratio_index, + # for rope 2d + token_height=[self.vae_image_info.token_height, self.vision_image_info.token_height], + token_width=[self.vae_image_info.token_width, self.vision_image_info.token_width], + # for bc + image_height=[self.vae_image_info.image_height, self.vision_image_info.image_height], + image_width=[self.vae_image_info.image_width, self.vision_image_info.image_width], + ) + + @property + def num_special_tokens(self): + return ( + 2 # + + + (1 if self.add_timestep_token else 0) + + (2 if self.add_image_shape_token else 0) + + 1 # + ) + + def copy(self, copy_image_tensor=True): + if copy_image_tensor and ( + self.vae_image_info.image_tensor is None or self.vision_image_info.image_tensor is None + ): + raise ValueError("image_tensor is None, cannot copy") + return JointImageInfo( + self.vae_image_info.copy(copy_image_tensor), + self.vision_image_info.copy(copy_image_tensor), + self.vision_encoder_kwargs, + ) + + def zeros_(self): + self.vae_image_info.zeros_() + self.vision_image_info.zeros_() + + +class JointImage: + def __init__(self, vae_image: ImageTensor, vision_image: ImageTensor): + self.vae_image = vae_image + self.vision_image = vision_image + self.i = JointImageInfo(vae_image.i, vision_image.i) + + +class Config: + def __init__(self, config): + if config is not None: + for key, value in config.items(): + setattr(self, key, value) + + +class LightProjector(nn.Module): + def __init__(self, config): + config = Config(config) + super().__init__() + + if config.projector_type == "linear": + modules = nn.Linear(config.input_dim, config.n_embed) + + elif config.projector_type == "mlp_gelu": + modules = [nn.Linear(config.input_dim, config.n_embed)] + for _ in range(1, config.depth): + modules.append(nn.GELU()) + modules.append(nn.Linear(config.n_embed, config.n_embed)) + modules = nn.Sequential(*modules) + + else: + raise ValueError(f"Unknown projector type: {config.projector_type}") + + self.layers = modules + + +class HunYuanRotary2DEmbedder: + r""" + A RoPE wrapper specifically designed for HunYuan-Image attention. + + This class implements Rotary Position Embedding (RoPE) specifically optimized for + the HunYuan-Image model's attention mechanism. It handles the application of rotary + embeddings to query and key tensors with support for custom position embeddings + and attention metadata. + + Example: + ```python + embedder = HunYuanRotaryEmbedder(num_heads=num_h, num_kv_heads=num_kv, head_dim=h_d) + q, k = embedder(q, k, hidden_states, custom_pos_emb, first_step, attn_meta) + ``` + + Args: + num_heads (`int`): + Number of attention heads in the model. + num_kv_heads (`int`): + Number of key-value heads (for grouped-query attention). + head_dim (`int`): + Dimension of each attention head. + + Methods: + __call__: Applies rotary position embedding to query and key tensors. + """ + + def __init__(self, num_heads: int, num_kv_heads: int, head_dim: int): + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + self.custom_pos_emb: tuple[torch.Tensor, torch.Tensor] | None = None + self.rope = RotaryEmbedding(is_neox_style=True) + + def _prepare_cos_sin( + self, + custom_pos_emb: tuple[torch.Tensor, torch.Tensor], + first_step: bool, + device: torch.device, + ): + """Returns cos/sin on the target device based on first_step and caching strategy.""" + if first_step: + cos_input, sin_input = custom_pos_emb + cos = cos_input.to(device) + sin = sin_input.to(device) + self.custom_pos_emb = None + else: + if self.custom_pos_emb is None: + cos_input, sin_input = custom_pos_emb + cos = cos_input.to(device) + sin = sin_input.to(device) + self.custom_pos_emb = (cos, sin) + else: + cos, sin = self.custom_pos_emb + return cos, sin + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + hidden_states: torch.Tensor, + custom_pos_emb: tuple[torch.Tensor, torch.Tensor], + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states_shape = hidden_states.shape + hidden_states = hidden_states.reshape(-1, hidden_states_shape[-1]) + if kwargs.get("mode", "gen_text") != "gen_image": + return q, k + + first_step = kwargs.get("first_step", False) + device = q.device + # 1. Prepare cos/sin + cos, sin = self._prepare_cos_sin(custom_pos_emb, first_step, device) + + # 2. Shape validation + query_lens: list[int] = kwargs.get("query_lens") + bs = len(query_lens) + q_len = query_lens[0] + assert hidden_states.shape[0] == bs * q_len, f"{hidden_states.shape[0]} != {bs * q_len}" + + # 3. Reshape + transpose for apply_rotary_pos_emb + # Assume q shape [B*L, H*D] -> [2, L, H, D] -> [2, H, L, D] + q = q.reshape(bs, q_len, self.num_heads, self.head_dim) + k = k.reshape(bs, q_len, self.num_kv_heads, self.head_dim) + + q = self.rope(q.to(torch.float32), cos, sin) + k = self.rope(k.to(torch.float32), cos, sin) + + # 5. Restore original shape + convert to bfloat16 + q = q.reshape(hidden_states.shape[0], self.num_heads * self.head_dim).to(torch.bfloat16) + k = k.reshape(hidden_states.shape[0], self.num_kv_heads * self.head_dim).to(torch.bfloat16) + hidden_states = hidden_states.reshape(hidden_states_shape) + return q, k + + +class ImageKVCacheManager: + """ + Manages specialized caching and updating of KV-Cache for image tokens in multimodal models. + """ + + def __init__(self, image_token_len: int = 4097): + """ + Args: + image_token_len: Number of tokens per image (including special placeholders), + default 4097 (timestamp + 4096 image tokens). + """ + self.image_token_len: int = image_token_len + self.image_kv_cache: tuple[torch.Tensor, torch.Tensor] = None + + def _save_image_kv_caches( + self, + key: torch.Tensor, + value: torch.Tensor, + seq_len: int, + ) -> None: + bs, q_len, num_kv_heads, head_dim = key.shape + assert q_len == seq_len, f"for first-step, {q_len} != {seq_len}" + + key = key.reshape(-1, num_kv_heads, head_dim) + value = value.reshape(-1, num_kv_heads, head_dim) + + cached_prompt_len = seq_len - self.image_token_len - 1 + cached_key = [key[:cached_prompt_len], key[seq_len - 1 : seq_len]] + cached_value = [value[:cached_prompt_len], value[seq_len - 1 : seq_len]] + + if bs > 1: + assert bs == 2, "for cfg case, bs must be 2" + cached_key.append(key[seq_len : seq_len + cached_prompt_len]) + cached_key.append(key[-1:]) + + cached_value.append(value[seq_len : seq_len + cached_prompt_len]) + cached_value.append(value[-1:]) + + cached_key = torch.cat(cached_key, dim=0) + cached_value = torch.cat(cached_value, dim=0) + self.image_kv_cache_map = (cached_key, cached_value) + + def _update_image_kv_caches( + self, + key: torch.Tensor, + value: torch.Tensor, + seq_len: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + cached_key, cached_value = self.image_kv_cache_map + bs, q_len, num_kv_heads, head_dim = key.shape + + cached_prompt_len = cached_key.shape[0] // bs - 1 + assert (cached_prompt_len + 1) == (seq_len - q_len), f"{cached_prompt_len + 1} != {seq_len - q_len}" + + key = key.reshape(-1, num_kv_heads, head_dim) + value = value.reshape(-1, num_kv_heads, head_dim) + + new_key = [ + cached_key[:cached_prompt_len], + key[:q_len], + cached_key[cached_prompt_len : cached_prompt_len + 1], + ] + new_value = [ + cached_value[:cached_prompt_len], + value[:q_len], + cached_value[cached_prompt_len : cached_prompt_len + 1], + ] + + if bs > 1: + assert bs == 2, "for cfg case, bs must be 2" + new_key.append(cached_key[cached_prompt_len + 1 : cached_prompt_len + 1 + cached_prompt_len]) + new_key.append(key[q_len:]) + new_key.append(cached_key[-1:]) + + new_value.append(cached_value[cached_prompt_len + 1 : cached_prompt_len + 1 + cached_prompt_len]) + new_value.append(value[q_len:]) + new_value.append(cached_value[-1:]) + + new_key = torch.cat(new_key, dim=0) + new_value = torch.cat(new_value, dim=0) + new_key = new_key.reshape(bs, seq_len, num_kv_heads, head_dim) + new_value = new_value.reshape(bs, seq_len, num_kv_heads, head_dim) + + return new_key.contiguous(), new_value.contiguous() + + def __call__( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + self.image_token_len = kwargs.get("num_image_tokens") + first_step = kwargs.get("first_step") + + query_lens = kwargs.get("query_lens") + seq_lens = kwargs.get("seq_lens") + bs = len(query_lens) + q_len = query_lens[0] + seq_len = seq_lens[0] + assert query.shape[0] == bs * q_len, f"{query.shape[0]} != {bs * q_len}" + + head_num_per_rank = query.shape[1] + kv_head_num_per_rank = key.shape[1] + repeat_num = head_num_per_rank // kv_head_num_per_rank + head_dim = query.shape[2] + + query = query.reshape(bs, q_len, head_num_per_rank, head_dim) + key = key.reshape(bs, q_len, kv_head_num_per_rank, head_dim) + value = value.reshape(bs, q_len, kv_head_num_per_rank, head_dim) + + if first_step: + self.image_kv_cache_map = None + self._save_image_kv_caches(key, value, seq_len) + else: + key, value = self._update_image_kv_caches(key, value, seq_len) + + query = query.transpose(1, 2).contiguous() + key = key.transpose(1, 2).contiguous() + value = value.transpose(1, 2).contiguous() + + key = repeat_kv(key, repeat_num) + value = repeat_kv(value, repeat_num) + + attention_mask = attention_mask.contiguous() + + attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=0.0) + + attn_output = attn_output.transpose(1, 2).contiguous() # [bs, q_len, heads, head_dim] + attn_output = attn_output.reshape(bs * q_len, head_num_per_rank, head_dim) + return attn_output + + +@dataclass +class CausalMMOutputWithPast(CausalLMOutputWithPast): + diffusion_prediction: torch.Tensor | None = None + + +class HunyuanImage3Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`HunyuanImage3Model`]. It is used to instantiate + an Hunyuan model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the Hunyuan-7B. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + vocab_size (`int`, *optional*, defaults to 32000): + Vocabulary size of the Hunyuan Image 3 model. Defines the number of different tokens that can be + represented by the `inputs_ids` passed when calling [`HunyuanImage3Model`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 11008): + Dimension of the MLP representations or shared MLP representations. + moe_intermediate_size (`int` or `List`, *optional*, defaults to 11008): + Dimension of the MLP representations in MoE. Use a list if you want a different size per layer. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + num_key_value_heads (`int`, *optional*): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): + The non-linear activation function (function or string) in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 2048): + The maximum sequence length that this model might ever be used with. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + rms_norm_eps (`float`, *optional*, defaults to 1e-06): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if `config.is_decoder=True`. + pad_token_id (`int`, *optional*): + Padding token id. + bos_token_id (`int`, *optional*, defaults to 1): + Beginning of stream token id. + eos_token_id (`int`, *optional*, defaults to 2): + End of stream token id. + pretraining_tp (`int`, *optional*, defaults to 1): + Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this + document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is + necessary to ensure exact reproducibility of the pretraining results. Please refer to [this + issue](https://github.com/pytorch/pytorch/issues/76232). + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings + rope_theta (`float`, *optional*, defaults to 10000.0): + The base period of the RoPE embeddings. + rope_scaling (`Dict`, *optional*): + Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling + strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is + `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update + `max_position_embeddings` to the expected new maximum. See the following thread for more information on how + these scaling strategies behave: + https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an + experimental feature, subject to breaking API changes in future versions. + attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): + Whether to use a bias in the query, key, value and output projection layers during self-attention. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + use_qk_norm (`bool`, *optional*, defaults to `False`): + Whether query and key in attention use norm + use_cla (`bool`, *optional*, defaults to `False`): + Whether to use CLA in attention + cla_share_factor (`int`, *optional*, defaults to 1): + The share factor of CLA + num_experts (`int` or `List`, *optional*, defaults to 1): + The number of experts for moe. If it is a list, it will be used as the number of experts for each layer. + num_shared_expert (`int` or `List`, *optional*, defaults to 1): + The number of shared experts for moe. If it is a list, it will be used as the number of shared experts + for each layer. + moe_topk (`int` or `List`, *optional*, defaults to 1): + The topk value for moe. If it is a list, it will be used as the topk value for each layer. + capacity_factor (Not used) (`float` or `List`, *optional*, defaults to 1.0): + The capacity factor for moe. If it is a list, it will be used as the capacity factor for each layer. + moe_layer_num_skipped (`int`, *optional*, defaults to 0): + First moe_layer_num_skipped layers do not use MoE. + """ + + model_type = "Hunyuan" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=290943, + hidden_size=4096, + intermediate_size: int = 11008, + moe_intermediate_size: int | list = None, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + attention_head_dim=None, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-5, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + eod_token_id=3, + im_start_id=4, + im_end_id=5, + text_start_id=6, + text_end_id=7, + image_token_id=8, + video_start_id=9, + video_end_id=10, + im_newline_id=11, + mask_init_id=12, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + mlp_bias=False, + attention_dropout=0.0, + use_qk_norm=False, + use_rotary_pos_emb=True, + use_cla=False, + cla_share_factor=1, + norm_type="hf_rms", + num_experts: int | list = 1, + use_mixed_mlp_moe=False, + num_shared_expert: int | list = 1, + moe_topk: int | list = 1, + capacity_factor: float | list = 1.0, + moe_drop_tokens=False, + moe_random_routing_dropped_token=False, + use_mla=False, + kv_lora_rank=512, + q_lora_rank=1536, + qk_rope_head_dim=64, + v_head_dim=128, + qk_nope_head_dim=128, + moe_layer_num_skipped=0, + norm_topk_prob=True, + routed_scaling_factor=1.0, + group_limited_greedy=False, + n_group=None, + topk_group=None, + add_classification_head=False, + class_num=0, + pool_type="last", + pad_id=-1, + # Added + moe_impl="eager", + vae_downsample_factor=(16, 16), # (h, w) + img_proj_type="unet", + patch_size=1, + patch_embed_hidden_dim=1024, + image_base_size=1024, + vae=None, + vit=None, + vit_processor=None, + vit_aligner=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.moe_impl = moe_impl + self.num_experts = num_experts + self.use_mixed_mlp_moe = use_mixed_mlp_moe + self.num_shared_expert = num_shared_expert + self.moe_topk = moe_topk + self.capacity_factor = capacity_factor + self.moe_drop_tokens = moe_drop_tokens + self.moe_random_routing_dropped_token = moe_random_routing_dropped_token + + if attention_head_dim is not None: + self.attention_head_dim = attention_head_dim + else: + self.attention_head_dim = self.hidden_size // num_attention_heads + + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.mlp_bias = mlp_bias + self.attention_dropout = attention_dropout + self.use_qk_norm = use_qk_norm + self.use_rotary_pos_emb = use_rotary_pos_emb + self.use_cla = use_cla + self.cla_share_factor = cla_share_factor + self.norm_type = norm_type + # MLA args + self.use_mla = use_mla + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.v_head_dim = v_head_dim + + # DeepSeek related args + self.moe_layer_num_skipped = moe_layer_num_skipped + self.norm_topk_prob = norm_topk_prob + self.routed_scaling_factor = routed_scaling_factor + self.group_limited_greedy = group_limited_greedy + self.n_group = n_group + self.topk_group = topk_group + self.add_classification_head = add_classification_head + self.class_num = class_num + self.pool_type = pool_type + self.pad_id = pad_id + + if self.class_num is not None: + self.dense_list = [self.hidden_size, self.class_num] + + # ViT args + self.vit = vit + self.vit_processor = vit_processor + self.vit_aligner = vit_aligner + + # Image Gen args + self.vae = vae + self.vae_downsample_factor = vae_downsample_factor + self.img_proj_type = img_proj_type + self.patch_size = patch_size + self.patch_embed_hidden_dim = patch_embed_hidden_dim + self.image_base_size = image_base_size + + # token id + self.eod_token_id = eod_token_id + self.im_start_id = im_start_id + self.im_end_id = im_end_id + self.text_start_id = text_start_id + self.text_end_id = text_end_id + self.image_token_id = image_token_id + self.video_start_id = video_start_id + self.video_end_id = video_end_id + self.im_newline_id = im_newline_id + self.mask_init_id = mask_init_id + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class HunyuanImage3ImageProcessor: + def __init__(self, config): + self.config = config + + self.reso_group = ResolutionGroup(base_size=config.image_base_size) + self.vae_processor = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize([0.5], [0.5]), # transform to [-1, 1] + ] + ) + self.vision_encoder_processor = Siglip2ImageProcessorFast.from_dict(config.vit_processor) + + def build_image_info(self, image_size): + # parse image size (HxW, H:W, or ) + if isinstance(image_size, str): + if image_size.startswith("")) + reso = self.reso_group[ratio_index] + image_size = reso.height, reso.width + elif "x" in image_size: + image_size = [int(s) for s in image_size.split("x")] + elif ":" in image_size: + image_size = [int(s) for s in image_size.split(":")] + else: + raise ValueError( + f"`image_size` should be in the format of 'HxW', 'H:W' or , got {image_size}." + ) + assert len(image_size) == 2, f"`image_size` should be in the format of 'HxW', got {image_size}." + elif isinstance(image_size, (list, tuple)): + assert len(image_size) == 2 and all(isinstance(s, int) for s in image_size), ( + f"`image_size` should be a tuple of two integers or a string in the format of 'HxW', got {image_size}." + ) + else: + raise ValueError( + f"`image_size` should be a tuple of two integers or a string in the format of 'WxH', got {image_size}." + ) + image_width, image_height = self.reso_group.get_target_size(image_size[1], image_size[0]) + token_height = image_height // (self.config.vae_downsample_factor[0] * self.config.patch_size) + token_width = image_width // (self.config.vae_downsample_factor[1] * self.config.patch_size) + base_size, ratio_idx = self.reso_group.get_base_size_and_ratio_index(image_size[1], image_size[0]) + image_info = ImageInfo( + image_type="gen_image", + image_width=image_width, + image_height=image_height, + token_width=token_width, + token_height=token_height, + base_size=base_size, + ratio_index=ratio_idx, + ) + return image_info + + +class HunYuanMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + prefix: str = "", + reduce_results: bool = True, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + input_size=hidden_size, + output_sizes=[intermediate_size] * 2, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + input_size=intermediate_size, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + reduce_results=reduce_results, + ) + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.") + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class HunYuanSparseMoeBlock(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: QuantizationConfig | None = None, + layer_id: int = -1, + prefix: str = "", + enable_eplb: bool = False, + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + self.n_routed_experts = config.num_experts + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than the number of experts {config.num_experts}." + ) + + # Get layer_id topk if config.moe_topk is a list + if isinstance(config.moe_topk, list): + assert layer_id >= 0 + assert len(config.moe_topk) > layer_id + top_k = config.moe_topk[layer_id] + else: + top_k = config.moe_topk + + # If it is moe, moe_intermediate_size is preferred + intermediate_size = config.intermediate_size + if config.moe_intermediate_size is not None: + intermediate_size = ( + config.moe_intermediate_size + if isinstance(config.moe_intermediate_size, int) + else config.moe_intermediate_size[layer_id] + ) + + self.enable_eplb = False + self.n_logical_experts = self.n_routed_experts + self.n_redundant_experts = 0 + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=None, + prefix=f"{prefix}.gate", + ) + if config.use_mixed_mlp_moe > 0: + # Get layer_id num_shared_expert if config.num_shared_expert is + # a list. + if isinstance(config.num_shared_expert, list): + assert layer_id >= 0 + assert len(config.num_shared_expert) > layer_id + num_shared_expert = config.num_shared_expert[layer_id] + else: + num_shared_expert = config.num_shared_expert + + self.shared_mlp = HunYuanMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size * num_shared_expert, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.shared_mlp", + ) + else: + self.shared_mlp = None + + self.experts = HunyuanFusedMoE( + shared_experts=self.shared_mlp, + num_experts=self.n_routed_experts, + top_k=top_k, + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + reduce_results=False, + renormalize=top_k > 1, + quant_config=quant_config, + prefix=f"{prefix}.experts", + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + pcp_size=1, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + # NOTE: hidden_states can have either 1D or 2D shape. + orig_shape = hidden_states.shape + hidden_dim = hidden_states.shape[-1] + hidden_states = hidden_states.view(-1, hidden_dim) + + # router_logits: (num_tokens, n_experts) + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=router_logits) + if self.shared_mlp is not None: + 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) + + return final_hidden_states.view(orig_shape) + + +class HunYuanAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + rope_theta: float = 10000, + rope_scaling: dict[str, Any] | None = None, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + bias: bool = False, + cache_config: CacheConfig | None = None, + prefix: str = "", + layer_id: int = -1, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + tp_size = get_tensor_model_parallel_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.num_key_value_heads = ( + config.num_key_value_heads // tp_size if config.num_key_value_heads else self.num_heads + ) + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + if hasattr(config, "head_dim") and config.head_dim: + self.head_dim = config.head_dim + elif hasattr(config, "attention_head_dim"): + self.head_dim = config.attention_head_dim + else: + self.head_dim = self.hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.use_qk_norm = getattr(config, "use_qk_norm", False) + self.layer_id = layer_id + + self.qkv_proj = QKVParallelLinear( + hidden_size=hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + + self.o_proj = RowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=hidden_size, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + if rope_scaling is not None: + # for t2t, fallback rotary_emb type from 'custom' to default. + rope_scaling["rope_type"] = "default" + self.rotary_emb = get_rope( + head_size=self.head_dim, + max_position=max_position_embeddings, + rope_parameters=rope_scaling, + is_neox_style=True, + ) + # which attention, from vllm or omni? + # try use Atten from diffusion + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=self.scaling, + num_kv_heads=self.num_kv_heads, + # cache_config=cache_config, + # quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + # default image_token_len = timestamp + 4096*image_tokes + self.image_attn = ImageKVCacheManager(image_token_len=4097) + self.image_rope2d_emb = HunYuanRotary2DEmbedder( + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + ) + + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + kv_states: tuple[torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + custom_pos_emb: tuple[torch.FloatTensor] | None = None, + **kwargs, + ) -> torch.Tensor: + bsz, q_len, _ = hidden_states.size() + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + past_key_value: Cache | None = kwargs.get("past_key_value", None) + if past_key_value is not None: + position_ids = kwargs.get("position_ids") + key_states = k.reshape(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = v.reshape(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + cache_kwargs = {"cache_position": position_ids} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_id, cache_kwargs) + # for image_generation + if kwargs.get("mode", "gen_text") == "gen_image": + # assert positions is None, "positions should be None for image attention" + q, k = self.image_rope2d_emb(q, k, hidden_states, custom_pos_emb, **kwargs) + else: + q, k = self.rotary_emb(positions, q, k) + if self.use_qk_norm: + q = self.query_layernorm(q.view(-1, self.num_heads, self.head_dim).contiguous()) + k = self.key_layernorm(k.view(-1, self.num_kv_heads, self.head_dim).contiguous()) + # for image_generation + if kwargs.get("mode", "gen_text") == "gen_image": + attn_output = self.image_attn(q, k, v, attention_mask=attention_mask, **kwargs) + else: + attn_output = self.attn(q, k, v) + # For o_proj + attn_output = attn_output.view(q.shape[0], -1) + output, _ = self.o_proj(attn_output) + output = output.reshape(bsz, q_len, -1) + 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): + from vllm.model_executor.layers.fused_moe.layer import get_forward_context + + ctx = get_forward_context() + if not ctx.remaining_moe_layers: + import re + + moe_names = [name for name in ctx.no_compile_layers.keys() if ".mlp.experts" in name] + + def get_layer_num(name): + match = re.search(r"layers\.(\d+)\.mlp", name) + return int(match.group(1)) if match else -1 + + moe_names.sort(key=get_layer_num, reverse=True) + ctx.remaining_moe_layers.extend(moe_names) + return super().forward(hidden_states, router_logits) + + +class HunyuanImage3DecoderLayer(nn.Module): + def __init__(self, config: HunyuanImage3Config, layer_idx: int, prefix: str = ""): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + self.intermediate_size = ( + config.intermediate_size + if isinstance(config.intermediate_size, int) + else config.intermediate_size[layer_idx] + ) + rope_theta = getattr(config, "rope_theta", 10000) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None and getattr(config, "original_max_position_embeddings", None): + rope_scaling["original_max_position_embeddings"] = config.original_max_position_embeddings + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + attention_bias = getattr(config, "attention_bias", False) or getattr(config, "bias", False) + + cla_factor = _get_cla_factor(config) + attention_type = ( + AttentionType.ENCODER_DECODER if layer_idx >= 0 and layer_idx % cla_factor != 0 else AttentionType.DECODER + ) + if attention_type == AttentionType.DECODER: + self.self_attn = HunYuanAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=getattr(config, "num_key_value_heads", config.num_attention_heads), + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=None, + bias=attention_bias, + cache_config=None, + prefix=f"{prefix}.self_attn", + layer_id=layer_idx, + ) + else: + raise RuntimeError(f"Unsupported attention type: {attention_type}") + + if ( + (isinstance(config.num_experts, int) and config.num_experts > 1) + or (isinstance(config.num_experts, list) and max(config.num_experts) > 1) + ) and layer_idx >= config.moe_layer_num_skipped: + self.mlp = HunYuanSparseMoeBlock(config, layer_id=layer_idx, prefix=f"{prefix}.mlp") + else: + self.mlp = HunYuanMLP(self.hidden_size, self.intermediate_size, config.hidden_act) + + 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, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_value: tuple[torch.Tensor] | None = None, + output_attentions: bool | None = False, + use_cache: bool | None = False, + custom_pos_emb: tuple[torch.FloatTensor] | None = None, + **kwargs, + ) -> tuple[torch.FloatTensor | Any]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): + attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, + query_sequence_length, key_sequence_length)` if default attention is used. + position_ids (`torch.LongTensor`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + custom_pos_emb (`Tuple[torch.FloatTensor]`, *optional*): custom position embedding for rotary + position embedding + """ + if "padding_mask" in kwargs: + logger.warning( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use " + "`attention_mask` instead.`" + ) + + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + custom_pos_emb=custom_pos_emb, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +class HunyuanImage3PreTrainedModel(PreTrainedModel): + config_class = HunyuanImage3Config + base_model_prefix = "" + supports_gradient_checkpointing = True + _no_split_modules = ["HunyuanImage3DecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class HunyuanImage3Model(nn.Module): + def __init__(self, config: HunyuanImage3Config, prefix: str = ""): + super().__init__() + quant_config = None + lora_config = None + self.num_redundant_experts = 0 + self.config = config + self.quant_config = quant_config + self.padding_idx = config.pad_token_id + lora_vocab = (lora_config.lora_extra_vocab_size * (lora_config.max_loras or 1)) if lora_config else 0 + self.vocab_size = config.vocab_size + lora_vocab + self.org_vocab_size = config.vocab_size + self.wte = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank): + self.embed_tokens = VocabParallelEmbedding( + self.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + quant_config=quant_config, + ) + else: + self.embed_tokens = PPMissingLayer() + # maybe should not use make_layers() + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: HunyuanImage3DecoderLayer( + config=config, + layer_idx=int(prefix.split(".")[-1]), + prefix=prefix, + ), + prefix=f"{prefix}.layers", + ) + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = PPMissingLayer() + + def _split_qkv_weight(self, qkv: torch.Tensor): + num_attention_heads = self.config.num_attention_heads + num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads) + num_key_value_groups = num_attention_heads // num_kv_heads + hidden_size = self.config.hidden_size + + if hasattr(self.config, "head_dim"): + attention_head_dim = self.config.head_dim + elif hasattr(self.config, "attention_head_dim"): + attention_head_dim = self.config.attention_head_dim + else: + attention_head_dim = self.config.hidden_size // num_attention_heads + + qkv = qkv.reshape(num_kv_heads, num_key_value_groups + 2, attention_head_dim, hidden_size) + q, k, v = torch.split(qkv, (num_key_value_groups, 1, 1), dim=1) + q = q.reshape(-1, hidden_size) + k = k.reshape(-1, hidden_size) + v = v.reshape(-1, hidden_size) + return torch.concat((q, k, v)) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + if _is_moe(self.config): + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + fused_moe_expert_mapping = HunyuanFusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + num_redundant_experts=self.num_redundant_experts, + ) + expert_weights_remapping = { + "gate_proj": ("gate_and_up_proj", 1, 2), + "up_proj": ("gate_and_up_proj", 0, 2), + } + return fused_moe_expert_mapping, expert_weights_remapping + else: + return [], {} + + # rename for delay load + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + cla_factor = _get_cla_factor(self.config) + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + num_attention_heads = self.config.num_attention_heads + num_kv_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads) + split_params_mapping = [ + (".gate_up_proj", ".gate_and_up_proj", 2, [(1, 1), (0, 1)], None), + ( + ".qkv_proj", + ".qkv_proj", + num_attention_heads + num_kv_heads * 2, + [("q", num_attention_heads), ("k", num_kv_heads), ("v", num_kv_heads)], + self._split_qkv_weight, + ), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + expert_params_mapping, expert_weights_remapping = self.get_expert_mapping() + + # List of unexpected keywords in weight names + unexpected_keywords = [ + "vae", + "vision_aligner", + "vision_model", + "final_layer", + "patch_embed", + "timestep_emb", + "time_embed", + "time_embed_2", + "guidance_emb", + "timestep_r_emb", + ] + + def contains_unexpected_keyword(name, keywords): + """Check if the name contains any unexpected keywords""" + for keyword in keywords: + if keyword in name: + return True + return False + + for name, loaded_weight in weights: + # print(f"Loading weight name: {name}, tp_rank: {tp_rank}", flush=True) + if contains_unexpected_keyword(name, unexpected_keywords): + print(f"Skipping unexpected weight name: {name}") + continue + if "rotary_emb.inv_freq" in name: + continue + if "gate_proj_bias" in name: + name = name.replace("gate_proj_bias", "gate_proj.bias") + if "up_proj_bias" in name: + name = name.replace("up_proj_bias", "up_proj.bias") + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + # With tie_word_embeddings, we can skip lm_head.weight + # The weight might appear unnecessarily in the files if the model is + # processed with quantization, LoRA, fine-tuning, etc. + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): + # Loading kv cache scales for compressed-tensors quantization + param = params_dict[scale_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + loaded_weight = loaded_weight[0] + weight_loader(param, loaded_weight) + continue + + is_found = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + # cross layer only have q_proj, skip qkv pack + if weight_name == ".q_proj": + match = re.search(r"layers\.\d+", name) + if match: + layer_id = int(match.group(0).split(".")[-1]) + if cla_factor > 1 and layer_id % cla_factor != 0: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + is_found = True + break + if is_found: + continue + + for ( + param_name, + weight_name, + den, + split_param, + func, + ) in split_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if is_pp_missing_parameter(name, self): + continue + + assert loaded_weight.shape[0] % den == 0 + units = loaded_weight.shape[0] // den + param = params_dict[name] + weight_loader = param.weight_loader + offset = 0 + for shard_id, num in split_param: + new_offset = offset + num * units + if func: + weight_loader(param, func(loaded_weight)[offset:new_offset], shard_id) + else: + weight_loader(param, loaded_weight[offset:new_offset], shard_id) + offset = new_offset + + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + is_expert_weight = False + is_found = False + found_num = 0 + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + offset = 0 + den = 1 + for ( + mapped_weight_substr, + origin_weight_info, + ) in expert_weights_remapping.items(): + if mapped_weight_substr in weight_name: + origin_weight_name, offset, den = origin_weight_info + weight_name = weight_name.replace(mapped_weight_substr, origin_weight_name) + break + if weight_name not in name: + continue + is_expert_weight = True + name_mapped = name.replace(weight_name, param_name) + found_num += 1 + if is_pp_missing_parameter(name_mapped, self): + continue + param = params_dict[name_mapped] + weight_loader = cast(Callable[..., bool], param.weight_loader) + assert loaded_weight.shape[0] % den == 0 + units = loaded_weight.shape[0] // den + + success = weight_loader( + param, + loaded_weight[offset * units : offset * units + units], + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + return_success=True, + ) + if success: + loaded_params.add(name_mapped) + is_found = True + if found_num == den: + break + if is_found: + continue + if is_expert_weight: + # We've checked that this is an expert weight + # However it's not mapped locally to this rank + # So we simply skip it + continue + # Remapping the name of FP8 kv-scale. + name = maybe_remap_kv_scale_name(name, params_dict) + if name is None: + continue + + if is_pp_missing_parameter(name, self): + continue + + if "mlp.gate.wg." in name: + name = name.replace("wg.", "") + if name == "ln_f.weight": + name = "norm.weight" + if name == "wte.weight": + name = "embed_tokens.weight" + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + custom_pos_emb: tuple[torch.FloatTensor] | None = None, + mode: str = "gen_text", + first_step: bool | None = None, + query_lens: list[int] | None = None, + seq_lens: list[int] | None = None, + num_image_tokens: int | None = None, + gen_timestep_scatter_index: torch.Tensor | None = None, + ) -> tuple | BaseModelOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if inputs_embeds is None: + inputs_embeds = self.wte(input_ids) + + # embed positions + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for layer_idx, decoder_layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = decoder_layer( + positions=None, + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + custom_pos_emb=custom_pos_emb, + mode=mode, + first_step=first_step, + query_lens=query_lens, + seq_lens=seq_lens, + num_image_tokens=num_image_tokens, + gen_timestep_scatter_index=gen_timestep_scatter_index, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +class ClassifierFreeGuidance: + def __init__( + self, + use_original_formulation: bool = False, + start: float = 0.0, + stop: float = 1.0, + ): + super().__init__() + self.use_original_formulation = use_original_formulation + + def __call__( + self, + pred_cond: torch.Tensor, + pred_uncond: torch.Tensor | None, + guidance_scale: float, + step: int, + ) -> torch.Tensor: + shift = pred_cond - pred_uncond + pred = pred_cond if self.use_original_formulation else pred_uncond + pred = pred + guidance_scale * shift + + return pred + + +@dataclass +class HunyuanImage3Text2ImagePipelineOutput(BaseOutput): + samples: list[Any] | np.ndarray + + +class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): + r""" + Pipeline for condition-to-sample generation using Stable Diffusion. + + This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods + implemented for all pipelines (downloading, saving, running on a particular device, etc.). + + Args: + model ([`ModelMixin`]): + A model to denoise the diffused latents. + scheduler ([`SchedulerMixin`]): + A scheduler to be used in combination with `diffusion_model` to denoise the diffused latents. Can be one of + [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`]. + """ + + model_cpu_offload_seq = "" + _optional_components = [] + _exclude_from_cpu_offload = [] + _callback_tensor_inputs = ["latents"] + + def __init__( + self, + model, + scheduler: SchedulerMixin, + vae, + progress_bar_config: dict[str, Any] | None = None, + ): + super().__init__() + + if progress_bar_config is None: + progress_bar_config = {} + if not hasattr(self, "_progress_bar_config"): + self._progress_bar_config = {} + self._progress_bar_config.update(progress_bar_config) + self.register_modules( + model=model, + scheduler=scheduler, + vae=vae, + ) + + # should be a tuple or a list corresponding to the size of latents (batch_size, channel, *size) + # if None, will be treated as a tuple of 1 + self.latent_scale_factor = self.model.config.vae_downsample_factor + self.image_processor = VaeImageProcessor(vae_scale_factor=self.latent_scale_factor) + + # Must start with APG_mode_ + self.cfg_operator = ClassifierFreeGuidance() + + @staticmethod + def denormalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor: + """ + Denormalize an image array to [0,1]. + """ + return (images / 2 + 0.5).clamp(0, 1) + + @staticmethod + def pt_to_numpy(images: torch.Tensor) -> np.ndarray: + """ + Convert a PyTorch tensor to a NumPy image. + """ + images = images.cpu().permute(0, 2, 3, 1).float().numpy() + return images + + @staticmethod + def numpy_to_pil(images: np.ndarray): + """ + Convert a numpy image or a batch of images to a PIL image. + """ + if images.ndim == 3: + images = images[None, ...] + images = (images * 255).round().astype("uint8") + if images.shape[-1] == 1: + # special case for grayscale (single channel) images + pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] + else: + pil_images = [Image.fromarray(image) for image in images] + + return pil_images + + def prepare_extra_func_kwargs(self, func, kwargs): + # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature + # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. + # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 + # and should be between [0, 1] + extra_kwargs = {} + + for k, v in kwargs.items(): + accepts = k in set(inspect.signature(func).parameters.keys()) + if accepts: + extra_kwargs[k] = v + return extra_kwargs + + def prepare_latents(self, batch_size, latent_channel, image_size, dtype, device, generator, latents=None): + if self.latent_scale_factor is None: + latent_scale_factor = (1,) * len(image_size) + elif isinstance(self.latent_scale_factor, int): + latent_scale_factor = (self.latent_scale_factor,) * len(image_size) + elif isinstance(self.latent_scale_factor, tuple) or isinstance(self.latent_scale_factor, list): + assert len(self.latent_scale_factor) == len(image_size), ( + "len(latent_scale_factor) should be the same as len(image_size)" + ) + latent_scale_factor = self.latent_scale_factor + else: + raise ValueError( + f"latent_scale_factor should be either None, int, tuple of int, or list of int, " + f"but got {self.latent_scale_factor}" + ) + + latents_shape = ( + batch_size, + latent_channel, + *[int(s) // f for s, f in zip(image_size, latent_scale_factor)], + ) + 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(latents_shape, generator=generator, device=device, dtype=dtype) + else: + latents = latents.to(device) + + # Check existence to make it compatible with FlowMatchEulerDiscreteScheduler + if hasattr(self.scheduler, "init_noise_sigma"): + # scale the initial noise by the standard deviation required by the scheduler + latents = latents * self.scheduler.init_noise_sigma + + return latents + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + def set_scheduler(self, new_scheduler): + self.register_modules(scheduler=new_scheduler) + + @torch.no_grad() + def __call__( + self, + batch_size: int, + image_size: list[int], + num_inference_steps: int = 50, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + guidance_scale: float = 5.0, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + return_dict: bool = True, + guidance_rescale: float = 0.0, + callback_on_step_end: Callable[[int, int, dict], None] + | PipelineCallback + | MultiPipelineCallbacks + | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + model_kwargs: dict[str, Any] | None = None, + **kwargs, + ): + r""" + The call function to the pipeline for generation. + + Args: + prompt (`str` or `List[str]`): + The text to guide image generation. + image_size (`Tuple[int]` or `List[int]`): + The size (height, width) of the generated image. + num_inference_steps (`int`, *optional*, defaults to 50): + The number of denoising steps. More denoising steps usually lead to a higher quality image at the + expense of slower inference. + timesteps (`List[int]`, *optional*): + Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument + in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is + passed will be used. Must be in descending order. + sigmas (`List[float]`, *optional*): + Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in + their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed + will be used. + guidance_scale (`float`, *optional*, defaults to 7.5): + A higher guidance scale value encourages the model to generate samples closely linked to the + `condition` at the expense of lower sample quality. Guidance scale is enabled when `guidance_scale > 1`. + generator (`torch.Generator` or `List[torch.Generator]`, *optional*): + A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make + generation deterministic. + latents (`torch.Tensor`, *optional*): + Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for sample + generation. Can be used to tweak the same generation with different conditions. If not provided, + a latents tensor is generated by sampling using the supplied random `generator`. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generated sample. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~DiffusionPipelineOutput`] instead of a + plain tuple. + guidance_rescale (`float`, *optional*, defaults to 0.0): + Guidance rescale factor from [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://arxiv.org/pdf/2305.08891.pdf). Guidance rescale factor should fix overexposure when + using zero terminal SNR. + callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*): + A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of + each denoising step during the inference. with the following arguments: `callback_on_step_end(self: + DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a + list of all tensors as specified by `callback_on_step_end_tensor_inputs`. + callback_on_step_end_tensor_inputs (`List`, *optional*): + The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list + will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the + `._callback_tensor_inputs` attribute of your pipeline class. + + Examples: + + Returns: + [`~DiffusionPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`~DiffusionPipelineOutput`] is returned, + otherwise a `tuple` is returned where the first element is a list with the generated samples. + """ + callback_steps = kwargs.pop("callback_steps", None) + pbar_steps = kwargs.pop("pbar_steps", None) + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + self._guidance_scale = guidance_scale + self._guidance_rescale = guidance_rescale + + cfg_factor = 1 + self.do_classifier_free_guidance + + # Define call parameters + device = self._execution_device + + # Prepare timesteps + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas, + ) + + # Prepare latent variables + latents = self.prepare_latents( + batch_size=batch_size, + latent_channel=self.model.config.vae["latent_channels"], + image_size=image_size, + dtype=torch.bfloat16, + device=device, + generator=generator, + latents=latents, + ) + + # Prepare extra step kwargs. + _scheduler_step_extra_kwargs = self.prepare_extra_func_kwargs(self.scheduler.step, {"generator": generator}) + + # Prepare model kwargs + input_ids = model_kwargs.pop("input_ids") + attention_mask = self.model._prepare_attention_mask_for_generation( # noqa + input_ids, + self.model.generation_config, + model_kwargs=model_kwargs, + ) + b, _, q_len1, seq_len = attention_mask.shape + query_lens = [q_len1] * b + seq_lens = [seq_len] * b + model_kwargs["query_lens"] = query_lens + model_kwargs["seq_lens"] = seq_lens + model_kwargs["attention_mask"] = attention_mask.to(latents.device) + + # Sampling loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * cfg_factor) + # latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + t_expand = t.repeat(latent_model_input.shape[0]) + + model_inputs = self.model.prepare_inputs_for_generation( + input_ids, + images=latent_model_input, + timestep=t_expand, + **model_kwargs, + ) + + with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True): + model_output = self.model.forward_call(**model_inputs, first_step=(i == 0)) + pred = model_output["diffusion_prediction"] + pred = pred.to(dtype=torch.float32) + + # perform guidance + if self.do_classifier_free_guidance: + pred_cond, pred_uncond = pred.chunk(2) + pred = self.cfg_operator(pred_cond, pred_uncond, self.guidance_scale, step=i) + + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0] + + if i != len(timesteps) - 1: + model_kwargs = self.model._update_model_kwargs_for_generation( # noqa + model_output, + model_kwargs, + ) + if input_ids.shape[1] != model_kwargs["position_ids"].shape[1]: + input_ids = torch.gather(input_ids, 1, index=model_kwargs["position_ids"]) + attention_mask = model_kwargs.get("attention_mask") + b, _, q_len1, seq_len = attention_mask.shape + query_lens = [q_len1] * b + seq_lens = [seq_len] * b + model_kwargs["query_lens"] = query_lens + model_kwargs["seq_lens"] = seq_lens + + 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) + + # 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() + + if hasattr(self.vae.config, "scaling_factor") and self.vae.config.scaling_factor: + latents = latents / self.vae.config.scaling_factor + if hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor: + latents = latents + self.vae.config.shift_factor + + if hasattr(self.vae, "ffactor_temporal"): + latents = latents.unsqueeze(2) + + with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True): + image = self.vae.decode(latents, return_dict=False, generator=generator)[0] + + if hasattr(self.vae, "ffactor_temporal"): + assert image.shape[2] == 1, "image should have shape [B, C, T, H, W] and T should be 1" + image = image.squeeze(2) + + do_denormalize = [True] * image.shape[0] + image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + + if not return_dict: + return (image,) + + return HunyuanImage3Text2ImagePipelineOutput(samples=image) + + +def timestep_embedding(t, dim, max_period=10000): + """ + Create sinusoidal timestep embeddings. + + Args: + t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. + dim (int): the dimension of the output. + max_period (int): controls the minimum frequency of the embeddings. + + Returns: + embedding (torch.Tensor): An (N, D) Tensor of positional embeddings. + + .. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + """ + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( + device=t.device + ) + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + act_layer=nn.GELU, + frequency_embedding_size=256, + max_period=10000, + out_size=None, + dtype=None, + device=None, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.frequency_embedding_size = frequency_embedding_size + self.max_period = max_period + if out_size is None: + out_size = hidden_size + + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs), + act_layer(), + nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs), + ) + nn.init.normal_(self.mlp[0].weight, std=0.02) + nn.init.normal_(self.mlp[2].weight, std=0.02) + + def forward(self, t): + t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).type(self.mlp[0].weight.dtype) + t_emb = self.mlp(t_freq) + return t_emb + + +class ResBlock(nn.Module): + r""" + A residual block that can optionally change the number of channels. + + Args: + in_channels (`int`): + The number of input channels. + emb_channels (`int`): + The number of timestep embedding channels. + dropout (`float`): + The rate of dropout. + out_channels (`int`, *optional*): + If specified, the number of output channels. + use_conv (`bool`, *optional*): + If True and out_channels is specified, use a spatial convolution instead of a + smaller 1x1 convolution to change the channels in the skip connection. + dims (`int`, *optional*): + Determines if the signal is 1D, 2D, or 3D. + up (`bool`, *optional*): + If True, use this block for upsampling. + down (`bool`, *optional*): + If True, use this block for downsampling. + + """ + + def __init__( + self, + in_channels, + emb_channels, + out_channels=None, + dropout=0.0, + use_conv=False, + dims=2, + up=False, + down=False, + device=None, + dtype=None, + ) -> None: + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + self.in_channels = in_channels + self.dropout = dropout + self.out_channels = out_channels or self.in_channels + self.use_conv = use_conv + + self.in_layers = nn.Sequential( + normalization(self.in_channels, **factory_kwargs), + nn.SiLU(), + conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs), # noqa: N802 + ) + + self.updown = up or down + self.h_upd = self.x_upd = nn.Identity() + + self.emb_layers = nn.Sequential(nn.SiLU(), linear(emb_channels, 2 * self.out_channels, **factory_kwargs)) + + self.out_layers = nn.Sequential( + normalization(self.out_channels, **factory_kwargs), + nn.SiLU(), + nn.Dropout(p=dropout), + zero_module(conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, **factory_kwargs)), # noqa: N802 + ) + + if self.out_channels == self.in_channels: + self.skip_connection = nn.Identity() + elif use_conv: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 3, padding=1, **factory_kwargs) # noqa: N802 + else: + self.skip_connection = conv_nd(dims, self.in_channels, self.out_channels, 1, **factory_kwargs) # noqa: N802 + + def forward(self, x, emb) -> torch.Tensor: + if self.updown: + in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] + h = in_rest(x) + h = self.h_upd(h) + x = self.x_upd(x) + h = in_conv(h) + else: + h = self.in_layers(x) + + emb_out = self.emb_layers(emb) + while len(emb_out.shape) < len(h.shape): + emb_out = emb_out[..., None] + + # Adaptive Group Normalization + out_norm, out_rest = self.out_layers[0], self.out_layers[1:] + scale, shift = torch.chunk(emb_out, 2, dim=1) + h = out_norm(h) * (1.0 + scale) + shift + h = out_rest(h) + + return self.skip_connection(x) + h + + +class UNetDown(nn.Module): + def __init__( + self, patch_size, in_channels, emb_channels, hidden_channels, out_channels, dropout=0.0, device=None, dtype=None + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + + self.patch_size = patch_size + assert self.patch_size in [1, 2, 4, 8] + + self.model = nn.ModuleList( + [ + conv_nd( + 2, in_channels=in_channels, out_channels=hidden_channels, kernel_size=3, padding=1, **factory_kwargs + ) # noqa: N802 + ] + ) + + if self.patch_size == 1: + self.model.append( + ResBlock( + in_channels=hidden_channels, + emb_channels=emb_channels, + out_channels=out_channels, + dropout=dropout, + **factory_kwargs, + ) + ) + else: + for i in range(self.patch_size // 2): + self.model.append( + ResBlock( + in_channels=hidden_channels, + emb_channels=emb_channels, + out_channels=hidden_channels if (i + 1) * 2 != self.patch_size else out_channels, + dropout=dropout, + down=True, + **factory_kwargs, + ) + ) + + def forward(self, x, t): + assert x.shape[2] % self.patch_size == 0 and x.shape[3] % self.patch_size == 0 + for module in self.model: + if isinstance(module, ResBlock): + x = module(x, t) + else: + x = module(x) + _, _, token_h, token_w = x.shape + x = rearrange(x, "b c h w -> b (h w) c") + return x, token_h, token_w + + +class UNetUp(nn.Module): + def __init__( + self, + patch_size, + in_channels, + emb_channels, + hidden_channels, + out_channels, + dropout=0.0, + device=None, + dtype=None, + out_norm=False, + ): + factory_kwargs = {"dtype": dtype, "device": device} + super().__init__() + + self.patch_size = patch_size + assert self.patch_size in [1, 2, 4, 8] + + self.model = nn.ModuleList() + + if self.patch_size == 1: + self.model.append( + ResBlock( + in_channels=in_channels, + emb_channels=emb_channels, + out_channels=hidden_channels, + dropout=dropout, + **factory_kwargs, + ) + ) + else: + for i in range(self.patch_size // 2): + self.model.append( + ResBlock( + in_channels=in_channels if i == 0 else hidden_channels, + emb_channels=emb_channels, + out_channels=hidden_channels, + dropout=dropout, + up=True, + **factory_kwargs, + ) + ) + + if out_norm: + self.model.append( + nn.Sequential( + normalization(hidden_channels, **factory_kwargs), + nn.SiLU(), + conv_nd( + 2, + in_channels=hidden_channels, + out_channels=out_channels, + kernel_size=3, + padding=1, + **factory_kwargs, + ), # noqa: N802 + ) + ) + else: + self.model.append( + conv_nd( + 2, + in_channels=hidden_channels, + out_channels=out_channels, + kernel_size=3, + padding=1, + **factory_kwargs, + ) # noqa: N802 + ) + + # batch_size, seq_len, model_dim + def forward(self, x, t, token_h, token_w): + x = rearrange(x, "b (h w) c -> b c h w", h=token_h, w=token_w) + for module in self.model: + if isinstance(module, ResBlock): + x = module(x, t) + else: + x = module(x) + return x diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py new file mode 100644 index 000000000..9a8f26d38 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py @@ -0,0 +1,1017 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import logging +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from transformers.generation.configuration_utils import GenerationConfig +from transformers.generation.utils import ALL_CACHE_NAMES, GenerationMixin +from transformers.models.siglip2 import Siglip2VisionConfig, Siglip2VisionModel +from transformers.utils.generic import ModelOutput +from vllm.config.vllm import get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_rank +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.transformers_utils.config import get_config + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.request import OmniDiffusionRequest + +from .autoencoder import AutoencoderKLConv3D +from .hunyuan_image_3_tokenizer import TokenizerWrapper +from .hunyuan_image_3_transformer import ( + CausalMMOutputWithPast, + HunyuanImage3ImageProcessor, + HunyuanImage3Model, + HunyuanImage3PreTrainedModel, + HunyuanImage3Text2ImagePipeline, + ImageInfo, + JointImageInfo, + LightProjector, + TimestepEmbedder, + UNetDown, + UNetUp, + build_batch_2d_rope, + get_full_state_dict, + real_batched_index_select, +) + +logger = logging.getLogger(__name__) + +BatchRaggedImages = torch.Tensor | list[torch.Tensor | list[torch.Tensor]] +BatchRaggedTensor = torch.Tensor | list[torch.Tensor] + + +def default(val, d): + return val if val is not None else d + + +def to_device(data, device): + if device is None: + return data + if isinstance(data, torch.Tensor): + return data.to(device) + elif isinstance(data, list): + return [to_device(x, device) for x in data] + else: + return data + + +class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin): + def __init__(self, od_config: OmniDiffusionConfig) -> None: + self.hf_config = get_config(od_config.model, trust_remote_code=True) + super().__init__(self.hf_config) + # update diffusion config + self.generation_config = GenerationConfig.from_pretrained(od_config.model) + self.od_config = od_config + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder=None, + revision=od_config.revision, + prefix="", + fall_back_to_pt=True, + ) + ] + self.model = HunyuanImage3Model(self.hf_config) + self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae) + self._pipeline = None + self._tkwrapper = TokenizerWrapper(od_config.model) + self.image_processor = HunyuanImage3ImageProcessor(self.hf_config) + self.hf_config.vit.pop("use_return_dict", None) + vision_config = Siglip2VisionConfig(**self.hf_config.vit) + self.vision_model = Siglip2VisionModel(vision_config).vision_model + # self.vision_model = vision_model.vision_model + self.vision_aligner = LightProjector(self.hf_config.vit_aligner) + self.timestep_emb = TimestepEmbedder(hidden_size=self.hf_config.hidden_size) + if self.hf_config.img_proj_type != "unet": + raise ValueError(f"Unknown img_proj_type: {self.hf_config.img_proj_type}") + + self.patch_embed = UNetDown( + patch_size=self.hf_config.patch_size, + emb_channels=self.hf_config.hidden_size, + in_channels=self.hf_config.vae["latent_channels"], + hidden_channels=self.hf_config.patch_embed_hidden_dim, + out_channels=self.hf_config.hidden_size, + ) + self.time_embed = TimestepEmbedder(hidden_size=self.hf_config.hidden_size) + self.final_layer = UNetUp( + patch_size=self.hf_config.patch_size, + emb_channels=self.hf_config.hidden_size, + in_channels=self.hf_config.hidden_size, + hidden_channels=self.hf_config.patch_embed_hidden_dim, + out_channels=self.hf_config.vae["latent_channels"], + out_norm=True, + ) + self.time_embed_2 = TimestepEmbedder(hidden_size=self.hf_config.hidden_size) + self.lm_head = nn.Linear(self.hf_config.hidden_size, self.hf_config.vocab_size, bias=False) + self.vllm_config = get_current_vllm_config() + self.post_init() + + def pre_load(self): + tp_rank = get_tensor_model_parallel_rank() + state_dict = get_full_state_dict(self.od_config.model) + non_layer_prefixes = [ + "vae", + "vision_model", + "vision_aligner", + "lm_head", + "patch_embed", + "timestep_emb", + "model.wte", + "model.ln_f", + "time_embed", + "time_embed_2", + "final_layer.model", + ] + filtered_sd = {} + for k, v in state_dict.items(): + if any(k.startswith(prefix) for prefix in non_layer_prefixes): + filtered_sd[k] = v + + missing, unexpected = self.load_state_dict(filtered_sd, strict=False) + + for prefix in non_layer_prefixes: + if hasattr(self, prefix.split(".")[0]): + module = dict(self.named_modules()).get(prefix) + if module: + module.to(f"cuda:{tp_rank}") + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + self.pre_load() + skip_prefixes = ["lm_head."] if self.hf_config.tie_word_embeddings else [] + # List of unexpected keywords in weight names + unexpected_keywords = [ + "vae", + "vision_aligner", + "vision_model", + "final_layer", + "patch_embed", + "timestep_emb", + "time_embed", + "time_embed_2", + "guidance_emb", + "timestep_r_emb", + ] + skip_prefixes.extend(unexpected_keywords) + loader = AutoWeightsLoader( + self, + skip_prefixes=skip_prefixes, + ) + return loader.load_weights(weights) + + def prepare_seed(self, seed=None, batch_size=1): + # random seed + if seed is not None: + return [seed + i for i in range(batch_size)] + else: + import random + + return [random.randint(0, 2**32 - 1) for _ in range(batch_size)] + + @property + def pipeline(self): + if self._pipeline is None: + # shift hard code + self.scheduler = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=1000, + shift=self.generation_config.flow_shift, + use_dynamic_shifting=False, + base_shift=0.5, + max_shift=1.15, + time_shift_type="exponential", + stochastic_sampling=False, + ) + self._pipeline = HunyuanImage3Text2ImagePipeline(model=self, scheduler=self.scheduler, vae=self.vae) + return self._pipeline + + @staticmethod + def get_pos_emb(custom_pos_emb, position_ids): + cos, sin = custom_pos_emb + cos = real_batched_index_select(cos, dim=1, idx=position_ids) + sin = real_batched_index_select(sin, dim=1, idx=position_ids) + return cos, sin + + def instantiate_vae_image_tokens( + self, + x: torch.Tensor, + images: BatchRaggedImages, + ts: BatchRaggedTensor, + image_mask: torch.Tensor, + ): + r""" + Instantiate the VAE image embeddings into the input embedding sequence. + Args: + x (`torch.Tensor`): + Input sequence tensor with shape `(batch_size, seq_len, n_embd)`. + images (`BatchRaggedImages`): + Batch of images to embed. Can be: + - A 4-D tensor (batch, channels, height, width) + - A list of 4-D tensors (variable number of images per batch) + - A list of lists of 3-D tensors (ragged batch structure) + ts (`BatchRaggedTensor`, *optional*): + Timestep tensor(s) for conditioning. Can be: + - A 1-D tensor (single timestep per batch) + - A list of 1-D tensors (variable timesteps per batch) + image_mask (`torch.Tensor`, *optional*): + Boolean mask tensor with shape `(batch_size, seq_len)` indicating which positions + should be replaced with image embeddings. + """ + batch_size, seq_len, n_embd = x.shape + + if isinstance(images, list): + index = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1) + t_emb = [] + for i, (image_i, t_i) in enumerate(zip(images, ts)): + if isinstance(image_i, torch.Tensor): + # time_embed needs a 1-D tensor as input + t_i_emb = self.time_embed(t_i) + # n_{i} x one_image_seq_len x n_embd + image_i_seq, _, _ = self.patch_embed(image_i, t_i_emb) + # 1 x (n_{i} * one_image_seq_len) + image_i_scatter_index = index[i : i + 1].masked_select(image_mask[i : i + 1].bool()).reshape(1, -1) + x[i : i + 1].scatter_( + dim=1, + index=image_i_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + # 1 x (n_{i} * one_image_seq_len) x n_embd + src=image_i_seq.reshape(1, -1, n_embd), # 1 x (n_{i} * one_image_seq_len) x n_embd + ) + t_emb.append(t_i_emb) + elif isinstance(image_i, list): + # time_embed needs a 1-D tensor as input + t_i_emb = self.time_embed(t_i) # n_{i} x d + image_i_seq_list = [] + for j in range(len(image_i)): + image_ij = image_i[j] + if image_ij.dim() == 4: + assert image_i[j].shape[0] == 1, "image_i[j] should have a batch dimension of 1" + elif image_ij.dim() == 3: + image_ij = image_ij.unsqueeze(0) + else: + raise ValueError(f"image_i[j] should have 3 or 4 dimensions, got {image_ij.dim()}") + # 1 x one_image_seq_len_{j} x n_embd + image_i_seq_j, _, _ = self.patch_embed(image_ij, t_i_emb[j : j + 1]) + image_i_seq_list.append(image_i_seq_j) + # 1 x sum_{j}(one_image_seq_len_{j}) x n_embd + image_i_seq = torch.cat(image_i_seq_list, dim=1) + # 1 x sum_{j}(one_image_seq_len_{j}) + image_i_scatter_index = index[i : i + 1].masked_select(image_mask[i : i + 1].bool()).reshape(1, -1) + x[i : i + 1].scatter_( + dim=1, + index=image_i_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + # 1 x sum_{j}(one_image_seq_len_{j}) x n_embd + src=image_i_seq.reshape(1, -1, n_embd), # 1 x sum_{j}(one_image_seq_len_{j}) x n_embd + ) + t_emb.append(t_i_emb) + else: + raise TypeError(f"image_i should be a torch.Tensor or a list, got {type(image_i)}") + token_h, token_w = None, None + else: + # images is a 4-D tensor + batch_size, seq_len, n_embd = x.shape + index = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1) + t_emb = self.time_embed(ts) + image_seq, token_h, token_w = self.patch_embed(images, t_emb) + image_scatter_index = index.masked_select(image_mask.bool()).reshape(batch_size, -1) + x.scatter_( + dim=1, + index=image_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + src=image_seq, + ) + + return x, token_h, token_w + + def instantiate_timestep_tokens( + self, + x: torch.Tensor, + t: BatchRaggedTensor, + timestep_scatter_index: BatchRaggedTensor, + ): + batch_size, seq_len, n_embd = x.shape + # batch_size x n x n_embd + timestep_scatter_src = self.timestep_emb(t.reshape(-1)).reshape(batch_size, -1, n_embd) + x.scatter_( + dim=1, + index=timestep_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + src=timestep_scatter_src, + ) + + return x + + def instantiate_vit_image_tokens( + self, + x: torch.Tensor, + cond_vit_images: torch.Tensor | list[torch.Tensor], + cond_vit_image_mask: torch.Tensor, + vit_kwargs: dict[str, Any], + ): + # 1. Forward the vit encoder and vit aligner to get the vit image embeddings and align them to the + # transformer hidden size + cond_vit_image_embeds = [] + for batch_idx, image in enumerate(cond_vit_images): + cur_kwargs = {k: v[batch_idx] for k, v in vit_kwargs.items()} + image_embed = self.vision_model(image, **cur_kwargs).last_hidden_state + image_embed = self.vision_aligner(image_embed) + n, seq_len, dim = image_embed.shape + image_embed = image_embed.reshape(n * seq_len, dim) + cond_vit_image_embeds.append(image_embed) + + # 2. Instantiate the vit image embeddings into the input sequence + batch_size, seq_len, n_embd = x.shape + index = torch.arange(seq_len, device=x.device).unsqueeze(0).repeat(batch_size, 1) + + for i, (image_embed, mask) in enumerate(zip(cond_vit_image_embeds, cond_vit_image_mask)): + image_scatter_index = index[i : i + 1].masked_select(mask.bool()).reshape(1, -1) + x[i : i + 1].scatter_( + dim=1, + index=image_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + src=image_embed.reshape(1, -1, n_embd), + ) + + return x + + def ragged_final_layer(self, x, image_mask, timestep, token_h, token_w, first_step): + bsz, seq_len, n_embd = x.shape + if first_step: + image_output = x.masked_select(image_mask.unsqueeze(-1).bool()).reshape(bsz, -1, n_embd) + else: + image_output = x[:, 1:, :] + timestep_emb = self.time_embed_2(timestep) + pred = self.final_layer(image_output, timestep_emb, token_h, token_w) + return pred + + @staticmethod + def build_batch_rope_image_info(output, sections): + rope_image_info = [] + for image_slices, sections_i in zip(output.all_image_slices, sections): + image_shapes = [] + for section in sections_i: + if "image" in section["type"]: + if isinstance(section["token_height"], list): + assert len(section["token_height"]) == len(section["token_width"]), ( + f"token_height and token_width should have the same length, " + f"but got {len(section['token_height'])} and {len(section['token_width'])}" + ) + image_shapes.extend(list(zip(section["token_height"], section["token_width"]))) + else: + image_shapes.append((section["token_height"], section["token_width"])) + assert len(image_slices) == len(image_shapes), ( + f"Size miss matching: Image slices({len(image_slices)}) != image shapes({len(image_shapes)})" + ) + rope_image_info.append(list(zip(image_slices, image_shapes))) + return rope_image_info + + def vae_encode(self, image, cfg_factor=1): + config = self.vae.config + + with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=True): + vae_encode_result = self.vae.encode(image) + if isinstance(vae_encode_result, torch.Tensor): + latents = vae_encode_result + else: + latents = vae_encode_result.latent_dist.sample() + if hasattr(config, "shift_factor") and config.shift_factor: + latents.sub_(config.shift_factor) + if hasattr(config, "scaling_factor") and config.scaling_factor: + latents.mul_(config.scaling_factor) + + if hasattr(self.vae, "ffactor_temporal"): + assert latents.shape[2] == 1, "latents should have shape [B, C, T, H, W] and T should be 1" + latents = latents.squeeze(2) + + # Here we always use t=0 to declare it is a clean conditional image + t = torch.zeros((latents.shape[0],)) + + if cfg_factor > 1: + t = t.repeat(cfg_factor) + latents = latents.repeat(cfg_factor, 1, 1, 1) + + return t, latents + + def _encode_cond_image( + self, + batch_cond_image_info_list: list[list[JointImageInfo]], + cfg_factor: int = 1, + ): + # VAE encode one by one, as we assume cond images have different sizes + batch_cond_vae_images, batch_cond_t, batch_cond_vit_images = [], [], [] + for cond_image_info_list in batch_cond_image_info_list: + cond_vae_image_list, cond_t_list, cond_vit_image_list = [], [], [] + for image_info in cond_image_info_list: + cond_t_, cond_vae_image_ = self.vae_encode( + image_info.vae_image_info.image_tensor.to(self.device), + ) + cond_vit_image_list.append(image_info.vision_image_info.image_tensor) + cond_vae_image_list.append(cond_vae_image_.squeeze(0)) + cond_t_list.append(cond_t_) + batch_cond_vae_images.append(cond_vae_image_list) + batch_cond_t.append(cond_t_list) + batch_cond_vit_images.append(torch.cat(cond_vit_image_list, dim=0)) + + # If only one cond image for each sample and all have the same size, we can batch them together + # In this case, cond_vae_images is a 4-D tensor. + if all([len(items) == 1 for items in batch_cond_vae_images]) and all( + items[0].shape == batch_cond_vae_images[0][0].shape for items in batch_cond_vae_images + ): + cond_vae_images = torch.stack([items[0] for items in batch_cond_vae_images], dim=0) + cond_t = torch.cat([items[0] for items in batch_cond_t], dim=0) + if cfg_factor > 1: + cond_t = cond_t.repeat(cfg_factor) + cond_vae_images = cond_vae_images.repeat(cfg_factor, 1, 1, 1) + else: + # In this case, cond_vae_images is a list of 4-D tensors or a list of lists of 3-D tensors. + cond_t = [torch.cat(item, dim=0) for item in batch_cond_t] + cond_vae_images = [] + for items in batch_cond_vae_images: + if all(items[0].shape == item.shape for item in items): + cond_vae_images.append(torch.stack(items, dim=0)) + else: + cond_vae_images.append(items) + if cfg_factor > 1: + cond_t = cond_t * cfg_factor + cond_vae_images = cond_vae_images * cfg_factor + + if cfg_factor > 1: + batch_cond_vit_images = batch_cond_vit_images * cfg_factor + + return cond_vae_images, cond_t, batch_cond_vit_images + + @staticmethod + def check_inputs(prompt=None, message_list=None): + if prompt is None and message_list is None: + raise ValueError("Either `prompt` or `message_list` should be provided.") + if prompt is not None and message_list is not None: + raise ValueError("Only one of `prompt` or `message_list` should be provided.") + if prompt is not None: + assert isinstance(prompt, str) or isinstance(prompt, list), ( + f"`prompt` should be a string or a list of strings, but got {type(prompt)}." + ) + if isinstance(prompt, list): + assert len(prompt) > 0 and all(isinstance(p, str) for p in prompt), ( + "`prompt` should be a non-empty list of strings." + ) + if message_list is not None: + if not isinstance(message_list, list): + raise ValueError(f"`message_list` should be a list of messages, but got {type(message_list)}.") + assert len(message_list) > 0, "`message_list` should be a non-empty list." + for message in message_list: + assert isinstance(message, list) or isinstance(message, dict), ( + f"Each message should be a list of dicts or a dict, but got {type(message)}." + ) + + def prepare_model_inputs( + self, + prompt=None, + mode="gen_image", + system_prompt=None, + cot_text=None, + num_inference_steps=50, + guidance_scale=5.0, + image_size="auto", + message_list=None, + device=None, + max_new_tokens=None, + **kwargs, + ): + # 1. Sanity check + self.check_inputs(prompt, message_list) + device = default(device, self.device) + + # 2. Format inputs + batch_message_list = message_list + batch_prompt = prompt + batch_cot_text = cot_text + batch_system_prompt = system_prompt + batch_gen_image_info = None + # TODO: construct with user input images + batch_cond_image_info = None + + # -- 2.1 message_list + if batch_message_list is not None: + if isinstance(batch_message_list[0], dict): + batch_message_list = [batch_message_list] + batch_size = len(batch_message_list) + + batch_gen_image_info = [ + [message["content"] for message in message_list_ if message["type"] == "gen_image"] + for message_list_ in batch_message_list + ] + # At most one gen_image is allowed for each message_list + batch_gen_image_info = [info[-1] if len(info) > 0 else None for info in batch_gen_image_info] + # Multiple cond images are allowed. + batch_cond_image_info = [ + [message["content"] for message in message_list_ if message["type"] == "joint_image"] + for message_list_ in batch_message_list + ] + + # -- 2.2 Prompt, cot text, system prompt + else: + if isinstance(batch_prompt, str): + batch_prompt = [batch_prompt] + batch_size = len(batch_prompt) + + if batch_cot_text is not None: + if isinstance(batch_cot_text, str): + batch_cot_text = [batch_cot_text] + else: + assert isinstance(batch_cot_text, list) and len(batch_cot_text) == batch_size, ( + "`cot_text` should be a string or a list of strings with the same length as `prompt`." + ) + + if batch_system_prompt is not None: + if isinstance(batch_system_prompt, str): + batch_system_prompt = [batch_system_prompt] + else: + assert isinstance(batch_system_prompt, list) and len(batch_system_prompt) == batch_size, ( + "`system_prompts` should be a string or a list of strings with the same length as `prompt`." + ) + + if mode == "gen_image": + batch_gen_image_info = [self.image_processor.build_image_info(image_size) for _ in range(batch_size)] + + # -- 2.3 seed + generator = kwargs.get("generator", None) + if generator is None: + seeds = self.prepare_seed(seed=kwargs.get("seed"), batch_size=batch_size) + generator = [torch.Generator(self.device).manual_seed(seed) for seed in seeds] + + # 3. apply chat template + cfg_factor = {"gen_text": 1, "gen_image": 2} + bot_task = kwargs.pop("bot_task", "auto") + # If `drop_think` enabled, always drop parts in the context. + drop_think = kwargs.get("drop_think", self.generation_config.drop_think) + # Apply batched prompt or batched message_list to build input sequence with associated info. + out = self._tkwrapper.apply_chat_template( + batch_prompt=batch_prompt, + batch_message_list=batch_message_list, + mode=mode, + batch_gen_image_info=batch_gen_image_info, + batch_cond_image_info=batch_cond_image_info, + batch_system_prompt=batch_system_prompt, + batch_cot_text=batch_cot_text, + max_length=kwargs.get("max_length"), + bot_task=bot_task, + image_base_size=self.config.image_base_size, + sequence_template="pretrain", + cfg_factor=cfg_factor[mode], + drop_think=drop_think, + ) + output, sections = out["output"], out["sections"] + + # 4. Encode conditional images + if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0: + cond_vae_images, cond_timestep, cond_vit_images = self._encode_cond_image( + batch_cond_image_info, cfg_factor[mode] + ) + # Pack vit kwargs. Siglip2-so requires spatial_shapes and attention_mask for inference. + vit_kwargs = {"spatial_shapes": [], "attention_mask": []} + for cond_image_info in batch_cond_image_info: + vit_kwargs["spatial_shapes"].append( + torch.stack([item.vision_encoder_kwargs["spatial_shapes"] for item in cond_image_info]) + ) + vit_kwargs["attention_mask"].append( + torch.stack([item.vision_encoder_kwargs["pixel_attention_mask"] for item in cond_image_info]) + ) + if cfg_factor[mode] > 1: + vit_kwargs["spatial_shapes"] = vit_kwargs["spatial_shapes"] * cfg_factor[mode] + vit_kwargs["attention_mask"] = vit_kwargs["attention_mask"] * cfg_factor[mode] + else: + cond_vae_images, cond_timestep, cond_vit_images = None, None, None + vit_kwargs = None + + # 5. Build position embeddings + rope_image_info = self.build_batch_rope_image_info(output, sections) + if mode == "gen_text": + seq_len = self.generation_config.max_length + else: + seq_len = output.tokens.shape[1] + cos, sin = build_batch_2d_rope( + image_infos=rope_image_info, + seq_len=seq_len, + n_elem=self.config.attention_head_dim, + device=device, + base=self.config.rope_theta, + ) + + # 6. Build kv cache + if bot_task == "img_ratio": + max_new_tokens = 1 + + # 7. Build position ids + batch_input_pos = torch.arange(0, output.tokens.shape[1], dtype=torch.long, device=device)[None].expand( + batch_size * cfg_factor[mode], -1 + ) # use expand to share indices to save memory + + # 8. Build model input kwargs + tkw = self._tkwrapper + if image_size == "auto": + extra_auto_stops = [tkw.special_token_map[f""] for i in range(33)] + else: + extra_auto_stops = [tkw.boi_token_id] + stop_token_id = dict( + auto=[tkw.eos_token_id] + extra_auto_stops, + image=[tkw.eos_token_id], + recaption=[tkw.end_recaption_token_id, tkw.end_answer_token_id, tkw.eos_token_id], + think=[tkw.end_recaption_token_id, tkw.end_answer_token_id, tkw.eos_token_id], + img_ratio=extra_auto_stops, + ) + model_input_kwargs = dict( + input_ids=output.tokens.to(device), + position_ids=batch_input_pos, + past_key_values=None, + custom_pos_emb=(cos, sin), + mode=mode, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + image_mask=to_device(output.gen_image_mask, device), + gen_timestep_scatter_index=to_device(output.gen_timestep_scatter_index, device), + cond_vae_images=to_device(cond_vae_images, device), + cond_timestep=to_device(cond_timestep, device), + cond_vae_image_mask=to_device(output.cond_vae_image_mask, device), + cond_vit_images=to_device(cond_vit_images, device), + cond_vit_image_mask=to_device(output.cond_vit_image_mask, device), + vit_kwargs={k: to_device(v, self.device) for k, v in vit_kwargs.items()} + if vit_kwargs is not None + else None, + cond_timestep_scatter_index=to_device(output.cond_timestep_scatter_index, device), + # for inner usage + tokenizer_output=output, + batch_gen_image_info=batch_gen_image_info, + generator=generator, + # generation config + eos_token_id=stop_token_id[bot_task], + max_new_tokens=max_new_tokens, + ) + + return model_input_kwargs + + def _prepare_attention_mask_for_generation( + self, + inputs_tensor: torch.Tensor, + generation_config: GenerationConfig, + model_kwargs: dict[str, Any], + ) -> torch.Tensor: + # create `4d` bool attention mask (b, 1, seqlen, seqlen) using this implementation to bypass the 2d requirement + # in the `transformers.generation_utils.GenerationMixin.generate`. + # This implementation can handle sequences with text and image modalities, where text tokens use causal + # attention and image tokens use full attention. + bsz, seq_len = inputs_tensor.shape + tokenizer_output = model_kwargs["tokenizer_output"] + batch_image_slices = [ + tokenizer_output.joint_image_slices[i] + tokenizer_output.gen_image_slices[i] for i in range(bsz) + ] + attention_mask = torch.ones(seq_len, seq_len, dtype=torch.bool).tril(diagonal=0).repeat(bsz, 1, 1) + for i in range(bsz): + for j, image_slice in enumerate(batch_image_slices[i]): + attention_mask[i, image_slice, image_slice] = True + attention_mask = attention_mask.unsqueeze(1) + return attention_mask + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + tokenizer_output=None, + batch_gen_image_info=None, + generator=None, + **kwargs, + ): + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and past_key_values is None: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + if input_ids.shape[1] != kwargs["position_ids"].shape[1]: # in decode steps + input_ids = torch.gather(input_ids, dim=1, index=kwargs["position_ids"]) + model_inputs = {"input_ids": input_ids} + + model_inputs.update( + { + "attention_mask": attention_mask, + "position_ids": kwargs["position_ids"], + "past_key_values": past_key_values, + "use_cache": kwargs.get("use_cache"), + "custom_pos_emb": kwargs["custom_pos_emb"], + "mode": kwargs["mode"], + "images": kwargs.get("images"), + "image_mask": kwargs.get("image_mask"), + "timestep": kwargs.get("timestep"), + "gen_timestep_scatter_index": kwargs.get("gen_timestep_scatter_index"), + "cond_vae_images": kwargs.get("cond_vae_images"), + "cond_timestep": kwargs.get("cond_timestep"), + "cond_vae_image_mask": kwargs.get("cond_vae_image_mask"), + "cond_vit_images": kwargs.get("cond_vit_images"), + "cond_vit_image_mask": kwargs.get("cond_vit_image_mask"), + "vit_kwargs": kwargs.get("vit_kwargs"), + "cond_timestep_scatter_index": kwargs.get("cond_timestep_scatter_index"), + "query_lens": kwargs.get("query_lens"), + "seq_lens": kwargs.get("seq_lens"), + "num_image_tokens": kwargs.get("num_image_tokens"), + } + ) + return model_inputs + + def _update_model_kwargs_for_generation( + self, + outputs: ModelOutput, + model_kwargs: dict[str, Any], + is_encoder_decoder: bool = False, + num_new_tokens: int = 1, + ) -> dict[str, Any]: + mode = model_kwargs["mode"] + + updated_model_kwargs = { + "mode": mode, + "custom_pos_emb": model_kwargs["custom_pos_emb"], + } + + # update past_key_values keeping its naming used in model code + for possible_cache_name in ALL_CACHE_NAMES: + if possible_cache_name in outputs: + # TODO (joao): remove output/input mismatch when these old models (xlnet, reformer) are deprecated + if possible_cache_name in ("past_buckets_states", "mems"): + cache_name = "past_key_values" + else: + cache_name = possible_cache_name + updated_model_kwargs[cache_name] = getattr(outputs, possible_cache_name) + break + + if "tokenizer_output" in model_kwargs: + if mode == "gen_text": + # When enable batching, we use right padding, which requires a real_pos to index the valid + # end position of the sequence. If tokenizer_output in model_kwargs, it means we are in the + # prefill step of generation. + real_pos = to_device(model_kwargs["tokenizer_output"].real_pos, self.device) + updated_model_kwargs["position_ids"] = real_pos + else: + # position ids + image_mask = model_kwargs["image_mask"] + bsz, seq_len = image_mask.shape + index = torch.arange(seq_len, device=image_mask.device).unsqueeze(0).repeat(bsz, 1) + position_ids = index.masked_select(image_mask.bool()).reshape(bsz, -1) + timestep_position_ids = index[ + torch.arange(bsz), model_kwargs["gen_timestep_scatter_index"][:, -1] + ].unsqueeze(-1) + updated_model_kwargs["position_ids"] = torch.cat([timestep_position_ids, position_ids], dim=1) + + # attention mask + mask_list = [] + for attention_mask_i, position_ids_i in zip( + model_kwargs["attention_mask"], updated_model_kwargs["position_ids"] + ): + mask_list.append(torch.index_select(attention_mask_i, dim=1, index=position_ids_i.reshape(-1))) + attention_mask = torch.stack(mask_list, dim=0) + updated_model_kwargs["attention_mask"] = attention_mask + updated_model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"] + + else: + if mode == "gen_text": + # Now we are in the decode steps. + updated_model_kwargs["position_ids"] = model_kwargs["position_ids"] + 1 + else: + updated_model_kwargs["position_ids"] = model_kwargs["position_ids"] + updated_model_kwargs["attention_mask"] = model_kwargs["attention_mask"] + updated_model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"] + + return updated_model_kwargs + + def _generate( + self, + generator: list[torch.Generator] | None = None, + **kwargs, + ): + mode = kwargs.get("mode", "gen_text") + # verbose > 1 not support + if mode == "gen_text": + raise NotImplementedError("Not support gen text for hunyuan image") + + elif mode == "gen_image": + batch_gen_image_info: list[ImageInfo] = kwargs.get("batch_gen_image_info") + if batch_gen_image_info is None: + raise ValueError("`batch_gen_image_info` should be provided when `mode` is `gen_image`.") + + image_info: ImageInfo = batch_gen_image_info[0] + num_image_tokens = ( + image_info.image_token_length + + (1 if image_info.add_timestep_token else 0) + + (1 if image_info.add_guidance_token else 0) + ) + kwargs["num_image_tokens"] = num_image_tokens + # 50 and 5.0 hard code + from vllm.forward_context import set_forward_context + + with set_forward_context(None, self.vllm_config): + results = self.pipeline( + batch_size=len(batch_gen_image_info), + image_size=[batch_gen_image_info[0].image_height, batch_gen_image_info[0].image_width], + num_inference_steps=kwargs.get("num_inference_steps", 50), + guidance_scale=kwargs.get("guidance_scale", 5.0), + generator=generator, + model_kwargs=kwargs, + ) + samples = results[0] + return samples + + else: + raise ValueError(f"Unknown mode {mode}, only `gen_text` and `gen_image` are supported.") + + @staticmethod + def _check_inputs(cond, target, check_list): + if cond: + for name, item in check_list: + assert item is not None, f"`{name}` should be provided when `{target}`." + + def forward_call( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + custom_pos_emb: tuple[torch.FloatTensor] | None = None, + mode: str = "gen_text", + first_step: bool | None = None, + # for gen image + images: BatchRaggedImages | None = None, + image_mask: torch.Tensor | None = None, + timestep: BatchRaggedTensor | None = None, + gen_timestep_scatter_index: torch.Tensor | None = None, + # for cond image + cond_vae_images: BatchRaggedImages | None = None, + cond_timestep: BatchRaggedTensor | None = None, + cond_vae_image_mask: torch.Tensor | None = None, + cond_vit_images: BatchRaggedImages | None = None, + cond_vit_image_mask: torch.Tensor | None = None, + vit_kwargs: dict[str, Any] | None = None, + cond_timestep_scatter_index: torch.Tensor | None = None, + query_lens: list[int] | None = None, + seq_lens: list[int] | None = None, + num_image_tokens: int | None = None, + ) -> tuple | CausalMMOutputWithPast: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + # Sanity Check of Inputs + self._check_inputs( + mode == "gen_image", + "in `gen_image` mode", + [ + ("images", images), + ("timestep", timestep), + ("gen_timestep_scatter_index", gen_timestep_scatter_index), + ], + ) + self._check_inputs( + mode == "gen_image" and first_step, + "in `gen_image` mode at the first step", + [ + ("image_mask", image_mask), + ], + ) + self._check_inputs( + cond_vae_images is not None, + "`cond_vae_images` is provided", + [ + ("cond_timestep", cond_timestep), + ("cond_vae_image_mask", cond_vae_image_mask), + ("cond_timestep_scatter_index", cond_timestep_scatter_index), + ], + ) + self._check_inputs( + cond_vit_images is not None, + "`cond_vit_images` is provided", + [ + ("cond_vit_image_mask", cond_vit_image_mask), + ("vit_kwargs", vit_kwargs), + ], + ) + + custom_pos_emb = self.get_pos_emb(custom_pos_emb, position_ids) + + inputs_embeds = self.model.wte(input_ids) + + bsz, seq_len, n_embd = inputs_embeds.shape + + # Instantiate placeholder tokens: , for the gen image + if mode == "gen_text": + # For gen_text, make sure gen_timestep_scatter_index is None + gen_timestep_scatter_index = None + token_h, token_w = None, None + else: + if first_step: + inputs_embeds, token_h, token_w = self.instantiate_vae_image_tokens( + inputs_embeds, images, timestep, image_mask + ) + inputs_embeds = self.instantiate_timestep_tokens(inputs_embeds, timestep, gen_timestep_scatter_index) + else: + t_emb = self.time_embed(timestep) + image_emb, token_h, token_w = self.patch_embed(images, t_emb) + timestep_emb = self.timestep_emb(timestep).reshape(bsz, -1, n_embd) + inputs_embeds = torch.cat([timestep_emb, image_emb], dim=1) + + # Instantiate placeholder tokens: , for cond images + # Should only run once with kv-cache enabled. + if cond_vae_images is not None: + inputs_embeds, _, _ = self.instantiate_vae_image_tokens( + inputs_embeds, cond_vae_images, cond_timestep, cond_vae_image_mask + ) + inputs_embeds = self.instantiate_timestep_tokens(inputs_embeds, cond_timestep, cond_timestep_scatter_index) + if cond_vit_images is not None: + inputs_embeds = self.instantiate_vit_image_tokens( + inputs_embeds, cond_vit_images, cond_vit_image_mask, vit_kwargs + ) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + custom_pos_emb=custom_pos_emb, + mode=mode, + first_step=first_step, + query_lens=query_lens, + seq_lens=seq_lens, + num_image_tokens=num_image_tokens, + gen_timestep_scatter_index=gen_timestep_scatter_index, + ) + hidden_states = outputs[0] + + if mode == "gen_text": + hidden_states = self.model.ln_f(hidden_states) + logits = self.lm_head(hidden_states) + logits = logits.float() + diffusion_prediction = None + else: + logits = None + hidden_states = hidden_states.to(input_ids.device) + assert hidden_states.numel() == bsz * seq_len * n_embd, ( + f"Shape mismatch: {hidden_states.shape} cannot reshape to ({bsz}, {seq_len}, {n_embd})" + ) + hidden_states = hidden_states.reshape(bsz, seq_len, n_embd) + diffusion_prediction = self.ragged_final_layer( + hidden_states, image_mask, timestep, token_h, token_w, first_step + ) + + if not return_dict: + output = (logits,) + outputs[1:] + (diffusion_prediction,) + return output + + output = CausalMMOutputWithPast( + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + diffusion_prediction=diffusion_prediction, + ) + + return output + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] = "", + image_size="auto", + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 50, + guidance_scale: float = 5.0, + system_prompt: str | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + **kwargs, + ) -> DiffusionOutput: + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + generator = req.sampling_params.generator or generator + 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 + if req.sampling_params.guidance_scale_provided: + guidance_scale = req.sampling_params.guidance_scale + if guidance_scale <= 1.0: + logger.warning("HunyuanImage3.0 does not support guidance_scale <= 1.0, will set it to 1.0 + epsilon.") + guidance_scale = 1.0 + np.finfo(float).eps + image_size = (height, width) + model_inputs = self.prepare_model_inputs( + prompt=prompt, + cot_text=None, + system_prompt=system_prompt, + mode="gen_image", + generator=generator, + image_size=image_size, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + ) + outputs = self._generate(**model_inputs, **kwargs) + return DiffusionOutput(output=outputs[0]) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 024352479..3d236ac31 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -85,6 +85,11 @@ _DIFFUSION_MODELS = { "pipeline_sd3", "StableDiffusion3Pipeline", ), + "HunyuanImage3ForCausalMM": ( + "hunyuan_image_3", + "pipeline_hunyuan_image_3", + "HunyuanImage3Pipeline", + ), "Flux2KleinPipeline": ( "flux2_klein", "pipeline_flux2_klein", diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 38f438062..0e12460e8 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -87,18 +87,17 @@ class DiffusionModelRunner: return nullcontext() # Load model within forward context - with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config): - load_config = LoadConfig() - model_loader = DiffusersPipelineLoader(load_config) - time_before_load = time.perf_counter() + load_config = LoadConfig() + model_loader = DiffusersPipelineLoader(load_config) + time_before_load = time.perf_counter() - with get_memory_context(): - with DeviceMemoryProfiler() as m: - self.pipeline = model_loader.load_model( - od_config=self.od_config, - load_device=load_device, - ) - time_after_load = time.perf_counter() + with get_memory_context(): + with DeviceMemoryProfiler() as m: + self.pipeline = model_loader.load_model( + od_config=self.od_config, + load_device=load_device, + ) + time_after_load = time.perf_counter() logger.info( "Model loading took %.4f GiB and %.6f seconds", diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 0fad90a5b..6f5bd00e6 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -13,10 +13,11 @@ from contextlib import AbstractContextManager, nullcontext import torch import zmq -from vllm.config import VllmConfig +from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger from vllm.utils.mem_utils import GiB_bytes +from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import ( DiffusionOutput, @@ -84,13 +85,16 @@ class DiffusionWorker: current_omni_platform.set_device(self.device) # Create vllm_config for parallel configuration - vllm_config = VllmConfig() + vllm_config = VllmConfig(compilation_config=CompilationConfig()) vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size self.vllm_config = vllm_config # Initialize distributed environment - with set_forward_context(vllm_config=vllm_config, omni_diffusion_config=self.od_config): + with ( + set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config), + set_current_vllm_config(self.vllm_config), + ): init_distributed_environment(world_size=world_size, rank=rank) logger.info(f"Worker {self.rank}: Initialized device and distributed environment.") @@ -104,16 +108,17 @@ class DiffusionWorker: tensor_parallel_size=parallel_config.tensor_parallel_size, pipeline_parallel_size=parallel_config.pipeline_parallel_size, ) + init_workspace_manager(self.device) - # Create model runner and load model - self.model_runner = DiffusionModelRunner( - vllm_config=self.vllm_config, - od_config=self.od_config, - device=self.device, - ) - self.model_runner.load_model( - memory_pool_context_fn=self._maybe_get_memory_pool_context, - ) + # Create model runner and load model + self.model_runner = DiffusionModelRunner( + vllm_config=self.vllm_config, + od_config=self.od_config, + device=self.device, + ) + self.model_runner.load_model( + memory_pool_context_fn=self._maybe_get_memory_pool_context, + ) assert self.model_runner.pipeline is not None self.lora_manager = DiffusionLoRAManager( pipeline=self.model_runner.pipeline, diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 4cd225078..66aebddc1 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/vllm_omni/distributed/omni_connectors/adapter.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: from .connectors.base import OmniConnectorBase from vllm_omni.entrypoints.stage_utils import OmniStageTaskType +from vllm_omni.metrics import OrchestratorAggregator from .utils.logging import get_connector_logger @@ -30,7 +31,7 @@ def try_send_via_connector( sampling_params: Any, original_prompt: Any, next_stage_queue_submit_fn: Callable[[dict[str, Any]], None], - metrics: Any, + metrics: OrchestratorAggregator, ) -> bool: """ Attempts to send data via OmniConnector. diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 1545e17a1..6a3d7393a 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -5,8 +5,6 @@ import copy import time import weakref from collections.abc import AsyncGenerator, Iterable, Sequence -from dataclasses import asdict -from pprint import pformat from typing import Any from vllm.config import VllmConfig @@ -23,9 +21,6 @@ from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ from vllm_omni.distributed.ray_utils.utils import try_close_ray from vllm_omni.engine.input_processor import OmniInputProcessor from vllm_omni.entrypoints.client_request_state import ClientRequestState -from vllm_omni.entrypoints.log_utils import ( - OrchestratorMetrics, -) from vllm_omni.entrypoints.omni import OmniBase from vllm_omni.entrypoints.omni_stage import OmniStage from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType @@ -37,6 +32,7 @@ from vllm_omni.inputs.data import OmniPromptType, OmniSamplingParams # Internal imports (our code) from vllm_omni.lora.request import LoRARequest +from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -298,10 +294,11 @@ class AsyncOmni(OmniBase): ) # Metrics/aggregation helper - metrics = OrchestratorMetrics( - num_stages, - self._enable_stats, - _wall_start_ts, + metrics = OrchestratorAggregator( + num_stages=num_stages, + log_stats=self.log_stats, + wall_start_ts=_wall_start_ts, + final_stage_id_for_e2e=final_stage_id_for_e2e, ) req_state = ClientRequestState(request_id) req_state.metrics = metrics @@ -328,8 +325,6 @@ class AsyncOmni(OmniBase): req_state, metrics, final_stage_id_for_e2e, - _req_start_ts, - _wall_start_ts, ): yield output else: @@ -338,8 +333,6 @@ class AsyncOmni(OmniBase): req_state, metrics, final_stage_id_for_e2e, - _req_start_ts, - _wall_start_ts, sampling_params_list, prompt, ): @@ -348,17 +341,15 @@ class AsyncOmni(OmniBase): logger.debug(f"[{self._name}] Request {request_id} finalized at stage-{final_stage_id_for_e2e}") try: # Finalize E2E metrics if not already done - if str(request_id) not in metrics.e2e_done: - metrics.on_finalize_request( - final_stage_id_for_e2e, - request_id, - _req_start_ts.get(request_id, _wall_start_ts), - ) + metrics.on_finalize_request( + final_stage_id_for_e2e, + request_id, + _req_start_ts.get(request_id, _wall_start_ts), + ) logger.debug(f"[{self._name}] All requests completed") # Summarize and print stats - summary = metrics.build_and_log_summary(final_stage_id_for_e2e) - logger.info("[Summary] %s", pformat(summary, sort_dicts=False)) + metrics.build_and_log_summary() except Exception as e: logger.exception(f"[{self._name}] Request {request_id} Failed to finalized/build/log summary: {e}") finally: @@ -374,10 +365,8 @@ class AsyncOmni(OmniBase): prompt: Any, sampling_params_list: list[SamplingParams], req_state: ClientRequestState, - metrics: OrchestratorMetrics, + metrics: OrchestratorAggregator, final_stage_id_for_e2e: int, - req_start_ts: dict[int, float], - wall_start_ts: float, ) -> AsyncGenerator[OmniRequestOutput, None]: all_stages_finished = {stage_id: False for stage_id in range(final_stage_id_for_e2e + 1)} submit_flag = True @@ -413,16 +402,17 @@ class AsyncOmni(OmniBase): all_stages_finished[stage_id] = finished if output_to_yield: + metrics.record_audio_generated_frames( + output_to_yield, engine_outputs.finished, stage_id, request_id + ) yield output_to_yield async def _process_sequential_results( self, request_id: str, req_state: ClientRequestState, - metrics: OrchestratorMetrics, + metrics: OrchestratorAggregator, final_stage_id_for_e2e: int, - req_start_ts: dict[int, float], - wall_start_ts: float, sampling_params_list: list[SamplingParams], prompt: Any, ) -> AsyncGenerator[OmniRequestOutput, None]: @@ -438,6 +428,9 @@ class AsyncOmni(OmniBase): metrics, ) if output_to_yield: + metrics.record_audio_generated_frames( + output_to_yield, engine_outputs.finished, stage_id, request_id + ) yield output_to_yield if not isinstance(engine_outputs, list): engine_outputs = [engine_outputs] @@ -446,7 +439,9 @@ class AsyncOmni(OmniBase): next_stage_id = stage_id + 1 if next_stage_id <= final_stage_id_for_e2e: next_stage: OmniStage = self.stage_list[next_stage_id] - next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt) + # Derive inputs for the next stage, record postprocess time + with metrics.stage_postprocess_timer(stage_id, request_id): + next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt) sp_next: SamplingParams = sampling_params_list[next_stage_id] # Check if we have a connector for this edge @@ -487,7 +482,7 @@ class AsyncOmni(OmniBase): result: dict[str, Any], stage: OmniStage, stage_id: int, - metrics: OrchestratorMetrics, + metrics: OrchestratorAggregator, ) -> tuple[Any, bool, OmniRequestOutput | None]: """ Process a single result dictionary from a stage. @@ -513,9 +508,16 @@ class AsyncOmni(OmniBase): metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) try: - _m = asdict(result.get("metrics")) - if _m is not None and finished: - metrics.on_stage_metrics(stage_id, req_id, _m) + _m: StageRequestStats = result.get("metrics") + if _m is not None: + # Accumulate generation time + metrics.accumulated_gen_time_ms[req_id][stage_id] += _m.stage_gen_time_ms + + # For diffusion stages, we also accumulate diffusion time + metrics.accumulate_diffusion_metrics(stage.stage_type, req_id, engine_outputs) + + if finished: + metrics.on_stage_metrics(stage_id, req_id, _m, stage.final_output_type) except Exception as e: logger.exception( f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}", diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py index 535f04f7d..11a3c07e1 100644 --- a/vllm_omni/entrypoints/async_omni_diffusion.py +++ b/vllm_omni/entrypoints/async_omni_diffusion.py @@ -98,6 +98,10 @@ class AsyncOmniDiffusion: od_config.model_class_name = "BagelPipeline" od_config.tf_model_config = TransformerConfig() od_config.update_multimodal_support() + elif architectures and len(architectures) == 1: + od_config.model_class_name = architectures[0] + else: + raise # Initialize engine self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index ac03de143..1be3b02f1 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -194,17 +194,17 @@ class OmniServeCommand(CLISubcommand): ) # diffusion model offload parameters - serve_parser.add_argument( + omni_config_group.add_argument( "--enable-cpu-offload", action="store_true", help="Enable CPU offloading for diffusion models.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--enable-layerwise-offload", action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - serve_parser.add_argument( + omni_config_group.add_argument( "--layerwise-num-gpu-layers", type=int, default=1, diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 3e68abb17..1c9103f79 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -1,6 +1,6 @@ import asyncio -from vllm_omni.entrypoints.log_utils import OrchestratorMetrics +from vllm_omni.metrics import OrchestratorAggregator class ClientRequestState: @@ -10,4 +10,4 @@ class ClientRequestState: self.request_id = request_id self.stage_id: int | None = None self.queue = queue if queue is not None else asyncio.Queue() - self.metrics: OrchestratorMetrics | None = None + self.metrics: OrchestratorAggregator | None = None diff --git a/vllm_omni/entrypoints/log_utils.py b/vllm_omni/entrypoints/log_utils.py deleted file mode 100644 index 2031038da..000000000 --- a/vllm_omni/entrypoints/log_utils.py +++ /dev/null @@ -1,591 +0,0 @@ -from __future__ import annotations - -import time -from dataclasses import dataclass -from pprint import pformat -from typing import Any - -from vllm.logger import init_logger - -logger = init_logger(__name__) - - -def log_transfer_tx( - from_stage: int, - to_stage: int, - request_id: Any, - size_bytes: int, - tx_time_ms: float, - used_shm: bool, -) -> None: - logger.info( - pformat( - { - "type": "transfer_stats", - "from_stage": from_stage, - "to_stage": to_stage, - "request_id": request_id, - "size_bytes": int(size_bytes), - "tx_time_ms": float(tx_time_ms), - "tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0), - "used_shm": bool(used_shm), - }, - sort_dicts=False, - ) - ) - - -def log_transfer_rx( - from_stage: int, - to_stage: int, - request_id: Any, - rx_bytes: int, - rx_decode_time_ms: float, - in_flight_time_ms: float, -) -> None: - logger.info( - pformat( - { - "type": "transfer_rx_stats", - "from_stage": from_stage, - "to_stage": to_stage, - "request_id": request_id, - "rx_bytes": int(rx_bytes), - "rx_decode_time_ms": float(rx_decode_time_ms), - "in_flight_time_ms": float(in_flight_time_ms), - "rx_time_per_kb_ms": ( - (float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0 - ), - }, - sort_dicts=False, - ) - ) - - -def log_transfer_total( - from_stage: int, - to_stage: int, - request_id: Any, - size_bytes: int, - tx_time_ms: float, - in_flight_time_ms: float, - rx_decode_time_ms: float, - total_time_ms: float, -) -> None: - logger.info( - pformat( - { - "type": "transfer_total_stats", - "from_stage": from_stage, - "to_stage": to_stage, - "request_id": request_id, - "size_bytes": int(size_bytes), - "tx_time_ms": float(tx_time_ms), - "in_flight_time_ms": float(in_flight_time_ms), - "rx_decode_time_ms": float(rx_decode_time_ms), - "total_time_ms": float(total_time_ms), - "total_time_per_kb_ms": ( - float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0 - ), - }, - sort_dicts=False, - ) - ) - - -def log_stage_request_stats( - stage_id: int, - request_id: Any, - batch_size: int, - num_tokens_out: int, - stage_gen_time_ms: float, - tokens_per_s: float, - rx_transfer_bytes: int, - rx_decode_time_ms: float, - rx_mbps: float, -) -> None: - logger.info( - pformat( - { - "type": "Request_stage_stats", - "stage_id": stage_id, - "request_id": request_id, - "batch_size": int(batch_size), - "num_tokens_out": int(num_tokens_out), - "stage_gen_time_ms": float(stage_gen_time_ms), - "tokens_per_s": float(tokens_per_s), - "rx_transfer_bytes": int(rx_transfer_bytes), - "rx_decode_time_ms": float(rx_decode_time_ms), - "rx_mbps": float(rx_mbps), - }, - sort_dicts=False, - ) - ) - - -def compute_and_log_stage_request_stats( - stage_id: int, - request_id: Any, - batch_size: int, - num_engine_outputs: int, - stage_gen_time_ms: float, - rx_transfer_bytes: int, - rx_decode_time_ms: float, -) -> None: - """Compute per-request metrics and log them in one call.""" - tokens_per_s = (num_engine_outputs * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0 - rx_mbps = ( - (float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0) - if rx_transfer_bytes > 0 - else 0.0 - ) - log_stage_request_stats( - stage_id, - request_id, - int(batch_size), - int(num_engine_outputs), - float(stage_gen_time_ms), - float(tokens_per_s), - int(rx_transfer_bytes), - float(rx_decode_time_ms), - float(rx_mbps), - ) - - -# ----------------- Aggregation helpers for orchestrator ----------------- - - -def record_stage_metrics( - per_request: dict[str, dict[str, Any]], - stage_req_counts: list[int], - stage_total_time_ms: list[float], - stage_total_tokens: list[int], - stage_id: int, - req_id: Any, - metrics: dict[str, Any], -) -> None: - try: - stage_req_counts[stage_id] += 1 - stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0)) - rid_key = str(req_id) - pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) - pr_stages = pr["stages"] # type: ignore[index] - stage_data: dict[str, Any] = { - "stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)), - "num_tokens_out": int(metrics.get("num_tokens_out", 0)), - } - # Only record num_tokens_in for stage 0 (initial prompt) - if stage_id == 0: - stage_data["num_tokens_in"] = int(metrics.get("num_tokens_in", 0)) - stage_total_tokens[stage_id] += int(metrics.get("num_tokens_in", 0)) - pr_stages[stage_id] = stage_data - except Exception: - pass - - -def aggregate_rx_and_maybe_total( - transfer_edge_req: dict[tuple[int, int, str], dict[str, float]], - transfer_agg: dict[tuple[int, int], dict[str, float]], - per_request: dict[str, dict[str, Any]], - stage_id: int, - req_id: Any, - rx_bytes: float, - rx_ms: float, - in_flight_ms: float, -) -> tuple[int, float, float] | None: - try: - # Update RX aggregates for (stage_id-1 -> stage_id) - if stage_id > 0: - key = (stage_id - 1, stage_id) - agg = transfer_agg.get(key) - if agg is None: - agg = { - "sum_bytes": 0.0, - "sum_ms": 0.0, - "count": 0.0, - "sum_rx_bytes": 0.0, - "sum_rx_ms": 0.0, - "rx_count": 0.0, - "sum_total_ms": 0.0, - "total_count": 0.0, - } - transfer_agg[key] = agg - agg["sum_rx_bytes"] += float(rx_bytes) - agg["sum_rx_ms"] += float(rx_ms) - agg["rx_count"] += 1.0 - - # Try combine with sender-side timing if present - rid_key = str(req_id) - s = transfer_edge_req.get((stage_id - 1, stage_id, rid_key)) - if s is None: - return None - tx_ms = float(s.get("tx_ms", 0.0)) - size_b = float(s.get("size_bytes", rx_bytes)) - total_ms = tx_ms + float(in_flight_ms) + float(rx_ms) - agg["sum_total_ms"] += total_ms - agg["total_count"] += 1.0 - # accumulate per-request transfer totals - try: - pr = per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) - pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index] - pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes) # type: ignore[index] - except Exception: - pass - return int(size_b), float(tx_ms), float(total_ms) - return None - except Exception: - return None - - -def record_sender_transfer_agg( - transfer_agg: dict[tuple[int, int], dict[str, float]], - transfer_edge_req: dict[tuple[int, int, str], dict[str, float]], - from_stage: int, - to_stage: int, - req_id: Any, - size_bytes: int, - tx_ms: float, -) -> None: - try: - key = (from_stage, to_stage) - agg = transfer_agg.get(key) - if agg is None: - agg = { - "sum_bytes": 0.0, - "sum_ms": 0.0, - "count": 0.0, - "sum_rx_bytes": 0.0, - "sum_rx_ms": 0.0, - "rx_count": 0.0, - "sum_total_ms": 0.0, - "total_count": 0.0, - } - transfer_agg[key] = agg - agg["sum_bytes"] += float(size_bytes) - agg["sum_ms"] += float(tx_ms) - agg["count"] += 1.0 - # Store sender-side timing for per-request combination - rid_key = str(req_id) - transfer_edge_req[(from_stage, to_stage, rid_key)] = { - "tx_ms": float(tx_ms), - "size_bytes": float(size_bytes), - } - except Exception: - pass - - -def count_tokens_from_outputs(engine_outputs: list[Any]) -> int: - total = 0 - for _ro in engine_outputs: - try: - outs = getattr(_ro, "outputs", None) - if outs and len(outs) > 0: - tokens = getattr(outs[0], "token_ids", None) - if tokens is not None: - total += len(tokens) - except Exception: - pass - return total - - -def build_stage_summary( - stage_req_counts: list[int], - stage_total_tokens: list[int], - stage_total_time_ms: list[float], -) -> list[dict[str, Any]]: - summary: list[dict[str, Any]] = [] - for sid in range(len(stage_req_counts)): - reqs = stage_req_counts[sid] - tokens = stage_total_tokens[sid] - total_ms = float(stage_total_time_ms[sid]) - avg_req = (total_ms / reqs) if reqs > 0 else 0.0 - avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0 - summary.append( - { - "stage_id": sid, - "requests": int(reqs), - "tokens": int(tokens), - "total_time_ms": total_ms, - "avg_time_per_request_ms": avg_req, - "avg_tokens_per_s": avg_tok, - } - ) - return summary - - -def build_transfer_summary( - transfer_agg: dict[tuple[int, int], dict[str, float]], -) -> list[dict[str, Any]]: - summary: list[dict[str, Any]] = [] - for (src, dst), agg in transfer_agg.items(): - sum_bytes = float(agg.get("sum_bytes", 0.0)) - sum_ms = float(agg.get("sum_ms", 0.0)) - samples = int(agg.get("count", 0.0)) - tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 - sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0)) - sum_rx_ms = float(agg.get("sum_rx_ms", 0.0)) - samples_rx = int(agg.get("rx_count", 0.0)) - rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0 - sum_total_ms = float(agg.get("sum_total_ms", 0.0)) - samples_total = int(agg.get("total_count", 0.0)) - total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0 - summary.append( - { - "from_stage": src, - "to_stage": dst, - "samples": samples, - "total_bytes": int(sum_bytes), - "total_time_ms": sum_ms, - "tx_mbps": tx_mbps, - "rx_samples": samples_rx, - "rx_total_bytes": int(sum_rx_bytes), - "rx_total_time_ms": sum_rx_ms, - "rx_mbps": rx_mbps, - "total_samples": samples_total, - "total_transfer_time_ms": sum_total_ms, - "total_mbps": total_mbps, - } - ) - return summary - - -@dataclass -class StageStats: - total_token: int - total_gen_time: float - - -@dataclass -class StageRequestMetrics: - num_tokens_in: int - num_tokens_out: int - stage_gen_time_ms: float - batch_id: int - batch_size: int - rx_decode_time_ms: float - rx_transfer_bytes: int - rx_in_flight_time_ms: float - - stage_stats: StageStats - - -class OrchestratorMetrics: - def __init__( - self, - num_stages: int, - enable_stats: bool, - wall_start_ts: float, - ) -> None: - self.num_stages = int(num_stages) - self.enable_stats = bool(enable_stats) - self.stage_total_time_ms: list[float] = [0.0 for _ in range(self.num_stages)] - self.stage_total_tokens: list[int] = [0 for _ in range(self.num_stages)] - self.stage_req_counts: list[int] = [0 for _ in range(self.num_stages)] - self.transfer_agg: dict[tuple[int, int], dict[str, float]] = {} - self.transfer_edge_req: dict[tuple[int, int, str], dict[str, float]] = {} - self.e2e_total_ms: float = 0.0 - self.e2e_total_tokens: int = 0 - self.e2e_count: int = 0 - self.e2e_done: set[str] = set() - self.per_request: dict[str, dict[str, Any]] = {} - self.sum_per_request_transfer_ms: float = 0.0 - self.wall_start_ts: float = float(wall_start_ts) - self.last_finish_ts: float = float(wall_start_ts) - self.stage_seen_batches: dict[int, set] = {sid: set() for sid in range(self.num_stages)} - self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)] - self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)] - - def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: dict[str, Any]) -> None: - record_stage_metrics( - self.per_request, - self.stage_req_counts, - self.stage_total_time_ms, - self.stage_total_tokens, - stage_id, - req_id, - metrics, - ) - if self.enable_stats: - compute_and_log_stage_request_stats( - stage_id=stage_id, - request_id=req_id, - batch_size=metrics.get("batch_size"), - num_engine_outputs=metrics.get("num_tokens_out"), - stage_gen_time_ms=metrics.get("stage_gen_time_ms"), - rx_decode_time_ms=metrics.get("rx_decode_time_ms"), - rx_transfer_bytes=metrics.get("rx_transfer_bytes"), - ) - if stage_stats := metrics.get("stage_stats", None): - total_token = int(stage_stats.get("total_token")) - total_gen_time = float(stage_stats.get("total_gen_time")) - _avg_tokens_per_s = (total_token * 1000.0 / total_gen_time) if total_gen_time > 0 else 0.0 - logger.info( - pformat( - { - "type": "Stage_running_avg", - "stage_id": stage_id, - "total_tokens": total_token, - "total_gen_time_ms": total_gen_time, - "avg_tokens_per_s": _avg_tokens_per_s, - }, - sort_dicts=False, - ) - ) - try: - batch_id_raw = metrics.get("batch_id", None) - if batch_id_raw is not None: - batch_id = int(batch_id_raw) - if batch_id not in self.stage_seen_batches[stage_id]: - self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0)) - self.stage_seen_batches[stage_id].add(batch_id) - except Exception: - pass - rx_b = float(metrics.get("rx_transfer_bytes", 0.0)) - rx_ms = float(metrics.get("rx_decode_time_ms", 0.0)) - in_flight_ms = float(metrics.get("rx_in_flight_time_ms", 0.0)) - combined = aggregate_rx_and_maybe_total( - self.transfer_edge_req, - self.transfer_agg, - self.per_request, - stage_id, - req_id, - rx_b, - rx_ms, - in_flight_ms, - ) - if self.enable_stats and stage_id > 0: - log_transfer_rx( - stage_id - 1, - stage_id, - req_id, - int(rx_b), - float(rx_ms), - float(in_flight_ms), - ) - if combined is not None: - size_b_c, tx_ms_c, total_ms_c = combined - log_transfer_total( - stage_id - 1, - stage_id, - req_id, - int(size_b_c), - float(tx_ms_c), - float(in_flight_ms), - float(rx_ms), - float(total_ms_c), - ) - - def on_forward( - self, - from_stage: int, - to_stage: int, - req_id: Any, - size_bytes: int, - tx_ms: float, - used_shm: bool, - ) -> None: - # Mark first input time for the destination stage if not set - if self.stage_first_ts[to_stage] is None: - self.stage_first_ts[to_stage] = time.time() - if self.enable_stats: - log_transfer_tx( - from_stage, - to_stage, - req_id, - int(size_bytes), - float(tx_ms), - bool(used_shm), - ) - record_sender_transfer_agg( - self.transfer_agg, - self.transfer_edge_req, - from_stage, - to_stage, - req_id, - int(size_bytes), - float(tx_ms), - ) - - def on_finalize_request( - self, - stage_id: int, - req_id: Any, - req_start_ts: float, - ) -> None: - rid_key = str(req_id) - _t0 = float(req_start_ts) - _t1 = time.time() - # Update last output time for this stage - prev_last = self.stage_last_ts[stage_id] - self.stage_last_ts[stage_id] = _t1 if prev_last is None else max(prev_last, _t1) - self.last_finish_ts = max(self.last_finish_ts, _t1) - e2e_ms = (_t1 - _t0) * 1000.0 - - # Sum tokens from all stages for this request - # Include input tokens from stage 0 + output tokens from all stages - pr = self.per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}) - total_tokens = 0 - stages_info = pr.get("stages", {}) - for sid, stage_data in stages_info.items(): - # Add input tokens only from stage 0 (initial prompt) - if sid == 0: - total_tokens += int(stage_data.get("num_tokens_in", 0)) - total_tokens += int(stage_data.get("num_tokens_out", 0)) - - self.e2e_total_ms += e2e_ms - self.e2e_total_tokens += total_tokens - self.e2e_count += 1 - self.e2e_done.add(rid_key) - per_req_record = { - "type": "request_level_metrics", - "request_id": rid_key, - "e2e_time_ms": e2e_ms, - "e2e_tpt": (e2e_ms / total_tokens) if total_tokens > 0 else 0.0, - "e2e_total_tokens": total_tokens, - "transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)), - "transfers_total_bytes": int(pr.get("transfers_bytes", 0)), - "stages": stages_info, - } - self.sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0)) - logger.info(pformat(per_req_record, sort_dicts=False)) - - def build_and_log_summary(self, final_stage_id_to_prompt: dict[str, int]) -> dict[str, Any]: - # Compute stage summary using wall time between first input and last output per stage - stage_summary: list[dict[str, Any]] = [] - for sid in range(self.num_stages): - first_ts = self.stage_first_ts[sid] - last_ts = self.stage_last_ts[sid] - total_ms = ( - (max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0 - ) - reqs = self.stage_req_counts[sid] - tokens = self.stage_total_tokens[sid] - avg_req = (total_ms / reqs) if reqs > 0 else 0.0 - avg_tok = (tokens * 1000.0 / total_ms) if total_ms > 0 else 0.0 - stage_summary.append( - { - "stage_id": sid, - "requests": int(reqs), - "tokens": int(tokens), - "total_time_ms": float(total_ms), - "avg_time_per_request_ms": float(avg_req), - "avg_tokens_per_s": float(avg_tok), - } - ) - transfer_summary = build_transfer_summary(self.transfer_agg) - e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0 - e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0 - wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0) - summary: dict[str, Any] = { - "e2e_requests": int(self.e2e_count), - "e2e_total_time_ms": float(wall_time_ms), - "e2e_sum_time_ms": float(self.e2e_total_ms), - "e2e_total_tokens": int(self.e2e_total_tokens), - "e2e_avg_time_per_request_ms": e2e_avg_req, - "e2e_avg_tokens_per_s": e2e_avg_tok, - "wall_time_ms": wall_time_ms, - "final_stage_id": final_stage_id_to_prompt, - "stages": stage_summary, - "transfers": transfer_summary, - } - return summary diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 97357dc3b..9e49f951e 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -8,8 +8,6 @@ import uuid import weakref from collections.abc import Callable, Generator, Sequence from concurrent.futures import ThreadPoolExecutor, as_completed -from dataclasses import asdict -from pprint import pformat from typing import Any, Literal, overload from omegaconf import OmegaConf @@ -30,7 +28,6 @@ from vllm_omni.distributed.ray_utils.utils import ( get_ray_queue_class, try_close_ray, ) -from vllm_omni.entrypoints.log_utils import OrchestratorMetrics from vllm_omni.entrypoints.omni_stage import OmniStage from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load @@ -42,6 +39,7 @@ from vllm_omni.entrypoints.utils import ( resolve_model_config_path, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams +from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -240,7 +238,7 @@ class OmniBase: ) # Initialize stats paths - self._enable_stats: bool = bool(log_stats) + self.log_stats: bool = bool(log_stats) self.worker_backend = worker_backend self.ray_address = ray_address @@ -670,10 +668,11 @@ class Omni(OmniBase): final_stage_id_to_prompt[rid] = final_stage_id_for_e2e # Metrics/aggregation helper - metrics = OrchestratorMetrics( + metrics = OrchestratorAggregator( num_stages, - self._enable_stats, + self.log_stats, _wall_start_ts, + final_stage_id_to_prompt, ) it = request_id_to_prompt.items() @@ -740,11 +739,15 @@ class Omni(OmniBase): # Mark last output time for this stage whenever we receive outputs metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time()) try: - _m = result.get("metrics") + _m: StageRequestStats = result.get("metrics") if _m is not None: - if not isinstance(_m, dict): - _m = asdict(_m) - metrics.on_stage_metrics(stage_id, req_id, _m) + # Accumulate generation time + metrics.accumulated_gen_time_ms[req_id][stage_id] += _m.stage_gen_time_ms + + # For diffusion stages, we also accumulate diffusion time + metrics.accumulate_diffusion_metrics(stage.stage_type, req_id, engine_outputs) + + metrics.on_stage_metrics(stage_id, req_id, _m, stage.final_output_type) if pbar: elapsed = pbar.format_dict["elapsed"] or 1e-6 # Aggregate total tokens/images across all stages @@ -781,8 +784,7 @@ class Omni(OmniBase): # End-to-end timing and time-per-token for final output # (only once per request at the designated final stage) try: - rid_key = str(req_id) - if stage_id == final_stage_id_to_prompt[req_id] and rid_key not in metrics.e2e_done: + if stage_id == final_stage_id_to_prompt[req_id]: metrics.on_finalize_request( stage_id, req_id, @@ -792,17 +794,42 @@ class Omni(OmniBase): logger.exception( f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}", ) - yield OmniRequestOutput( + output_to_yield = OmniRequestOutput( stage_id=stage_id, final_output_type=stage.final_output_type, # type: ignore[attr-defined] request_output=engine_outputs, ) + # Record audio generated frames with unified signature + try: + finished = ( + engine_outputs.finished + if hasattr(engine_outputs, "finished") + else ( + engine_outputs[0].finished + if isinstance(engine_outputs, list) + and engine_outputs + and hasattr(engine_outputs[0], "finished") + else False + ) + ) + metrics.record_audio_generated_frames(output_to_yield, finished, stage_id, req_id) + except Exception as e: + logger.exception( + f"[{self._name}] Failed to record audio metrics for req {req_id} at stage {stage_id}: {e}", + ) + + yield output_to_yield + next_stage_id = stage_id + 1 if next_stage_id <= final_stage_id_to_prompt[req_id]: next_stage: OmniStage = self.stage_list[next_stage_id] try: - next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]]) + # Derive inputs for the next stage, record preprocess time + with metrics.stage_postprocess_timer(stage_id, req_id): + next_inputs = next_stage.process_engine_inputs( + self.stage_list, [request_id_to_prompt[req_id]] + ) except Exception as e: logger.exception( f"[{self._name}] Process engine inputs error for req {req_id}" @@ -856,8 +883,7 @@ class Omni(OmniBase): # Summarize and print stats try: - summary = metrics.build_and_log_summary(final_stage_id_to_prompt) - logger.info("[Summary] %s", pformat(summary, sort_dicts=False)) + metrics.build_and_log_summary() except Exception as e: logger.exception(f"[{self._name}] Failed to build/log summary: {e}") diff --git a/vllm_omni/entrypoints/omni_diffusion.py b/vllm_omni/entrypoints/omni_diffusion.py index 5ad9a91c8..9f6dde15b 100644 --- a/vllm_omni/entrypoints/omni_diffusion.py +++ b/vllm_omni/entrypoints/omni_diffusion.py @@ -80,6 +80,8 @@ class OmniDiffusion: od_config.model_class_name = "BagelPipeline" od_config.tf_model_config = TransformerConfig() od_config.update_multimodal_support() + elif architectures and len(architectures) == 1: + od_config.model_class_name = architectures[0] else: raise diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 0356ea231..118a0ba7f 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -81,7 +81,7 @@ class OmniLLM(LLM): self.worker_backend = kwargs.get("worker_backend", "multi_process") self.ray_address = kwargs.get("ray_address", None) self.batch_timeout = batch_timeout - self._enable_stats: bool = bool(log_stats) + self.log_stats: bool = bool(log_stats) # Load stage configurations if stage_configs_path is None: diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 4935c43cc..ca5b2f43e 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -39,7 +39,6 @@ from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_acto from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM -from vllm_omni.entrypoints.log_utils import count_tokens_from_outputs from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion from vllm_omni.entrypoints.omni_llm import OmniLLM from vllm_omni.entrypoints.stage_utils import ( @@ -51,6 +50,7 @@ from vllm_omni.entrypoints.stage_utils import ( set_stage_devices, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams, OmniTokensPrompt +from vllm_omni.metrics import count_tokens_from_outputs from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -844,7 +844,7 @@ def _stage_worker( try: sent_ts = float(t.get("sent_ts", None)) if isinstance(t, dict) else None if sent_ts is not None: - _in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0 + _in_flight_ms_by_rid[rid] = max(0.0, (_recv_dequeue_ts - sent_ts) * 1000.0) else: _in_flight_ms_by_rid[rid] = 0.0 except Exception: @@ -1205,7 +1205,7 @@ async def _stage_worker_async( try: sent_ts = float(task.get("sent_ts", None)) if isinstance(task, dict) else None if sent_ts is not None: - _in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0 + _in_flight_ms_by_rid[rid] = max(0.0, (_recv_dequeue_ts - sent_ts) * 1000.0) else: _in_flight_ms_by_rid[rid] = 0.0 except Exception: @@ -1384,13 +1384,11 @@ def make_request_stats( rx_transfer_bytes: int, rx_in_flight_time_ms: float, ): - from vllm_omni.entrypoints.log_utils import ( - StageRequestMetrics, - ) + from vllm_omni.metrics import StageRequestStats num_tokens_in = count_prompt_tokens_from_outputs(req_output) num_tokens_out = count_tokens_from_outputs(req_output) - return StageRequestMetrics( + return StageRequestStats( num_tokens_in=num_tokens_in, num_tokens_out=num_tokens_out, stage_gen_time_ms=stage_gen_time_ms, @@ -1404,9 +1402,9 @@ def make_request_stats( def make_stage_stats(_agg_total_tokens: int, _agg_total_gen_time_ms: float): - from vllm_omni.entrypoints.log_utils import StageStats + from vllm_omni.metrics import StageStats - return StageStats(total_token=_agg_total_tokens, total_gen_time=_agg_total_gen_time_ms) + return StageStats(total_token=_agg_total_tokens, total_gen_time_ms=_agg_total_gen_time_ms) def output_strip(r_output: RequestOutput | OmniRequestOutput, omni_stage: OmniStage): diff --git a/vllm_omni/metrics/__init__.py b/vllm_omni/metrics/__init__.py new file mode 100644 index 000000000..deceb2333 --- /dev/null +++ b/vllm_omni/metrics/__init__.py @@ -0,0 +1,9 @@ +from .stats import OrchestratorAggregator, StageRequestStats, StageStats +from .utils import count_tokens_from_outputs + +__all__ = [ + "OrchestratorAggregator", + "StageStats", + "StageRequestStats", + "count_tokens_from_outputs", +] diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py new file mode 100644 index 000000000..037db2ea8 --- /dev/null +++ b/vllm_omni/metrics/stats.py @@ -0,0 +1,561 @@ +from __future__ import annotations + +import time +from collections import defaultdict +from collections.abc import Callable +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.metrics.utils import _build_field_defs, _build_row, _format_table + +logger = init_logger(__name__) + + +@dataclass +class StageStats: + total_token: int = 0 + total_gen_time_ms: float = 0.0 + + @property + def avg_tokens_per_s(self) -> float: + return (self.total_token * 1000.0 / self.total_gen_time_ms) if self.total_gen_time_ms > 0 else 0.0 + + +@dataclass +class StageRequestStats: + batch_id: int + batch_size: int + num_tokens_in: int + num_tokens_out: int + stage_gen_time_ms: float + rx_transfer_bytes: int + rx_decode_time_ms: float + rx_in_flight_time_ms: float + stage_stats: StageStats + stage_id: int | None = None + final_output_type: str | None = None + request_id: str | None = None + postprocess_time_ms: float = 0.0 + diffusion_metrics: dict[str, int] = None + audio_generated_frames: int = 0 + + @property + def rx_mbps(self) -> float: + return ( + (float(self.rx_transfer_bytes) * 8.0) / (max(float(self.rx_decode_time_ms), 1e-6) * 1000.0) + if self.rx_transfer_bytes > 0 + else 0.0 + ) + + @property + def tokens_per_s(self) -> float: + return (self.num_tokens_out * 1000.0 / self.stage_gen_time_ms) if (self.stage_gen_time_ms > 0) else 0.0 + + +@dataclass +class TransferEdgeStats: + from_stage: int + to_stage: int + request_id: str + size_bytes: int + tx_time_ms: float + used_shm: bool = False + rx_decode_time_ms: float = 0.0 + in_flight_time_ms: float = 0.0 + + @property + def total_time_ms(self) -> float: + return float(self.tx_time_ms) + float(self.rx_decode_time_ms) + float(self.in_flight_time_ms) + + +@dataclass +class RequestE2EStats: + request_id: str + e2e_total_ms: float + e2e_total_tokens: int + transfers_total_time_ms: float + transfers_total_bytes: int + + @property + def e2e_tpt(self) -> float: + return (self.e2e_total_ms / self.e2e_total_tokens) if self.e2e_total_tokens > 0 else 0.0 + + +# === Field Configuration === +# Fields requiring unit conversion: original_field_name -> (display_name, transform_fn) +FIELD_TRANSFORMS: dict[str, tuple[str, Callable[[Any], Any]]] = { + "rx_transfer_bytes": ("rx_transfer_kbytes", lambda v: v / 1024.0), + "size_bytes": ("size_kbytes", lambda v: v / 1024.0), + "transfers_total_bytes": ("transfers_total_kbytes", lambda v: v / 1024.0), +} + +# Fields to exclude from table display for each event type +STAGE_EXCLUDE = { + "stage_stats", + "stage_id", + "request_id", + "rx_transfer_bytes", + "rx_decode_time_ms", + "rx_in_flight_time_ms", + "final_output_type", +} +TRANSFER_EXCLUDE = {"from_stage", "to_stage", "request_id", "used_shm"} +E2E_EXCLUDE = {"request_id"} + +# Decide the order of overall summary fields, or None for auto +OVERALL_FIELDS: list[str] | None = None +STAGE_FIELDS = _build_field_defs(StageRequestStats, STAGE_EXCLUDE, FIELD_TRANSFORMS) +TRANSFER_FIELDS = _build_field_defs(TransferEdgeStats, TRANSFER_EXCLUDE, FIELD_TRANSFORMS) +E2E_FIELDS = _build_field_defs(RequestE2EStats, E2E_EXCLUDE, FIELD_TRANSFORMS) + + +class OrchestratorAggregator: + def __init__( + self, + num_stages: int, + log_stats: bool, + wall_start_ts: float, + final_stage_id_for_e2e: dict[str, int] | int, + ) -> None: + self.num_stages = int(num_stages) + self.log_stats = bool(log_stats) + self.final_stage_id_for_e2e = final_stage_id_for_e2e + self.init_run_state(wall_start_ts) + self.stage_events: dict[str, list[StageRequestStats]] = {} + self.transfer_events: dict[ + tuple[int, int, str], TransferEdgeStats + ] = {} # Key: (from_stage, to_stage, request_id) + self.e2e_events: list[RequestE2EStats] = [] + + def init_run_state(self, wall_start_ts: float) -> None: + # Per-run aggregates and timing state + self.stage_total_tokens = [0 for _ in range(self.num_stages)] + self.e2e_total_ms = 0.0 + self.e2e_total_tokens = 0 + self.e2e_count = 0 + self.e2e_done = set() + self.wall_start_ts = float(wall_start_ts) + self.last_finish_ts = float(wall_start_ts) + self.stage_first_ts = [None for _ in range(self.num_stages)] + self.stage_last_ts = [None for _ in range(self.num_stages)] + self.accumulated_gen_time_ms: defaultdict[str, defaultdict[int, float]] = defaultdict( + lambda: defaultdict(float) + ) # {request_id: {stage_id:accumulated_gen_time_ms}} + self.diffusion_metrics: defaultdict[str, defaultdict[str, float]] = defaultdict( + lambda: defaultdict(float) + ) # {request_id: {diffusion_metrics_key: accumulated_metrics_data}} + + def _get_or_create_transfer_event( + self, + from_stage: int, + to_stage: int, + request_id: str, + ) -> TransferEdgeStats: + key = (from_stage, to_stage, request_id) + evt = self.transfer_events.get(key) + if evt is None: + evt = TransferEdgeStats( + from_stage=from_stage, + to_stage=to_stage, + request_id=request_id, + size_bytes=0, + tx_time_ms=0.0, + used_shm=False, + rx_decode_time_ms=0.0, + in_flight_time_ms=0.0, + ) + self.transfer_events[key] = evt + return evt + + def record_transfer_tx( + self, + from_stage: int, + to_stage: int, + request_id: Any, + size_bytes: int, + tx_time_ms: float, + used_shm: bool, + ) -> TransferEdgeStats | None: + try: + evt = self._get_or_create_transfer_event( + int(from_stage), + int(to_stage), + str(request_id), + ) + # Accumulate tx metrics + evt.size_bytes += int(size_bytes) + evt.tx_time_ms += float(tx_time_ms) + evt.used_shm = evt.used_shm or bool(used_shm) + return evt + except Exception: + return None + + def record_transfer_rx( + self, + stats: StageRequestStats, + ) -> TransferEdgeStats | None: + try: + if stats.stage_id is None or stats.stage_id <= 0: + return None + from_stage = int(stats.stage_id) - 1 + to_stage = int(stats.stage_id) + rid_key = str(stats.request_id) + evt = self._get_or_create_transfer_event(from_stage, to_stage, rid_key) + # Accumulate rx metrics + if evt.size_bytes == 0: + # size_bytes has been recorded in tx phase + evt.size_bytes = int(stats.rx_transfer_bytes) + evt.rx_decode_time_ms += float(stats.rx_decode_time_ms) + evt.in_flight_time_ms += float(stats.rx_in_flight_time_ms) + return evt + except Exception: + return None + + def record_audio_generated_frames( + self, + output_to_yield: Any, + finished: bool, + stage_id: int, + request_id: str, + ) -> None: + if ( + output_to_yield.final_output_type == "audio" + and finished + and (multimodal_output := output_to_yield.multimodal_output.get("audio")) is not None + ): + nframes = int(multimodal_output[-1].shape[0]) + stage_events_for_req = self.stage_events.get(request_id, []) + if stage_events_for_req: + for stage_event in stage_events_for_req: + if stage_event.stage_id == stage_id: + stage_event.audio_generated_frames += nframes + break + else: + logger.warning( + "Failed to record audio generated frames for request %s at stage %s: no stage event found", + request_id, + stage_id, + ) + + def _as_stage_request_stats( + self, + stage_id: int, + req_id: str, + metrics: StageRequestStats, + final_output_type: str | None = None, + ) -> StageRequestStats: + "Convert dict to StageRequestStats if needed." + stats = metrics + stats.stage_id = stage_id + stats.request_id = req_id + stats.final_output_type = final_output_type + stats.diffusion_metrics = ( + {k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()} + if req_id in self.diffusion_metrics + else None + ) + return stats + + def on_stage_metrics( + self, + stage_id: int, + req_id: Any, + metrics: StageRequestStats, + final_output_type: str | None = None, + ) -> None: + stats = self._as_stage_request_stats(stage_id, req_id, metrics, final_output_type) + self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_out) + if stats.stage_id == 0: + self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_in) + self.stage_events.setdefault(str(stats.request_id), []).append(stats) + + self.record_transfer_rx(stats) + + def record_stage_postprocess_time(self, stage_id: int, req_id: Any, postproc_time_ms: float) -> None: + if req_id in self.stage_events: + for stats in self.stage_events[req_id]: + if stats.stage_id == stage_id: + stats.postprocess_time_ms = float(postproc_time_ms) + break + else: + logger.warning( + "Failed to record postprocess time for request %s at stage %s: no stage event found", + req_id, + stage_id, + ) + + @contextmanager + def stage_postprocess_timer(self, stage_id: int, req_id: Any): + """Context manager for measuring and recording stage postprocessing time. + + Usage: + with metrics.stage_postprocess_timer(stage_id, request_id): + next_inputs = next_stage.process_engine_inputs(...) + """ + _t0 = time.perf_counter() + try: + yield + finally: + _postproc_ms = (time.perf_counter() - _t0) * 1000.0 + self.record_stage_postprocess_time(stage_id, req_id, _postproc_ms) + + def accumulate_diffusion_metrics(self, stage_type: str, req_id: Any, engine_outputs: Any) -> None: + """Accumulate diffusion metrics for a request. + + Handles extraction and accumulation of diffusion stage metrics. + + Args: + req_id: Request ID + engine_outputs: Engine output object containing metrics + """ + if stage_type != "diffusion": + return + engine_output = engine_outputs[0] if isinstance(engine_outputs, list) and engine_outputs else engine_outputs + diffusion_metrics: dict = getattr(engine_output, "metrics", {}) + if isinstance(diffusion_metrics, list): + diffusion_metrics = diffusion_metrics[0] + if diffusion_metrics: + for key, value in diffusion_metrics.items(): + self.diffusion_metrics[req_id][key] += value + + def on_forward( + self, + from_stage: int, + to_stage: int, + req_id: Any, + size_bytes: int, + tx_ms: float, + used_shm: bool, + ) -> None: + # Mark first input time for the destination stage if not set + if self.stage_first_ts[to_stage] is None: + self.stage_first_ts[to_stage] = time.time() + self.record_transfer_tx( + from_stage=from_stage, + to_stage=to_stage, + request_id=req_id, + size_bytes=size_bytes, + tx_time_ms=tx_ms, + used_shm=used_shm, + ) + + def on_finalize_request( + self, + stage_id: int, + req_id: Any, + req_start_ts: float, + ) -> None: + rid_key = str(req_id) + if rid_key in self.e2e_done: + return # Already finalized + _t0 = float(req_start_ts) + _t1 = time.time() + # Update last output time for this stage + prev_last = self.stage_last_ts[stage_id] + self.stage_last_ts[stage_id] = _t1 if prev_last is None else max(prev_last, _t1) + self.last_finish_ts = max(self.last_finish_ts, _t1) + e2e_ms = (_t1 - _t0) * 1000.0 + + # Sum tokens from all stages for this request + # Include input tokens from stage 0 + output tokens from all stages + total_tokens = 0 + if rid_key in self.stage_events: + for evt in self.stage_events[rid_key]: + if evt.stage_id == 0: + total_tokens += int(evt.num_tokens_in) + total_tokens += int(evt.num_tokens_out) + + self.e2e_total_ms += e2e_ms + self.e2e_total_tokens += total_tokens + self.e2e_count += 1 + self.e2e_done.add(rid_key) + per_req_record = RequestE2EStats( + request_id=rid_key, + e2e_total_ms=e2e_ms, + e2e_total_tokens=total_tokens, + transfers_total_time_ms=float( + sum(evt.total_time_ms for evt in self.transfer_events.values() if evt.request_id == rid_key) + ), + transfers_total_bytes=int( + sum(evt.size_bytes for evt in self.transfer_events.values() if evt.request_id == rid_key) + ), + ) + self.e2e_events.append(per_req_record) + + def build_and_log_summary(self) -> dict[str, Any]: + if not self.log_stats: + return {} + wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0) + e2e_avg_req = (wall_time_ms / self.e2e_count) if self.e2e_count > 0 else 0.0 + e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / wall_time_ms) if wall_time_ms > 0 else 0.0 + + if isinstance(self.final_stage_id_for_e2e, int): + final_stage_id_map: dict[str, int] = {"*": int(self.final_stage_id_for_e2e)} + else: + final_stage_id_map = self.final_stage_id_for_e2e + + stage_wall_time_ms = [ + ((self.stage_last_ts[i] - self.stage_first_ts[i]) * 1000.0) + if (self.stage_first_ts[i] is not None and self.stage_last_ts[i] is not None) + else 0.0 + for i in range(self.num_stages) + ] + + overall_summary = { + "e2e_requests": int(self.e2e_count), + "e2e_wall_time_ms": float(wall_time_ms), + "e2e_total_tokens": int(self.e2e_total_tokens), + "e2e_avg_time_per_request_ms": float(e2e_avg_req), + "e2e_avg_tokens_per_s": float(e2e_avg_tok), + } + # Add stage_wall_time_ms as separate fields for each stage + for idx, wall_time in enumerate(stage_wall_time_ms): + overall_summary[f"e2e_stage_{idx}_wall_time_ms"] = wall_time + + # Print overall summary + # filter out all-zero fields for logging + overall_fields = [] + for k in OVERALL_FIELDS or list(overall_summary.keys()): + v = overall_summary.get(k, None) + if v not in (0, 0.0, 0.000, None, ""): + overall_fields.append(k) + if overall_fields: + logger.info( + "\n%s", + _format_table("Overall Summary", overall_summary, overall_fields), + ) + + all_request_ids = sorted(set(self.stage_events.keys()) | {e.request_id for e in self.e2e_events}) + + result_stage_table = [] + result_trans_table = [] + result_e2e_table = [] + + for rid in all_request_ids: + # === E2E table (single column) === + e2e_evt = next((e for e in self.e2e_events if e.request_id == rid), None) + if e2e_evt: + e2e_data = _build_row(e2e_evt, E2E_FIELDS) + result_e2e_table.append({"request_id": rid, **e2e_data}) + + # filter out all-zero fields for logging + nonzero_e2e_fields = set() + for k, v in e2e_data.items(): + if v not in (0, 0.000, None, ""): + nonzero_e2e_fields.add(k) + value_fields_e2e = sorted(nonzero_e2e_fields) + + if value_fields_e2e: + logger.info( + "\n%s", + _format_table( + f"RequestE2EStats [request_id={rid}]", + e2e_data, + value_fields=value_fields_e2e, + ), + ) + + # === Stage table (columns = stage_id) === + stage_evts = sorted( + self.stage_events.get(rid, []), + key=lambda e: e.stage_id if e.stage_id is not None else -1, + ) + # if any stage has diffusion_metrics, remove postprocess_time_ms field + # because it is already included in diffusion_metrics + local_exclude = STAGE_EXCLUDE.copy() + has_diffusion_metrics = any(getattr(evt, "diffusion_metrics", None) for evt in stage_evts) + if has_diffusion_metrics: + local_exclude.add("postprocess_time_ms") + local_stage_fields = _build_field_defs(StageRequestStats, local_exclude, FIELD_TRANSFORMS) + + # if diffusion_metrics is present, expand it into multiple columns + # then remove diffusion_metrics from the table + stage_rows = [] + for evt in stage_evts: + row = {"stage_id": evt.stage_id, **_build_row(evt, local_stage_fields)} + if evt.diffusion_metrics: + row.update(evt.diffusion_metrics) + row.pop("diffusion_metrics", None) # Remove the dict itself + stage_rows.append(row) + + result_stage_table.append({"request_id": rid, "stages": stage_rows}) + + if stage_rows: + # filter out all-zero fields for logging + all_value_fields = set() + for row in stage_rows: + for k in row.keys(): + if k != "stage_id": + all_value_fields.add(k) + value_fields_list = [] + for field in sorted(all_value_fields): + all_zero = True + for row in stage_rows: + v = row.get(field, None) + if v not in (0, 0.0, 0.000, None, ""): + all_zero = False + break + if not all_zero: + value_fields_list.append(field) + + if value_fields_list: + logger.info( + "\n%s", + _format_table( + f"StageRequestStats [request_id={rid}]", + stage_rows, + column_key="stage_id", + value_fields=value_fields_list, + ), + ) + + # === Transfer table (columns = edge) === + transfer_evts = sorted( + [e for e in self.transfer_events.values() if e.request_id == rid], + key=lambda e: (e.from_stage, e.to_stage), + ) + transfer_rows = [ + {"edge": f"{evt.from_stage}->{evt.to_stage}", **_build_row(evt, TRANSFER_FIELDS)} + for evt in transfer_evts + ] + result_trans_table.append({"request_id": rid, "transfers": transfer_rows}) + + if transfer_rows: + # filter out all-zero fields for logging + all_value_fields = set() + for row in transfer_rows: + for k in row.keys(): + if k != "edge": + all_value_fields.add(k) + value_fields_list = [] + for field in sorted(all_value_fields): + all_zero = True + for row in transfer_rows: + v = row.get(field, None) + if v not in (0, 0.0, 0.000, None, ""): + all_zero = False + break + if not all_zero: + value_fields_list.append(field) + + if value_fields_list: + logger.info( + "\n%s", + _format_table( + f"TransferEdgeStats [request_id={rid}]", + transfer_rows, + column_key="edge", + value_fields=value_fields_list, + ), + ) + + return { + "final_stage_id": final_stage_id_map, + "overall_summary": overall_summary, + "stage_table": result_stage_table, + "trans_table": result_trans_table, + "e2e_table": result_e2e_table, + } diff --git a/vllm_omni/metrics/utils.py b/vllm_omni/metrics/utils.py new file mode 100644 index 000000000..ea83d959e --- /dev/null +++ b/vllm_omni/metrics/utils.py @@ -0,0 +1,143 @@ +from collections.abc import Callable +from dataclasses import fields +from typing import Any + +from prettytable import PrettyTable + + +def _build_field_defs( + cls: type, + exclude: set[str], + transforms: dict[str, tuple[str, Callable]] | None = None, +) -> list[tuple[str, Callable[[Any], Any]]]: + """Auto-generate field definitions from a dataclass. + + Args: + cls: The dataclass type to extract fields from. + exclude: Set of field names to exclude from output. + transforms: Optional mapping of field transformations. + Format: {original_name: (display_name, transform_fn)} + + Returns: + List of (display_name, getter_fn) tuples for table generation. + """ + transforms = transforms or {} + result = [] + for f in fields(cls): + if f.name in exclude: + continue + if f.name in transforms: + display_name, transform_fn = transforms[f.name] + # Capture variables in closure to avoid late binding issues + result.append((display_name, lambda e, fn=transform_fn, n=f.name: fn(getattr(e, n)))) + else: + result.append((f.name, lambda e, n=f.name: getattr(e, n))) + return result + + +def _build_row(evt: Any, field_defs: list[tuple[str, Callable]]) -> dict[str, Any]: + """Build a row dict from an event object using field definitions. + + Args: + evt: The event object (dataclass instance). + field_defs: List of (field_name, getter_fn) tuples. + + Returns: + Dict mapping field names to their values. + """ + return {name: getter(evt) for name, getter in field_defs} + + +def _get_field_names(field_defs: list[tuple[str, Callable]]) -> list[str]: + """Extract field names from field definitions. + + Args: + field_defs: List of (field_name, getter_fn) tuples. + + Returns: + List of field names. + """ + return [name for name, _ in field_defs] + + +def _format_table( + title: str, + data: dict[str, Any] | list[dict[str, Any]], + value_fields: list[str], + column_key: str | None = None, + column_prefix: str = "", +) -> str: + """Format a table for display. + + Supports two modes: + 1. Single-column mode: data is a dict, displays as Field | Value + 2. Multi-column mode: data is a list of dicts, displays as Field | col1 | col2 | ... + + Args: + title: Table title. + data: Either a single dict (single-column) or list of dicts (multi-column). + value_fields: List of field names to display as rows. + column_key: Key in each dict used as column header (required for multi-column mode). + column_prefix: Optional prefix for column headers (multi-column mode only). + + Returns: + Formatted table string. + """ + if not data: + return f"[{title}] " + + def _format_value(value: Any) -> str: + if isinstance(value, bool): + return str(value) + if isinstance(value, int): + return f"{value:,}" + if isinstance(value, float): + return f"{value:,.3f}" + if isinstance(value, list): + return ", ".join(str(f"{v:,.3f}") for v in value) + return str(value) + + table = PrettyTable() + + # Single-column mode: data is a dict + if isinstance(data, dict): + table.field_names = ["Field", "Value"] + table.align["Field"] = "l" + table.align["Value"] = "r" + for field in value_fields: + if field in data: + if isinstance(data[field], dict): + for sub_key, sub_value in data[field].items(): + table.add_row([f"{sub_key}", _format_value(sub_value)]) + else: + table.add_row([field, _format_value(data[field])]) + + # Multi-column mode: data is a list of dicts + else: + if column_key is None: + raise ValueError("column_key is required for multi-column mode") + col_headers = [f"{column_prefix}{row.get(column_key, '?')}" for row in data] + table.field_names = ["Field"] + col_headers + table.align["Field"] = "l" + for col in col_headers: + table.align[col] = "r" + for field in value_fields: + row_values = [_format_value(r.get(field, "")) for r in data] + table.add_row([field] + row_values) + + return "\n".join([f"[{title}]", table.get_string()]) + + +def count_tokens_from_outputs(engine_outputs: list[Any]) -> int: + total = 0 + for _ro in engine_outputs: + try: + outs = getattr(_ro, "outputs", None) + if outs and len(outs) > 0: + tokens = getattr(outs[0], "token_ids", None) + if tokens is not None: + total += len(tokens) + except Exception: + # Ignore any issues with individual outputs to keep token counting best-effort. + pass + return total