diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index bba4e2513..8edafec1e 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -64,6 +64,7 @@ th { | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | | | `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | | +| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | | | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | | | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 702ea4e87..18792a9d6 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -127,6 +127,7 @@ The following tables show which models support each feature: | **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **SenseNova-U1** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | | **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ | | **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/examples/offline_inference/sensenova_u1/README.md b/examples/offline_inference/sensenova_u1/README.md new file mode 100644 index 000000000..8052c90d3 --- /dev/null +++ b/examples/offline_inference/sensenova_u1/README.md @@ -0,0 +1,208 @@ +# SenseNova-U1-8B-MoT + +## Architecture + +SenseNova-U1 is a unified Qwen3-based LLM with Mixture-of-Tokenizers (MoT) attention. Unlike two-stage pipelines (e.g., BAGEL), it handles text encoding, optional reasoning ("think mode"), and flow-matching-based image denoising entirely within a single diffusion stage. + +| Feature | Description | +| :------ | :---------- | +| **Model type** | Unified LLM (single-stage diffusion pipeline) | +| **Base LLM** | Qwen3 with MoT attention and 3D RoPE | +| **Image generation** | Flow-matching Euler sampler, no separate VAE | +| **Think mode** | Optional chain-of-thought reasoning before image generation | +| **Parallelism** | Tensor Parallelism (TP) with fused QKV and fused gate/up projections | +| **Modalities** | text2img, img2img, img2text (understanding), text2text | + +## Quick Start + +```bash +cd examples/offline_inference/sensenova_u1 + +# Text-to-image +python end2end.py --prompt "A cute cat" --think + +# Image-to-image editing +python end2end.py --prompt "Turn this into an oil painting" \ + --image input.png --think + +# Image understanding (img2text) +python end2end.py --modality img2text \ + --prompt "Describe this image in detail" \ + --image photo.jpg + +# Text-to-text (pure chat) +python end2end.py --modality text2text \ + --prompt "What is the capital of France?" + +# Custom resolution (image generation) +python end2end.py --prompt "A futuristic cityscape at sunset" \ + --width 2048 --height 1024 --think +``` + +> **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU. + +## Text-to-Image (t2i) + +Standard text-to-image with optional think mode: + +```bash +python end2end.py \ + --prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light" \ + --width 1536 --height 2720 \ + --think --print-think +``` + +## Image-to-Image Editing (img2img) + +Pass one or more `--image` paths to trigger img2img mode. The model uses the +input image(s) as visual context for editing: + +```bash +# Single input image +python end2end.py \ + --prompt "Add a sunset sky in the background" \ + --image photo.jpg \ + --width 2048 --height 2048 \ + --think + +# Multiple reference images +python end2end.py \ + --prompt "Combine the style of Image-1 with the content of Image-2" \ + --image style_ref.png content_ref.png \ + --width 2048 --height 2048 \ + --think +``` + +### Dual CFG for img2img + +img2img supports dual classifier-free guidance (text CFG + image CFG): + +```bash +python end2end.py \ + --prompt "Make the person smile" \ + --image portrait.jpg \ + --cfg-scale 4.0 \ + --img-cfg-scale 2.0 \ + --think +``` + +| `img_cfg_scale` | Behavior | +| :-------------- | :------- | +| `1.0` (default) | Text-only CFG: guidance = image-condition → full-condition | +| `== cfg_scale` | Standard CFG: guidance = unconditional → full-condition | +| Other value | Dual CFG: separate text and image guidance strengths | + +## Full Parameter Reference + +### General Parameters + +| Parameter | Default | Description | +| :-------- | :------ | :---------- | +| `--modality` | `auto` | Task modality: `auto`, `text2img`, `img2img`, `img2text`, `text2text` | +| `--prompt` | "A cute cat..." | Text prompt / editing instruction / question | +| `--image` | None | Input image path(s) for img2img or img2text | + +### Image Generation Parameters (text2img / img2img) + +| Parameter | Default | Description | +| :-------- | :------ | :---------- | +| `--height` | 2048 | Height of generated image (pixels) | +| `--width` | 2048 | Width of generated image (pixels) | +| `--seed` | 42 | Random seed for reproducibility | +| `--num-steps` | 50 | Number of denoising steps | +| `--cfg-scale` | 4.0 | Text classifier-free guidance scale | +| `--img-cfg-scale` | 1.0 | Image CFG scale (img2img only, 1.0 = disabled) | +| `--cfg-norm` | "none" | CFG normalization: none, global, channel, cfg_zero_star (t2i only) | +| `--timestep-shift` | 3.0 | Timestep shift for flow-matching schedule | +| `--t-eps` | 0.02 | Epsilon for timestep schedule | +| `--think` | False | Enable think mode | +| `--print-think` | False | Print think text to stdout | + +### Text Generation Parameters (text2text / img2text) + +| Parameter | Default | Description | +| :-------- | :------ | :---------- | +| `--max-tokens` | 512 | Maximum number of tokens to generate | +| `--do-sample` | False | Use sampling instead of greedy decoding | +| `--temperature` | 0.7 | Sampling temperature (higher = more diverse) | + +### Infrastructure Parameters + +| Parameter | Default | Description | +| :-------- | :------ | :---------- | +| `--model` | `SenseNova/SenseNova-U1-8B-MoT` | HuggingFace model ID or local path | +| `--output` | `.` | Output directory for saved images | +| `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism | +| `--enforce-eager` | False | Disable torch.compile | +| `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM | + +## Reducing GPU Memory Usage + +For hardware with limited VRAM, enable **module-wise CPU offload** with +`--enable-cpu-offload`. The pipeline implements `SupportsModuleOffload`, so the +vision encoder (`vision_model`) and the Qwen3 LLM (`language_model`) are +swapped between CPU and GPU on demand: + +- During text/vision encoding, the LLM is on CPU. +- During the diffusion loop, the vision encoder is on CPU. +- Lightweight FM modules stay resident on GPU. + +This lowers peak VRAM at the cost of extra CPU<->GPU transfers (use pinned +memory) — useful for running SenseNova-U1 on consumer-grade GPUs. + +```bash +# Text-to-image with CPU offload +python end2end.py \ + --prompt "A cute cat sitting on a windowsill" \ + --width 2048 --height 2048 \ + --enable-cpu-offload --think + +# Image-to-image editing with CPU offload +python end2end.py \ + --prompt "Turn this into an oil painting" \ + --image input.png \ + --width 2048 --height 2048 \ + --enable-cpu-offload --think + +# Image understanding (img2text) with CPU offload +python end2end.py \ + --modality img2text \ + --prompt "Describe this image in detail" \ + --image photo.jpg \ + --enable-cpu-offload +``` + +> **Notes** +> - CPU offload is **single-GPU only** (incompatible with `--tensor-parallel-size > 1`). +> - First-step latency is higher because of the cold-start CPU<->GPU transfers. +> - For more details and other offloading strategies, see +> [CPU Offloading for Diffusion Models](../../../docs/user_guide/diffusion/cpu_offload_diffusion.md). + +## Reproducing the E2E Test + +The following command reproduces the pixel-validated CI test case: + +```bash +python end2end.py \ + --prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look. The portrait should feel polished and natural, with sharp eyes, realistic skin texture, accurate facial anatomy, and premium lighting that keeps the face as the main focus." \ + --width 1536 --height 2720 \ + --seed 42 --num-steps 50 \ + --cfg-scale 4.0 --timestep-shift 3.0 --cfg-norm none \ + --think --print-think \ + --output outputs +``` + +The corresponding pytest: + +```bash +pytest -s -v tests/e2e/offline_inference/test_sensenova_u1_text2img.py \ + -m "advanced_model" --run-level "advanced_model" +``` + +The img2img counterpart lives at +`tests/e2e/offline_inference/test_sensenova_u1_img2img.py`. + +## Online Serving + +For OpenAI-compatible API serving, see +[`examples/online_serving/sensenova_u1/`](../../examples/online_serving/sensenova_u1/). diff --git a/examples/offline_inference/sensenova_u1/end2end.py b/examples/offline_inference/sensenova_u1/end2end.py new file mode 100644 index 000000000..12f28da6f --- /dev/null +++ b/examples/offline_inference/sensenova_u1/end2end.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end offline inference for SenseNova-U1-8B-MoT. + +Supports four modalities: text2img, img2img, img2text, text2text. + +Text-to-image: + python end2end.py --prompt "A cute cat" --think + +Image-to-image (editing): + python end2end.py --prompt "Turn this into an oil painting" \ + --image input.png --think + +Image understanding: + python end2end.py --modality img2text \ + --prompt "Describe this image" --image photo.jpg + +Text chat: + python end2end.py --modality text2text \ + --prompt "What is the capital of France?" + +See README.md for more examples. +""" + +import argparse +import os + +from PIL import Image + +from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + +def parse_args(): + parser = argparse.ArgumentParser( + description="SenseNova-U1 text-to-image / image-to-image / text / understanding via vLLM-Omni.", + ) + parser.add_argument( + "--model", + default="SenseNova/SenseNova-U1-8B-MoT", + help="HuggingFace model ID or local path.", + ) + parser.add_argument( + "--modality", + default="auto", + choices=["auto", "text2img", "img2img", "img2text", "text2text"], + help="Task modality. 'auto' infers from --image (img2img/img2text if images, else t2i/t2t).", + ) + parser.add_argument( + "--prompt", + default="A cute cat sitting on a windowsill, soft natural light", + help="Text prompt for generation or editing instruction.", + ) + parser.add_argument( + "--image", + nargs="+", + metavar="PATH", + default=None, + help="Input image path(s) for img2img or img2text.", + ) + parser.add_argument( + "--output", + type=str, + default=".", + help="Output directory to save generated images.", + ) + # Image dimensions + parser.add_argument("--height", type=int, default=2048, help="Height of generated image.") + parser.add_argument("--width", type=int, default=2048, help="Width of generated image.") + + # Generation parameters + parser.add_argument("--seed", type=int, default=42, help="Random seed for deterministic results.") + parser.add_argument("--num-steps", type=int, default=50, help="Number of denoising steps.") + parser.add_argument("--cfg-scale", type=float, default=4.0, help="Text classifier-free guidance scale.") + parser.add_argument( + "--img-cfg-scale", + type=float, + default=1.0, + help="Image CFG scale for img2img (1.0 = disabled). " + "When both --cfg-scale and --img-cfg-scale differ, dual CFG is used.", + ) + parser.add_argument( + "--cfg-norm", + type=str, + default="none", + choices=["none", "global", "channel", "cfg_zero_star"], + help="CFG normalization mode. cfg_zero_star is only for t2i.", + ) + parser.add_argument( + "--timestep-shift", + type=float, + default=3.0, + help="Timestep shift for flow-matching schedule.", + ) + parser.add_argument( + "--t-eps", + type=float, + default=0.02, + help="Epsilon for flow-matching timestep schedule.", + ) + + # Text generation parameters (for text2text / img2text) + parser.add_argument( + "--max-tokens", + type=int, + default=2048, + help="Max tokens for text generation (text2text / img2text).", + ) + parser.add_argument( + "--do-sample", + action="store_true", + help="Use sampling instead of greedy decoding for text generation.", + ) + parser.add_argument( + "--temperature", + type=float, + default=0.7, + help="Sampling temperature for text generation.", + ) + + # Think mode + parser.add_argument( + "--think", + action="store_true", + help="Enable think mode: the model reasons about the prompt before generating the image.", + ) + parser.add_argument( + "--print-think", + action="store_true", + help="Print the generated think text to stdout.", + ) + + # Advanced + parser.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help="Number of GPUs for tensor parallelism.", + ) + parser.add_argument( + "--enforce-eager", + action="store_true", + help="Disable torch.compile and force eager execution.", + ) + parser.add_argument( + "--enable-cpu-offload", + action="store_true", + help=( + "Enable model-level (sequential) CPU offloading. " + "Mutually exclusive between vision_model and language_model: " + "only one is on GPU at a time, reducing peak VRAM at the cost " + "of extra CPU<->GPU transfers." + ), + ) + + nullify_stage_engine_defaults(parser) + return parser.parse_args() + + +def _resolve_modality(args): + """Determine the effective modality from CLI args.""" + if args.modality != "auto": + return args.modality + has_images = args.image is not None + # Default: image-producing tasks for backward compatibility + return "img2img" if has_images else "text2img" + + +def main(): + args = parse_args() + os.makedirs(args.output, exist_ok=True) + + modality = _resolve_modality(args) + is_text_output = modality in ("text2text", "img2text") + has_images = modality in ("img2img", "img2text") + + omni = Omni( + model=args.model, + tensor_parallel_size=args.tensor_parallel_size, + enforce_eager=args.enforce_eager, + enable_cpu_offload=args.enable_cpu_offload, + ) + + extra_args = { + "cfg_scale": args.cfg_scale, + "cfg_norm": args.cfg_norm, + "timestep_shift": args.timestep_shift, + "cfg_interval": (0.0, 1.0), + "batch_size": 1, + "think": args.think, + "t_eps": args.t_eps, + } + if modality == "img2img": + extra_args["img_cfg_scale"] = args.img_cfg_scale + if is_text_output: + extra_args["max_tokens"] = args.max_tokens + extra_args["do_sample"] = args.do_sample + extra_args["temperature"] = args.temperature + + sampling_params = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + seed=args.seed, + num_inference_steps=args.num_steps, + extra_args=extra_args, + ) + + print(f"\n{'=' * 60}") + print(f"SenseNova-U1 Configuration ({modality}):") + print(f" Model : {args.model}") + if not is_text_output: + print(f" Image size : {args.width}x{args.height}") + print(f" Steps : {args.num_steps}") + print(f" CFG scale : {args.cfg_scale}") + if modality == "img2img": + print(f" Img CFG scale : {args.img_cfg_scale}") + if has_images: + print(f" Input images : {args.image}") + if is_text_output: + print(f" Max tokens : {args.max_tokens}") + print(f" Temperature : {args.temperature}") + print(f" Seed : {args.seed}") + print(f" Think mode : {args.think}") + print(f" TP size : {args.tensor_parallel_size}") + print(f"{'=' * 60}\n") + + # Build prompt dict + if has_images: + if not args.image: + raise ValueError(f"{modality} requires --image.") + input_images = [Image.open(p).convert("RGB") for p in args.image] + if is_text_output: + prompt_dict = { + "prompt": args.prompt, + "multi_modal_data": {"image": input_images}, + "modalities": ["text"], + } + else: + prompt_dict = { + "prompt": args.prompt, + "multi_modal_data": {"image": input_images}, + "modalities": ["img2img"], + } + elif is_text_output: + prompt_dict = {"prompt": args.prompt, "modalities": ["text"]} + else: + prompt_dict = {"prompt": args.prompt, "modalities": ["image"]} + + outputs = list( + omni.generate( + prompts=prompt_dict, + sampling_params_list=sampling_params, + ) + ) + + for req_output in outputs: + custom = getattr(req_output, "_custom_output", {}) or {} + if args.print_think and custom.get("think_text"): + print(f"[Think]\n{custom['think_text']}\n") + + if is_text_output: + text = custom.get("text_output", "") + if not text: + text = getattr(req_output, "text", "") or "" + print(f"[Response]\n{text}") + else: + images = getattr(req_output, "images", None) or [] + if not images: + print("[Warning] No images generated.") + continue + for j, img in enumerate(images): + save_path = os.path.join(args.output, f"sensenova_u1_output_{j}.png") + img.save(save_path) + print(f"[Output] Saved {img.size[0]}x{img.size[1]} image to {save_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/sensenova_u1/README.md b/examples/online_serving/sensenova_u1/README.md new file mode 100644 index 000000000..161e5de3b --- /dev/null +++ b/examples/online_serving/sensenova_u1/README.md @@ -0,0 +1,163 @@ +# SenseNova-U1 Online Serving + +## Launch the Server + +```bash +vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 +``` + +Or use the convenience script: + +```bash +cd examples/online_serving/sensenova_u1 +bash run_server.sh +``` + +### Tensor Parallelism (TP) + +```bash +vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 \ + --tensor-parallel-size 2 +``` + +## Send Requests + +```bash +cd examples/online_serving/sensenova_u1 +``` + +### Text to Image (text2img) + +**Python client:** + +```bash +python openai_chat_client.py \ + --prompt "A beautiful sunset over mountains" \ + --modality text2img \ + --height 2048 --width 2048 --num-steps 50 +``` + +**curl:** + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [{"role": "user", "content": [{"type": "text", "text": "A beautiful sunset over mountains"}]}], + "modalities": ["image"], + "height": 2048, + "width": 2048, + "num_inference_steps": 50, + "seed": 42 + }' +``` + +### Image to Image (img2img) + +**Python client:** + +```bash +python openai_chat_client.py \ + --prompt "Turn this into an oil painting" \ + --modality img2img \ + --image-url /path/to/input.jpg \ + --height 2048 --width 2048 +``` + +**curl:** + +```bash +IMAGE_BASE64=$(base64 -w 0 input.jpg) + +cat < payload.json +{ + "messages": [{ + "role": "user", + "content": [ + {"type": "text", "text": "Turn this into an oil painting"}, + {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}} + ] + }], + "modalities": ["image"], + "height": 2048, + "width": 2048, + "num_inference_steps": 50, + "seed": 42 +} +EOF + +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d @payload.json +``` + +### Image to Text (img2text) + +**Python client:** + +```bash +python openai_chat_client.py \ + --prompt "Describe this image in detail" \ + --modality img2text \ + --image-url /path/to/image.jpg +``` + +**curl:** + +```bash +IMAGE_BASE64=$(base64 -w 0 image.jpg) + +cat < payload.json +{ + "messages": [{ + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image in detail"}, + {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}} + ] + }], + "modalities": ["text"] +} +EOF + +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d @payload.json +``` + +### Text to Text (text2text) + +**Python client:** + +```bash +python openai_chat_client.py \ + --prompt "What is the capital of France?" \ + --modality text2text +``` + +**curl:** + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [{"role": "user", "content": [{"type": "text", "text": "What is the capital of France?"}]}], + "modalities": ["text"] + }' +``` + +## Python Client Arguments + +| Argument | Default | Description | +| :------- | :------ | :---------- | +| `--prompt` / `-p` | `A cute cat` | Text prompt | +| `--output` / `-o` | `sensenova_u1_output.png` | Output file path | +| `--server` / `-s` | `http://localhost:8091` | Server URL | +| `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) | +| `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` | +| `--height` | `2048` | Image height (image generation only) | +| `--width` | `2048` | Image width (image generation only) | +| `--num-steps` | `50` | Number of inference steps (image generation only) | +| `--seed` | `42` | Random seed | +| `--cfg-scale` | `4.0` | CFG scale (image generation only) | +| `--think` | `False` | Enable think mode | diff --git a/examples/online_serving/sensenova_u1/openai_chat_client.py b/examples/online_serving/sensenova_u1/openai_chat_client.py new file mode 100644 index 000000000..df9afa83d --- /dev/null +++ b/examples/online_serving/sensenova_u1/openai_chat_client.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python3 +""" +SenseNova-U1 OpenAI-compatible chat client for all four task modalities. + +Usage: + # Text-to-image + python openai_chat_client.py --prompt "A cute cat" --modality text2img + + # Image-to-image editing + python openai_chat_client.py --prompt "Turn this into an oil painting" \ + --modality img2img --image-url /path/to/image.jpg + + # Image understanding (img2text) + python openai_chat_client.py --prompt "Describe this image in detail" \ + --modality img2text --image-url /path/to/image.jpg + + # Text-to-text (chat) + python openai_chat_client.py --prompt "What is the capital of France?" \ + --modality text2text +""" + +import argparse +import base64 +from pathlib import Path + +import requests + + +def _encode_image(image_url: str) -> str: + """Encode a local file or URL to a base64 data URI.""" + if Path(image_url).exists(): + with open(image_url, "rb") as f: + b64_data = base64.b64encode(f.read()).decode("utf-8") + return f"data:image/jpeg;base64,{b64_data}" + return image_url + + +def generate( + prompt: str, + server_url: str = "http://localhost:8091", + image_url: str | None = None, + modality: str = "text2img", + **kwargs: object, +) -> bytes | str | None: + """Send a request to the SenseNova-U1 server. + + All keyword arguments (height, width, seed, cfg_scale, think, …) are + forwarded as top-level fields in the request payload. The serving layer + maps standard keys (height, width, seed, num_inference_steps) to sampling + params and forwards model-specific keys (think, cfg_norm, img_cfg_scale, + …) to ``extra_args`` based on the pipeline's ``EXTRA_BODY_PARAMS``. + + Returns: + Image bytes (for image outputs) or text string (for text outputs). + """ + content = [{"type": "text", "text": prompt}] + + if image_url: + content.append({"type": "image_url", "image_url": {"url": _encode_image(image_url)}}) + + messages = [{"role": "user", "content": content}] + payload: dict = {"messages": messages} + + if modality in ("text2img", "img2img"): + payload["modalities"] = ["image"] + else: + payload["modalities"] = ["text"] + + for key, val in kwargs.items(): + if val is not None and val is not False: + payload[key] = val + + try: + print(f"Sending {modality} request to {server_url}...") + response = requests.post( + f"{server_url}/v1/chat/completions", + headers={"Content-Type": "application/json"}, + json=payload, + timeout=600, + ) + response.raise_for_status() + data = response.json() + + # Print think text if present in response metrics + metrics = data.get("metrics") or {} + think_text = metrics.get("think_text") + if think_text: + print(f"\n[Think]\n{think_text}\n") + + choices = data.get("choices", []) + + for choice in choices: + choice_content = choice.get("message", {}).get("content") + + if isinstance(choice_content, list) and len(choice_content) > 0: + first_item = choice_content[0] + if isinstance(first_item, dict) and "image_url" in first_item: + img_url_str = first_item["image_url"].get("url", "") + if img_url_str.startswith("data:image"): + _, b64_data = img_url_str.split(",", 1) + return base64.b64decode(b64_data) + + for choice in choices: + choice_content = choice.get("message", {}).get("content") + if isinstance(choice_content, str) and choice_content: + return choice_content + + print(f"Unexpected response format: {choices}") + return None + + except Exception as e: + print(f"Error: {e}") + return None + + +def main(): + parser = argparse.ArgumentParser(description="SenseNova-U1 multimodal chat client") + parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt") + parser.add_argument("--output", "-o", default="sensenova_u1_output.png", help="Output file (for image results)") + parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL") + parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path") + parser.add_argument( + "--modality", + "-m", + default="text2img", + choices=["text2img", "img2img", "img2text", "text2text"], + help="Task modality", + ) + # Standard generation parameters + parser.add_argument("--height", type=int, default=2048, help="Image height") + parser.add_argument("--width", type=int, default=2048, help="Image width") + parser.add_argument("--num-steps", type=int, default=50, help="Inference steps") + parser.add_argument("--seed", type=int, default=42, help="Random seed") + # Model-specific parameters (forwarded via EXTRA_BODY_PARAMS) + parser.add_argument("--cfg-scale", type=float, default=4.0, help="CFG scale") + parser.add_argument("--img-cfg-scale", type=float, help="Image CFG scale (img2img)") + parser.add_argument("--cfg-norm", type=str, help="CFG normalization mode") + parser.add_argument("--timestep-shift", type=float, help="Timestep shift") + parser.add_argument("--think", action="store_true", help="Enable think mode") + + args = parser.parse_args() + is_image = args.modality in ("text2img", "img2img") + + print(f"Mode: {args.modality}") + if args.image_url: + print(f"Input Image: {args.image_url}") + + extra: dict[str, object] = { + "seed": args.seed, + "think": args.think or None, + } + if is_image: + extra.update( + height=args.height, + width=args.width, + num_inference_steps=args.num_steps, + cfg_scale=args.cfg_scale, + ) + if args.img_cfg_scale is not None: + extra["img_cfg_scale"] = args.img_cfg_scale + if args.cfg_norm is not None: + extra["cfg_norm"] = args.cfg_norm + if args.timestep_shift is not None: + extra["timestep_shift"] = args.timestep_shift + + result = generate( + prompt=args.prompt, + server_url=args.server, + image_url=args.image_url, + modality=args.modality, + **extra, + ) + + if result: + if isinstance(result, bytes): + output_path = Path(args.output) + output_path.write_bytes(result) + print(f"Image saved to: {output_path}") + print(f"Size: {len(result) / 1024:.1f} KB") + elif isinstance(result, str): + print(f"[Response]\n{result}") + else: + print("Failed to generate response") + exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/sensenova_u1/run_server.sh b/examples/online_serving/sensenova_u1/run_server.sh new file mode 100644 index 000000000..8edf54cc3 --- /dev/null +++ b/examples/online_serving/sensenova_u1/run_server.sh @@ -0,0 +1,12 @@ +#!/bin/bash +# SenseNova-U1 online serving startup script + +MODEL="${MODEL:-SenseNova/SenseNova-U1-8B-MoT}" +PORT="${PORT:-8091}" + +echo "Starting SenseNova-U1 server..." +echo "Model: $MODEL" +echo "Port: $PORT" + +vllm serve "$MODEL" --omni \ + --port "$PORT" diff --git a/recipes/SenseNova/SenseNova-U1.md b/recipes/SenseNova/SenseNova-U1.md new file mode 100644 index 000000000..e88da2eb3 --- /dev/null +++ b/recipes/SenseNova/SenseNova-U1.md @@ -0,0 +1,187 @@ +# SenseNova-U1 for unified image generation and understanding + +## Summary + +- Vendor: SenseNova +- Model: `SenseNova/SenseNova-U1-8B-MoT` +- Task: text2img, img2img, img2text (visual understanding), text2text (chat) +- Mode: Offline inference, Online serving (OpenAI-compatible API) +- Maintainer: Community + +## When to use this recipe + +Use this recipe to run SenseNova-U1-8B-MoT via vLLM-Omni. SenseNova-U1 is a +unified Qwen3-based LLM with Mixture-of-Tokenizers (MoT) attention that handles +text encoding, optional chain-of-thought reasoning, flow-matching image +denoising, and visual understanding in a single pipeline — no separate text +encoder or VAE needed. It supports four task modalities: text-to-image, +image-to-image editing (with dual CFG), image-to-text understanding, and +text-to-text chat. + +## References + +- Offline example: + [`examples/offline_inference/sensenova_u1/end2end.py`](../../examples/offline_inference/sensenova_u1/end2end.py) +- Online serving: + [`examples/online_serving/sensenova_u1/`](../../examples/online_serving/sensenova_u1/) +- E2E tests: + [`tests/e2e/offline_inference/test_sensenova_u1_text2img.py`](../../tests/e2e/offline_inference/test_sensenova_u1_text2img.py), + [`tests/e2e/offline_inference/test_sensenova_u1_img2img.py`](../../tests/e2e/offline_inference/test_sensenova_u1_img2img.py) +- HuggingFace model page: + [SenseNova/SenseNova-U1-8B-MoT](https://huggingface.co/SenseNova/SenseNova-U1-8B-MoT) + +## Hardware Support + +## GPU + +### 1x H200 (144GB) + +#### Environment + +- OS: Linux +- Python: 3.12 +- Driver / runtime: NVIDIA 590.48.01, CUDA 13.1 +- vLLM-Omni version: 0.18.1.dev + +#### Command + +```bash +python examples/offline_inference/sensenova_u1/end2end.py \ + --prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look. The portrait should feel polished and natural, with sharp eyes, realistic skin texture, accurate facial anatomy, and premium lighting that keeps the face as the main focus." \ + --width 1536 --height 2720 \ + --seed 42 --num-steps 50 \ + --cfg-scale 4.0 --timestep-shift 3.0 --cfg-norm none \ + --think --print-think \ + --output outputs +``` + +#### Verification + +```bash +pytest -s -v tests/e2e/offline_inference/test_sensenova_u1_text2img.py \ + -m "advanced_model" --run-level "advanced_model" +``` + +#### Notes + +- E2E latency: **32.1s** (1536×2720, 50 steps, think mode, CFG scale 4.0) +- Peak VRAM: **35.9 GB** reserved, 35.1 GB allocated +- Model loading: 32.8 GiB, 8.7s +- No deploy YAML needed — the engine auto-generates a single-stage diffusion config. +- Think mode (`--think`) is recommended for higher image quality. + +#### Image-to-Image Editing (img2img) + +```bash +python examples/offline_inference/sensenova_u1/end2end.py \ + --prompt "Turn this into an oil painting" \ + --image input.png \ + --width 2048 --height 2048 \ + --seed 42 --num-steps 50 \ + --cfg-scale 4.0 --img-cfg-scale 1.0 --cfg-norm none \ + --think --print-think \ + --output outputs +``` + +- img2img uses dual CFG: `--cfg-scale` controls text guidance, `--img-cfg-scale` + controls image guidance (1.0 = image CFG disabled). +- Pass multiple `--image` paths for multi-reference editing. + +#### Image Understanding (img2text) + +```bash +python examples/offline_inference/sensenova_u1/end2end.py \ + --modality img2text \ + --prompt "Describe this image in detail" \ + --image photo.jpg \ + --max-tokens 512 +``` + +#### Text-to-Text Chat (text2text) + +```bash +python examples/offline_inference/sensenova_u1/end2end.py \ + --modality text2text \ + --prompt "Explain the theory of relativity in simple terms" \ + --max-tokens 256 +``` + +- For img2text and text2text, image generation parameters (height, width, + num-steps, cfg-scale) are ignored. +- Use `--do-sample --temperature 0.7` for more diverse text responses. + +### 2x H200 (144GB) — TP=2 + +#### Environment + +- OS: Linux +- Python: 3.12 +- Driver / runtime: NVIDIA 590.48.01, CUDA 13.1 +- vLLM-Omni version: 0.18.1.dev + +#### Command + +```bash +python examples/offline_inference/sensenova_u1/end2end.py \ + --prompt "Close portrait of an elderly woman by a farmhouse window, textured skin, gentle smile, warm natural light, emotional documentary look. The portrait should feel polished and natural, with sharp eyes, realistic skin texture, accurate facial anatomy, and premium lighting that keeps the face as the main focus." \ + --width 1536 --height 2720 \ + --seed 42 --num-steps 50 \ + --cfg-scale 4.0 --timestep-shift 3.0 --cfg-norm none \ + --think --print-think \ + --tensor-parallel-size 2 \ + --output outputs +``` + +#### Verification + +Verify the output image is generated at `outputs/sensenova_u1_output_0.png` +with the expected 1536×2720 resolution. + +#### Notes + +- E2E latency: **28.3s** (1536×2720, 50 steps, think mode, CFG scale 4.0) +- Peak VRAM (per GPU): **18.2 GB** reserved, 17.9 GB allocated +- Model loading: 16.5 GiB per GPU, 7.0s +- TP=2 provides ~12% speedup over TP=1; limited by serial CFG dual-forward + and communication overhead. +- The LLM transformer uses `QKVParallelLinear` and `MergedColumnParallelLinear` + for fused QKV and gate/up projections with TP support. + +## Online Serving + +SenseNova-U1 supports all four modalities via the OpenAI-compatible +`/v1/chat/completions` API. + +### Launch + +```bash +vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 +``` + +### Send Requests + +```bash +cd examples/online_serving/sensenova_u1 + +# Text-to-image +python openai_chat_client.py \ + --prompt "A beautiful sunset" --modality text2img + +# Image-to-image editing +python openai_chat_client.py \ + --prompt "Turn this into an oil painting" \ + --modality img2img --image-url input.jpg + +# Image understanding +python openai_chat_client.py \ + --prompt "Describe this image" \ + --modality img2text --image-url photo.jpg + +# Text chat +python openai_chat_client.py \ + --prompt "What is the capital of France?" \ + --modality text2text +``` + +For full API documentation and curl examples, see +[`examples/online_serving/sensenova_u1/README.md`](../../examples/online_serving/sensenova_u1/README.md). diff --git a/tests/e2e/offline_inference/test_sensenova_u1_img2img.py b/tests/e2e/offline_inference/test_sensenova_u1_img2img.py new file mode 100644 index 000000000..10f5f05dc --- /dev/null +++ b/tests/e2e/offline_inference/test_sensenova_u1_img2img.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end test for SenseNova-U1 img2img generation. + +This test validates that the SenseNova-U1 model generates edited images from +an input image and text prompt that match expected reference pixel values +within a ±10 tolerance. + +Equivalent to running: + python examples/offline_inference/sensenova_u1/end2end.py \ + --prompt "Turn this into an oil painting" \ + --image 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg \ + --width 2560 --height 1669 \ + --seed 42 --num-steps 50 \ + --cfg-scale 4.0 --img-cfg-scale 1.0 --cfg-norm none \ + --think --print-think +""" + +import os +from typing import Any + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest +from PIL import Image +from vllm.assets.image import ImageAsset + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +# Reference pixel data extracted from the known-good output image generated by: +# python examples/offline_inference/sensenova_u1/end2end.py \ +# --prompt "Turn this into an oil painting" \ +# --image 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg \ +# --width 2560 --height 1669 --seed 42 --num-steps 50 \ +# --cfg-scale 4.0 --img-cfg-scale 1.0 --cfg-norm none --think +REFERENCE_PIXELS = [ + {"position": (100, 100), "rgb": (95, 131, 174)}, + {"position": (400, 50), "rgb": (142, 153, 174)}, + {"position": (700, 100), "rgb": (127, 149, 174)}, + {"position": (150, 400), "rgb": (99, 141, 183)}, + {"position": (512, 336), "rgb": (184, 177, 165)}, + {"position": (700, 400), "rgb": (106, 139, 168)}, + {"position": (100, 600), "rgb": (170, 172, 155)}, + {"position": (400, 600), "rgb": (159, 171, 169)}, + {"position": (700, 600), "rgb": (139, 172, 187)}, + {"position": (256, 256), "rgb": (136, 158, 180)}, +] + +PIXEL_TOLERANCE = 10 + +DEFAULT_PROMPT = "Turn this into an oil painting" + +EXPECTED_OUTPUT_SIZE = (2560, 1664) + + +def _load_input_image() -> Image.Image: + """Load the test input image via vllm's ImageAsset.""" + return ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") + + +def _build_sampling_params() -> OmniDiffusionSamplingParams: + """Build sampling parameters for SenseNova-U1 img2img generation.""" + return OmniDiffusionSamplingParams( + height=1669, + width=2560, + seed=42, + num_inference_steps=50, + extra_args={ + "cfg_scale": 4.0, + "img_cfg_scale": 1.0, + "cfg_norm": "none", + "timestep_shift": 3.0, + "cfg_interval": (0.0, 1.0), + "batch_size": 1, + "think": True, + "t_eps": 0.02, + }, + ) + + +def _extract_generated_image(omni_outputs: list) -> Image.Image | None: + """Extract the generated image from Omni outputs (single-stage DiT).""" + for req_output in omni_outputs: + if images := getattr(req_output, "images", None): + return images[0] + return None + + +def _validate_pixels( + image: Image.Image, + reference_pixels: list[dict[str, Any]] = REFERENCE_PIXELS, + tolerance: int = PIXEL_TOLERANCE, +) -> None: + """Validate that image pixels match expected reference values. + + Args: + image: The PIL Image to validate. + reference_pixels: List of dicts with 'position' (x, y) and 'rgb' (R, G, B). + tolerance: Maximum allowed difference per color channel. + + Raises: + AssertionError: If any pixel differs beyond tolerance. + """ + for ref in reference_pixels: + x, y = ref["position"] + expected = ref["rgb"] + actual = image.getpixel((x, y))[:3] + assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), ( + f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}" + ) + + +def _generate_sensenova_u1_img2img( + omni: Omni, + input_image: Image.Image, + prompt: str = DEFAULT_PROMPT, +) -> Image.Image: + """Generate an edited image using SenseNova-U1 model with img2img pipeline. + + Args: + omni: The Omni instance to use for generation. + input_image: The input PIL Image for img2img. + prompt: The text prompt for image editing. + + Returns: + The generated PIL Image. + + Raises: + AssertionError: If no image is generated or size is incorrect. + """ + sampling_params = _build_sampling_params() + + omni_outputs = list( + omni.generate( + prompts={ + "prompt": prompt, + "multi_modal_data": {"image": input_image}, + "modalities": ["img2img"], + }, + sampling_params_list=sampling_params, + ) + ) + + generated_image = _extract_generated_image(omni_outputs) + assert generated_image is not None, "No images generated" + assert generated_image.size == EXPECTED_OUTPUT_SIZE, f"Expected {EXPECTED_OUTPUT_SIZE}, got {generated_image.size}" + + return generated_image + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +def test_sensenova_u1_img2img(run_level): + """Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML).""" + input_image = _load_input_image() + with OmniRunner( + "SenseNova/SenseNova-U1-8B-MoT", + stage_configs_path=None, + ) as runner: + generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image) + if run_level == "advanced_model": + _validate_pixels(generated_image) diff --git a/tests/e2e/offline_inference/test_sensenova_u1_text2img.py b/tests/e2e/offline_inference/test_sensenova_u1_text2img.py new file mode 100644 index 000000000..334381067 --- /dev/null +++ b/tests/e2e/offline_inference/test_sensenova_u1_text2img.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end test for SenseNova-U1 text2img generation. + +This test validates that the SenseNova-U1 model generates images that match +expected reference pixel values within a ±10 tolerance. + +Equivalent to running: + python SenseNova-U1/examples/t2i/inference.py \ + --model_path SenseNova/SenseNova-U1-8B-MoT \ + --prompt "Close portrait of an elderly woman ..." \ + --width 1536 --height 2720 \ + --cfg_scale 4.0 --cfg_norm none --timestep_shift 3.0 \ + --num_steps 50 --seed 42 --think +""" + +import os +from typing import Any + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest +from PIL import Image + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +# Reference pixel data extracted from the known-good output image generated by: +# python examples/offline_inference/sensenova_u1/end2end.py \ +# --prompt "Close portrait of an elderly woman ..." \ +# --width 1536 --height 2720 --seed 42 --num-steps 50 \ +# --cfg-scale 4.0 --timestep-shift 3.0 --cfg-norm none --think +REFERENCE_PIXELS = [ + {"position": (100, 100), "rgb": (247, 249, 250)}, + {"position": (768, 200), "rgb": (176, 135, 97)}, + {"position": (400, 600), "rgb": (193, 188, 180)}, + {"position": (1200, 2000), "rgb": (117, 98, 84)}, + {"position": (750, 500), "rgb": (186, 135, 84)}, + {"position": (300, 1360), "rgb": (202, 157, 107)}, + {"position": (1000, 1800), "rgb": (63, 32, 15)}, + {"position": (500, 2400), "rgb": (135, 116, 104)}, + {"position": (768, 1360), "rgb": (70, 34, 16)}, + {"position": (200, 900), "rgb": (208, 201, 191)}, +] + +PIXEL_TOLERANCE = 10 + +DEFAULT_PROMPT = ( + "Close portrait of an elderly woman by a farmhouse window, textured skin, " + "gentle smile, warm natural light, emotional documentary look. The portrait " + "should feel polished and natural, with sharp eyes, realistic skin texture, " + "accurate facial anatomy, and premium lighting that keeps the face as the " + "main focus." +) + +EXPECTED_OUTPUT_SIZE = (1536, 2720) + + +def _build_sampling_params() -> OmniDiffusionSamplingParams: + """Build sampling parameters for SenseNova-U1 text2img generation.""" + return OmniDiffusionSamplingParams( + height=EXPECTED_OUTPUT_SIZE[1], + width=EXPECTED_OUTPUT_SIZE[0], + seed=42, + num_inference_steps=50, + extra_args={ + "cfg_scale": 4.0, + "cfg_norm": "none", + "timestep_shift": 3.0, + "cfg_interval": (0.0, 1.0), + "batch_size": 1, + "think": True, + "t_eps": 0.02, + }, + ) + + +def _extract_generated_image(omni_outputs: list) -> Image.Image | None: + """Extract the generated image from Omni outputs (single-stage DiT).""" + for req_output in omni_outputs: + if images := getattr(req_output, "images", None): + return images[0] + return None + + +def _validate_pixels( + image: Image.Image, + reference_pixels: list[dict[str, Any]] = REFERENCE_PIXELS, + tolerance: int = PIXEL_TOLERANCE, +) -> None: + """Validate that image pixels match expected reference values. + + Args: + image: The PIL Image to validate. + reference_pixels: List of dicts with 'position' (x, y) and 'rgb' (R, G, B). + tolerance: Maximum allowed difference per color channel. + + Raises: + AssertionError: If any pixel differs beyond tolerance. + """ + for ref in reference_pixels: + x, y = ref["position"] + expected = ref["rgb"] + actual = image.getpixel((x, y))[:3] + assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), ( + f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}" + ) + + +def _generate_sensenova_u1_image( + omni: Omni, + prompt: str = DEFAULT_PROMPT, +) -> Image.Image: + """Generate an image using SenseNova-U1 model with configured parameters. + + Args: + omni: The Omni instance to use for generation. + prompt: The text prompt for image generation. + + Returns: + The generated PIL Image. + + Raises: + AssertionError: If no image is generated or size is incorrect. + """ + sampling_params = _build_sampling_params() + + omni_outputs = list( + omni.generate( + prompts={"prompt": prompt, "modalities": ["image"]}, + sampling_params_list=sampling_params, + ) + ) + + generated_image = _extract_generated_image(omni_outputs) + assert generated_image is not None, "No images generated" + assert generated_image.size == EXPECTED_OUTPUT_SIZE, f"Expected {EXPECTED_OUTPUT_SIZE}, got {generated_image.size}" + + return generated_image + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}) +def test_sensenova_u1_text2img(run_level): + """Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML).""" + with OmniRunner( + "SenseNova/SenseNova-U1-8B-MoT", + stage_configs_path=None, + ) as runner: + generated_image = _generate_sensenova_u1_image(runner.omni) + if run_level == "advanced_model": + _validate_pixels(generated_image) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 5a87e7762..bea0b7295 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -766,6 +766,10 @@ class OmniDiffusionConfig: self.model_class_name = "BagelPipeline" self.set_tf_model_config(TransformerConfig()) self.update_multimodal_support() + elif model_type == "neo_chat": + self.model_class_name = "SenseNovaU1Pipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() elif model_type == "nextstep": if self.model_class_name is None: self.model_class_name = "NextStep11Pipeline" diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 1e421316a..612e1fa09 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -69,6 +69,34 @@ def supports_audio_output(model_class_name: str) -> bool: return bool(getattr(model_cls, "support_audio_output", False)) +def get_extra_body_params(model_class_name: str) -> frozenset[str]: + """Return the set of extra_body keys accepted by a pipeline. + + Each pipeline can declare ``EXTRA_BODY_PARAMS: ClassVar[frozenset[str]]`` + to advertise which request-level parameters should be forwarded from + ``extra_body`` to ``OmniDiffusionSamplingParams.extra_args``. + Returns an empty frozenset when the pipeline does not declare any. + """ + model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) + if model_cls is None: + return frozenset() + return frozenset(getattr(model_cls, "EXTRA_BODY_PARAMS", frozenset())) + + +def get_extra_output_params(model_class_name: str) -> frozenset[str]: + """Return the set of custom_output keys to expose in API response metrics. + + Each pipeline can declare ``EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]]`` + to advertise which ``DiffusionOutput.custom_output`` keys should be + copied into the response ``metrics`` dict. + Returns an empty frozenset when the pipeline does not declare any. + """ + model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) + if model_cls is None: + return frozenset() + return frozenset(getattr(model_cls, "EXTRA_OUTPUT_PARAMS", frozenset())) + + class DiffusionEngine: """The diffusion engine for vLLM-Omni diffusion models.""" @@ -230,6 +258,11 @@ class DiffusionEngine: "postprocess_time_ms": postprocess_time * 1000, } + # Detect text output: when the pipeline returns a string (e.g., + # SenseNova-U1 / BAGEL single-stage img2text / text2text), wrap it + # as a text-type response instead of an image. + is_text_output = isinstance(output_data, str) and custom_output.get("text_output") is not None + # Handle single request or multiple requests is_audio_output = supports_audio_output(self.od_config.model_class_name) if is_audio_output and model_audio_sample_rate is None: @@ -247,6 +280,19 @@ class DiffusionEngine: prompt = request.prompts[0] request_id = request.request_ids[0] if request.request_ids else "" + if is_text_output: + return [ + OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + prompt=prompt, + metrics=metrics, + custom_output=custom_output, + final_output_type="text", + stage_durations=output.stage_durations, + peak_memory_mb=output.peak_memory_mb, + ), + ] if is_audio_output: request_audio_payload = outputs[0] if len(outputs) == 1 else outputs return [ diff --git a/vllm_omni/diffusion/models/sensenova_u1/__init__.py b/vllm_omni/diffusion/models/sensenova_u1/__init__.py new file mode 100644 index 000000000..93b148197 --- /dev/null +++ b/vllm_omni/diffusion/models/sensenova_u1/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from .pipeline_sensenova_u1 import SenseNovaU1Pipeline, get_sensenova_u1_post_process_func + +__all__ = [ + "SenseNovaU1Pipeline", + "get_sensenova_u1_post_process_func", +] diff --git a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py new file mode 100644 index 000000000..c3ca053ca --- /dev/null +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -0,0 +1,1475 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SenseNova-U1 Pipeline for vLLM-Omni. + +SenseNova-U1 is a unified Qwen3-based model that uses Mixture-of-Tokenizers +(MoT) attention for text-to-image generation via flow matching in patch space. +It has no separate VAE or text encoder — the Qwen3 LLM itself serves as both +the text encoder (via KV cache) and the denoising backbone (via MoT branches). + +Key integration points: +- Transformer layers ported with TP support (QKVParallelLinear, + MergedColumnParallelLinear, RowParallelLinear) in sensenova_u1_transformer.py. +- Vision model (NEOVisionModel) and FM modules kept as standard nn.Module + since they are lightweight (no transformer blocks). +- Weight loading uses stacked_params_mapping for fused QKV and gate_up. +""" + +from __future__ import annotations + +import json +import math +import os +from collections.abc import Iterable +from types import SimpleNamespace +from typing import ClassVar + +import numpy as np +import torch +import torch.nn as nn +import torchvision.transforms as T +from PIL import Image +from transformers import AutoTokenizer +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest + +from .sensenova_u1_transformer import ( + SenseNovaU1ForCausalLM, + clear_flash_kv_cache, + create_block_causal_mask, + prepare_flash_kv_cache, +) + +logger = init_logger(__name__) + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +NORM_MEAN = (0.5, 0.5, 0.5) +NORM_STD = (0.5, 0.5, 0.5) + +IMAGENET_MEAN = (0.485, 0.456, 0.406) +IMAGENET_STD = (0.229, 0.224, 0.225) + +IMG_START_TOKEN = "" +IMG_END_TOKEN = "" +IMG_CONTEXT_TOKEN = "" + +SYSTEM_MESSAGE_FOR_GEN = ( + "You are an image generation and editing assistant that accurately understands and executes " + "user intent.\n\nYou support two modes:\n\n1. Think Mode:\nIf the task requires reasoning, you " + "MUST start with a block. Put all reasoning inside the block using plain text. " + "DO NOT include any image tags. Keep it reasonable and directly useful for producing the final " + "image.\n\n2. Non-Think Mode:\nIf no reasoning is needed, directly produce the final image.\n\n" + "Task Types:\n\nA. Text-to-Image Generation:\n" + "- Generate a high-quality image based on the user's description.\n" + "- Ensure visual clarity, semantic consistency, and completeness.\n" + "- DO NOT introduce elements that contradict or override the user's intent.\n\n" + "B. Image Editing:\n" + "- Use the provided image(s) as input or reference for modification or transformation.\n" + "- The result can be an edited image or a new image based on the reference(s).\n" + "- Preserve all unspecified attributes unless explicitly changed.\n\n" + "General Rules:\n" + "- For any visible text in the image, follow the language specified for the rendered text in " + "the user's description, not the language of the prompt. If no language is specified, use the " + "user's input language." +) + + +# --------------------------------------------------------------------------- +# Config parsed from config.json +# --------------------------------------------------------------------------- + + +def _resolve_model_path(model_path: str) -> str: + """Resolve a HuggingFace model ID or local path to a local directory.""" + if os.path.isdir(model_path): + return model_path + from huggingface_hub import snapshot_download + + return snapshot_download(model_path) + + +def _round_by_factor(number: float, factor: int) -> int: + return round(number / factor) * factor + + +def _ceil_by_factor(number: float, factor: int) -> int: + return math.ceil(number / factor) * factor + + +def _floor_by_factor(number: float, factor: int) -> int: + return math.floor(number / factor) * factor + + +def _smart_resize( + height: int, + width: int, + factor: int = 32, + min_pixels: int = 65536, + max_pixels: int = 4194304, +) -> tuple[int, int]: + """Rescale so that H/W are divisible by *factor* and total pixels in [min, max].""" + if max(height, width) / min(height, width) > 200: + raise ValueError(f"absolute aspect ratio must be < 200, got {max(height, width) / min(height, width)}") + h_bar = max(factor, _round_by_factor(height, factor)) + w_bar = max(factor, _round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = max(factor, _floor_by_factor(height / beta, factor)) + w_bar = max(factor, _floor_by_factor(width / beta, factor)) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = _ceil_by_factor(height * beta, factor) + w_bar = _ceil_by_factor(width * beta, factor) + return h_bar, w_bar + + +def _load_image_native( + image, + patch_size: int = 16, + downsample_ratio: float = 0.5, + min_pixels: int = 65536, + max_pixels: int = 4194304, +): + """Load, smart-resize, ImageNet-normalize, and patchify an input image.""" + if not isinstance(image, Image.Image): + image = Image.open(image) + if image.mode == "RGBA": + bg = Image.new("RGB", image.size, (255, 255, 255)) + bg.paste(image, mask=image.split()[3]) + image = bg.convert("RGB") + else: + image = image.convert("RGB") + + transform = T.Compose( + [ + T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img), + T.ToTensor(), + T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD), + ] + ) + + size_factor = int(patch_size // downsample_ratio) + resized_h, resized_w = _smart_resize( + image.height, + image.width, + factor=size_factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + image = image.resize((resized_w, resized_h)) + + pixel_values = transform(image).to(torch.float32) + c, h, w = pixel_values.shape + grid_h = h // patch_size + grid_w = w // patch_size + flatten_pv = ( + pixel_values.view(c, grid_h, patch_size, grid_w, patch_size) + .permute(1, 3, 0, 2, 4) + .reshape(grid_h * grid_w, c * patch_size**2) + ) + grid_hw = torch.tensor([[grid_h, grid_w]]) + return flatten_pv, grid_hw + + +def _load_sensenova_u1_config(model_path: str): + """Parse config.json and return (top_config, llm_config, vision_config) namespaces.""" + local_path = _resolve_model_path(model_path) + cfg_path = os.path.join(local_path, "config.json") + with open(cfg_path, encoding="utf-8") as f: + raw = json.load(f) + + llm_raw = raw.get("llm_config", {}) + vis_raw = raw.get("vision_config", {}) + + # LLM config needs these extra fields for 3D RoPE + llm_cfg = SimpleNamespace( + hidden_size=llm_raw.get("hidden_size", 4096), + intermediate_size=llm_raw.get("intermediate_size", 11008), + num_hidden_layers=llm_raw.get("num_hidden_layers", 32), + num_attention_heads=llm_raw.get("num_attention_heads", 32), + num_key_value_heads=llm_raw.get("num_key_value_heads", llm_raw.get("num_attention_heads", 32)), + head_dim=llm_raw.get("head_dim", llm_raw.get("hidden_size", 4096) // llm_raw.get("num_attention_heads", 32)), + hidden_act=llm_raw.get("hidden_act", "silu"), + rms_norm_eps=llm_raw.get("rms_norm_eps", 1e-6), + vocab_size=llm_raw.get("vocab_size", 152064), + rope_theta=llm_raw.get("rope_theta", 1000000.0), + rope_theta_hw=llm_raw.get("rope_theta_hw", 10000.0), + max_position_embeddings=llm_raw.get("max_position_embeddings", 40960), + max_position_embeddings_hw=llm_raw.get("max_position_embeddings_hw", 10000), + attention_bias=llm_raw.get("attention_bias", True), + layer_types=llm_raw.get("layer_types", ["full_attention"] * llm_raw.get("num_hidden_layers", 32)), + tie_word_embeddings=llm_raw.get("tie_word_embeddings", False), + ) + + vis_cfg = SimpleNamespace( + num_channels=vis_raw.get("num_channels", 3), + patch_size=vis_raw.get("patch_size", 16), + hidden_size=vis_raw.get("hidden_size", 1024), + llm_hidden_size=vis_raw.get("llm_hidden_size", [llm_cfg.hidden_size]), + downsample_ratio=vis_raw.get("downsample_ratio", [0.5]), + rope_theta_vision=vis_raw.get("rope_theta_vision", 10000.0), + max_position_embeddings_vision=vis_raw.get("max_position_embeddings_vision", 10000), + output_hidden_states=vis_raw.get("output_hidden_states", False), + use_return_dict=vis_raw.get("use_return_dict", True), + ) + + top_cfg = SimpleNamespace( + downsample_ratio=raw.get("downsample_ratio", 0.5), + template=raw.get("template", "neo1_0"), + fm_head_layers=raw.get("fm_head_layers", 2), + fm_head_dim=raw.get("fm_head_dim", 4096), + fm_head_mlp_ratio=raw.get("fm_head_mlp_ratio", 1.0), + use_pixel_head=raw.get("use_pixel_head", False), + noise_scale=raw.get("noise_scale", 1.0), + noise_scale_mode=raw.get("noise_scale_mode", "none"), + noise_scale_base_image_seq_len=raw.get("noise_scale_base_image_seq_len", 256), + noise_scale_max_value=raw.get("noise_scale_max_value", 10.0), + add_noise_scale_embedding=raw.get("add_noise_scale_embedding", False), + time_schedule=raw.get("time_schedule", "standard"), + time_shift_type=raw.get("time_shift_type", "exponential"), + base_shift=raw.get("base_shift", 0.5), + max_shift=raw.get("max_shift", 1.15), + base_image_seq_len=raw.get("base_image_seq_len", 256), + max_image_seq_len=raw.get("max_image_seq_len", 4096), + concat_time_token_num=raw.get("concat_time_token_num", 0), + t_eps=raw.get("t_eps", 0.02), + ) + + return top_cfg, llm_cfg, vis_cfg + + +# --------------------------------------------------------------------------- +# Vision Embedding — 2D RoPE + Conv patch embed (no transformer blocks) +# --------------------------------------------------------------------------- + + +def _precompute_rope_freqs_sincos(dim, max_position, base=10000.0, device=None): + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim)) + t = torch.arange(max_position, device=device).float() + freqs = torch.outer(t, inv_freq) + return torch.cos(freqs), torch.sin(freqs) + + +def _apply_rotary_emb_1d(x, cos_cached, sin_cached, positions): + cos = cos_cached[positions] + sin = sin_cached[positions] + x1, x2 = x[..., 0::2], x[..., 1::2] + out = torch.empty_like(x) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out + + +def _build_abs_positions_from_grid_hw(grid_hw, device=None): + device = grid_hw.device + B = grid_hw.shape[0] + H, W = grid_hw[:, 0], grid_hw[:, 1] + N = H * W + N_total = N.sum() + patch_to_sample = torch.repeat_interleave(torch.arange(B, device=device), N) + pid = torch.arange(N_total, device=device) + pid = pid - torch.cumsum(torch.cat([torch.tensor([0], device=device), N[:-1]]), dim=0)[patch_to_sample] + W_per_patch = W[patch_to_sample] + return pid % W_per_patch, pid // W_per_patch + + +class NEOVisionEmbeddings(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.embed_dim = config.hidden_size + llm_hidden = config.llm_hidden_size + if isinstance(llm_hidden, (list, tuple)): + llm_hidden = llm_hidden[0] + self.llm_embed_dim = llm_hidden + ds_ratio = config.downsample_ratio + if isinstance(ds_ratio, (list, tuple)): + ds_ratio = ds_ratio[0] + self.downsample_factor = int(1 / ds_ratio) + self.patch_size = config.patch_size + + self.patch_embedding = nn.Conv2d( + config.num_channels, self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size + ) + self.dense_embedding = nn.Conv2d( + self.embed_dim, self.llm_embed_dim, kernel_size=self.downsample_factor, stride=self.downsample_factor + ) + self.gelu = nn.GELU() + + rope_dim = self.embed_dim // 2 + cos_x, sin_x = _precompute_rope_freqs_sincos( + rope_dim, config.max_position_embeddings_vision, base=config.rope_theta_vision + ) + cos_y, sin_y = _precompute_rope_freqs_sincos( + rope_dim, config.max_position_embeddings_vision, base=config.rope_theta_vision + ) + self.register_buffer("cos_cached_x", cos_x, persistent=False) + self.register_buffer("sin_cached_x", sin_x, persistent=False) + self.register_buffer("cos_cached_y", cos_y, persistent=False) + self.register_buffer("sin_cached_y", sin_y, persistent=False) + + def forward(self, pixel_values, grid_hw=None): + pixel_values = pixel_values.view(-1, 3, self.patch_size, self.patch_size) + patch_embeds = self.gelu(self.patch_embedding(pixel_values)).view(-1, self.embed_dim) + + # 2D RoPE + abs_x, abs_y = _build_abs_positions_from_grid_hw(grid_hw) + dim_half = self.embed_dim // 2 + p1 = _apply_rotary_emb_1d( + patch_embeds[:, :dim_half].float(), + self.cos_cached_x.to(patch_embeds.device), + self.sin_cached_x.to(patch_embeds.device), + abs_x, + ) + p2 = _apply_rotary_emb_1d( + patch_embeds[:, dim_half:].float(), + self.cos_cached_y.to(patch_embeds.device), + self.sin_cached_y.to(patch_embeds.device), + abs_y, + ) + patch_embeds = torch.cat([p1, p2], dim=-1).to(self.patch_embedding.weight.dtype) + + patches_list = [] + cur = 0 + for i in range(grid_hw.shape[0]): + h, w = grid_hw[i] + pe = patch_embeds[cur : cur + h * w].view(h, w, -1).unsqueeze(0) + pe = self.dense_embedding(pe.permute(0, 3, 1, 2)).permute(0, 2, 3, 1) + patches_list.append(pe.view(-1, pe.shape[-1])) + cur += h * w + + return torch.cat(patches_list, dim=0) + + +class NEOVisionModel(nn.Module): + def __init__(self, config): + super().__init__() + self.embeddings = NEOVisionEmbeddings(config) + + def forward(self, pixel_values=None, grid_hw=None, **_kwargs): + return self.embeddings(pixel_values, grid_hw=grid_hw) + + +# --------------------------------------------------------------------------- +# FM modules +# --------------------------------------------------------------------------- + + +class TimestepEmbedder(nn.Module): + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + + @staticmethod + def timestep_embedding(t, dim, max_period=10000.0): + half = dim // 2 + freqs = torch.exp(-math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half).to(t.device) + args = t[:, None].float() * freqs[None] + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1) + return emb + + def forward(self, t): + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + return self.mlp(t_freq.to(self.mlp[0].weight.dtype)) + + +class ConvDecoder(nn.Module): + def __init__(self, input_dim=4096, hidden_dim=1024): + super().__init__() + self.ps1 = nn.PixelShuffle(2) + self.conv1 = nn.Conv2d(input_dim // 4, hidden_dim, kernel_size=3, padding=1) + self.act1 = nn.GELU() + self.ps2 = nn.PixelShuffle(2) + self.conv2 = nn.Conv2d(hidden_dim // 4, 192, kernel_size=3, padding=1) + self.ps3 = nn.PixelShuffle(8) + + def forward(self, x): + x = self.act1(self.conv1(self.ps1(x))) + x = self.ps3(self.conv2(self.ps2(x))) + return x + + +# --------------------------------------------------------------------------- +# Conversation template (inlined from fastchat MPT style) +# --------------------------------------------------------------------------- + + +def _build_t2i_query(prompt_text, system_message=None, append_text=None): + sys_msg = system_message or "" + system_prompt = f"<|im_start|>system\n{sys_msg}<|im_end|>\n" if sys_msg else "" + query = system_prompt + f"<|im_start|>user\n{prompt_text}<|im_end|>\n<|im_start|>assistant\n" + if append_text is not None: + query += append_text + return query + + +# --------------------------------------------------------------------------- +# Patchify / Unpatchify +# --------------------------------------------------------------------------- + + +def _patchify(images, patch_size, channel_first=False): + h, w = images.shape[2] // patch_size, images.shape[3] // patch_size + x = images.reshape(images.shape[0], 3, h, patch_size, w, patch_size) + if channel_first: + x = torch.einsum("nchpwq->nhwcpq", x) + else: + x = torch.einsum("nchpwq->nhwpqc", x) + return x.reshape(images.shape[0], h * w, patch_size**2 * 3) + + +def _unpatchify(x, patch_size, h=None, w=None): + if h is None or w is None: + h = w = int(x.shape[1] ** 0.5) + else: + h = h // patch_size + w = w // patch_size + x = x.reshape(x.shape[0], h, w, patch_size, patch_size, 3) + x = torch.einsum("nhwpqc->nchpwq", x) + return x.reshape(x.shape[0], 3, h * patch_size, w * patch_size) + + +# --------------------------------------------------------------------------- +# Image denormalization for output +# --------------------------------------------------------------------------- + + +def _denorm(x): + mean = torch.tensor(NORM_MEAN, device=x.device, dtype=x.dtype).view(1, 3, 1, 1) + std = torch.tensor(NORM_STD, device=x.device, dtype=x.dtype).view(1, 3, 1, 1) + return (x * std + mean).clamp(0, 1) + + +def _to_pil(batch): + arr = _denorm(batch.float()).permute(0, 2, 3, 1).cpu().numpy() + arr = (arr * 255.0).round().astype(np.uint8) + return [Image.fromarray(a) for a in arr] + + +def get_sensenova_u1_post_process_func(od_config: OmniDiffusionConfig): + def post_process_func(x): + return x + + return post_process_func + + +# --------------------------------------------------------------------------- +# CFG helpers +# --------------------------------------------------------------------------- + + +def _optimized_scale(positive_flat, negative_flat): + # Manual fp32 cast keeps numerics stable without locking us to a CUDA-only autocast region. + pos = positive_flat.float() + neg = negative_flat.float() + dot = torch.sum(pos * neg, dim=1, keepdim=True) + sq_norm = torch.sum(neg**2, dim=1, keepdim=True) + 1e-8 + return dot / sq_norm + + +# --------------------------------------------------------------------------- +# Pipeline +# --------------------------------------------------------------------------- + + +class SenseNovaU1Pipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerMixin): + """SenseNova-U1 text-to-image and image-to-image pipeline for vllm-omni. + + Builds the full model graph internally: + - language_model: SenseNovaU1ForCausalLM (TP-aware) + - vision_model: NEOVisionModel (understanding branch) + - fm_modules: ModuleDict with vision_model_mot_gen, timestep_embedder, fm_head, etc. + + img2img (image editing) is triggered when `multi_modal_data["image"]` is + present in the prompt dict. The pipeline then uses triple KV caches + (condition / img_condition / uncondition) with dual CFG + (``cfg_scale`` + ``img_cfg_scale``). + """ + + support_image_input = True + + # Model-specific parameters accepted via ``extra_body`` in the online + # serving API. The serving layer reads this set and forwards matching + # keys from the request payload to ``OmniDiffusionSamplingParams.extra_args``. + EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset( + { + "think", + "cfg_scale", + "cfg_norm", + "timestep_shift", + "t_eps", + "img_cfg_scale", + "max_tokens", + } + ) + + # Keys from ``DiffusionOutput.custom_output`` that should be surfaced in + # the API response ``metrics`` dict. The serving layer reads this set and + # copies matching entries from ``custom_output`` into the response. + EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset( + { + "think_text", + } + ) + + # CPU-offload protocol: language_model carries the denoising blocks; the + # vision and FM modules are lightweight encoders pinned on GPU during the + # diffusion loop. There is no separate VAE. + _dit_modules: ClassVar[list[str]] = ["language_model"] + _encoder_modules: ClassVar[list[str]] = ["vision_model"] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = ["fm_modules"] + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.device = get_local_device() + model_path = od_config.model + self.local_model_path = _resolve_model_path(model_path) + + self.top_cfg, self.llm_cfg, self.vis_cfg = _load_sensenova_u1_config(model_path) + + self.tokenizer = AutoTokenizer.from_pretrained(self.local_model_path) + self.img_context_token_id = self.tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN) + self.img_start_token_id = self.tokenizer.convert_tokens_to_ids(IMG_START_TOKEN) + + # Language model (TP-aware) + self.language_model = SenseNovaU1ForCausalLM( + self.llm_cfg, + prefix="language_model", + ) + + # Vision model (understanding branch) + self.vision_model = NEOVisionModel(self.vis_cfg) + + # FM modules + llm_hidden_size = self.llm_cfg.hidden_size + patch_size = self.vis_cfg.patch_size + merge_size = int(1 / self.top_cfg.downsample_ratio) + output_dim = 3 * (patch_size * merge_size) ** 2 + + vision_model_mot_gen = NEOVisionModel(self.vis_cfg) + timestep_embedder = TimestepEmbedder(llm_hidden_size) + + if self.top_cfg.use_pixel_head: + fm_head = ConvDecoder(llm_hidden_size) + else: + # The reference implementation uses a fixed 2-layer GELU MLP with a 4096-d + # intermediate, regardless of `fm_head_layers` / `fm_head_dim` in config.json + # (those fields are unused in the released checkpoint). The deeper variant + # has not been ported yet. + fm_intermediate_dim = 4096 + fm_head = nn.Sequential( + nn.Linear(llm_hidden_size, fm_intermediate_dim, bias=True), + nn.GELU(), + nn.Linear(fm_intermediate_dim, output_dim, bias=True), + ) + + self.fm_modules = nn.ModuleDict( + { + "vision_model_mot_gen": vision_model_mot_gen, + "timestep_embedder": timestep_embedder, + "fm_head": fm_head, + } + ) + + if self.top_cfg.add_noise_scale_embedding: + self.fm_modules["noise_scale_embedder"] = TimestepEmbedder(llm_hidden_size) + + # Config shortcuts + self.patch_size = patch_size + self.merge_size = merge_size + self.downsample_ratio = self.top_cfg.downsample_ratio + + # Weight sources for diffusers_loader + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=self.local_model_path, + subfolder=None, + revision=od_config.revision, + prefix="", + fall_back_to_pt=False, + ), + ] + + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler + ) + + # ----------------------------------------------------------------------- + # Helpers + # ----------------------------------------------------------------------- + + def _extract_feature(self, pixel_values, gen_model=False, grid_hw=None): + if gen_model: + return self.fm_modules["vision_model_mot_gen"](pixel_values=pixel_values, grid_hw=grid_hw) + return self.vision_model(pixel_values=pixel_values, grid_hw=grid_hw) + + def _build_t2i_text_inputs(self, query): + model_inputs = self.tokenizer(query, return_tensors="pt") + input_ids = model_inputs["input_ids"].to(self.device) + t_idx = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=self.device) + h_idx = torch.zeros_like(t_idx) + w_idx = torch.zeros_like(t_idx) + indexes = torch.stack([t_idx, h_idx, w_idx], dim=0) + attention_mask = {"full_attention": create_block_causal_mask(indexes[0])} + return input_ids, indexes, attention_mask + + def _get_thw_indexes(self, input_ids, grid_hw=None): + """Build 3D RoPE indexes (t, h, w) from input_ids, assigning spatial + positions to ```` tokens when *grid_hw* is provided.""" + img_start_shift = torch.cat( + [ + torch.zeros(1, dtype=torch.long, device=input_ids.device), + (input_ids == self.img_start_token_id).long(), + ], + dim=0, + )[:-1] + not_img_token = (input_ids != self.img_context_token_id).long() + t_indexes = (img_start_shift + not_img_token).cumsum(0) - 1 + h_indexes = torch.zeros_like(t_indexes) + w_indexes = torch.zeros_like(t_indexes) + + if grid_hw is not None: + selected = input_ids == self.img_context_token_id + if selected.long().sum() > 0: + ds = int(1 / self.downsample_ratio) + abs_w, abs_h = _build_abs_positions_from_grid_hw( + grid_hw // ds, + device=t_indexes.device, + ) + h_indexes[selected] = abs_h.to(t_indexes.device, t_indexes.dtype) + w_indexes[selected] = abs_w.to(t_indexes.device, t_indexes.dtype) + return torch.stack([t_indexes, h_indexes, w_indexes], dim=0) + + def _build_it2i_inputs(self, query, pixel_values=None, grid_hw=None): + """Tokenize *query*, inject ViT features at ```` positions, + and return ``(input_embeds, indexes, attention_mask)``.""" + model_inputs = self.tokenizer(query, return_tensors="pt") + input_ids = model_inputs["input_ids"].to(self.device) + indexes = self._get_thw_indexes(input_ids[0], grid_hw) + attention_mask = {"full_attention": create_block_causal_mask(indexes[0])} + # Route through ``language_model.__call__`` so CPU-offload hooks fire. + input_embeds = self.language_model(input_ids=input_ids, embed_only=True).inputs_embeds + B, N, C = input_embeds.shape + if pixel_values is not None: + vit_embeds = self._extract_feature(pixel_values, grid_hw=grid_hw) + input_embeds = input_embeds.reshape(B * N, C) + flat_ids = input_ids.reshape(B * N) + selected = flat_ids == self.img_context_token_id + assert selected.sum() != 0, "No tokens found in query" + input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device) + input_embeds = input_embeds.reshape(B, N, C) + return input_embeds, indexes, attention_mask + + def _it2i_prefix_forward(self, input_embeds, indexes, attention_mask): + """Run the prefix (text + vision embeddings) through the LM to get KV cache.""" + out = self.language_model( + inputs_embeds=input_embeds, + indexes=indexes, + attention_mask=attention_mask, + use_cache=True, + compute_logits=False, + ) + return out.past_key_values, out.hidden_states + + def _build_t2i_image_indexes(self, token_h, token_w, text_len, device): + t_image = torch.full((token_h * token_w,), text_len, dtype=torch.long, device=device) + idx = torch.arange(token_h * token_w, device=device, dtype=torch.long) + h_image = idx // token_w + w_image = idx % token_w + return torch.stack([t_image, h_image, w_image], dim=0) + + def _t2i_prefix_forward(self, input_ids, indexes, attention_mask): + out = self.language_model( + input_ids=input_ids, + indexes=indexes, + attention_mask=attention_mask, + use_cache=True, + compute_logits=False, + ) + return out.past_key_values, out.hidden_states + + def _t2i_predict_v( + self, input_embeds, indexes_image, attn_mask, past_key_values, t, z, image_token_num, image_size=None, **_kw + ): + B, L = z.shape[0], z.shape[1] + outputs = self.language_model( + inputs_embeds=input_embeds, + image_gen_indicators=torch.ones( + (input_embeds.shape[0], input_embeds.shape[1]), dtype=torch.bool, device=input_embeds.device + ), + indexes=indexes_image, + attention_mask=attn_mask, + past_key_values=past_key_values, + update_cache=False, + use_cache=True, + compute_logits=False, + ) + last_hidden_state = outputs.hidden_states + + if self.top_cfg.use_pixel_head: + merge_size = self.merge_size + token_h = image_size[1] // (self.patch_size * merge_size) + token_w = image_size[0] // (self.patch_size * merge_size) + img_reshaped = last_hidden_state[:, -image_token_num:].view(B, token_h, token_w, -1) + img_2d = torch.einsum("b h w c -> b c h w", img_reshaped).contiguous() + smoothed_img_2d = self.fm_modules["fm_head"](img_2d) + smoothed_reshaped = smoothed_img_2d.view( + B, 3, token_h, self.patch_size * merge_size, token_w, self.patch_size * merge_size + ) + smoothed_reshaped = torch.einsum("b c h p w q -> b h w p q c", smoothed_reshaped) + x_pred = smoothed_reshaped.contiguous().view( + B, L, self.patch_size * merge_size * self.patch_size * merge_size * 3 + ) + else: + x_pred = self.fm_modules["fm_head"](last_hidden_state[:, -image_token_num:].view(B, L, -1)).view(B, L, -1) + + v_pred = (x_pred - z) / (1 - t).clamp_min(self.top_cfg.t_eps) + return v_pred + + def _apply_time_schedule(self, t, image_seq_len, timestep_shift): + sigma = 1 - t + shift = timestep_shift + sigma = shift * sigma / (1 + (shift - 1) * sigma) + return 1 - sigma + + def _ar_step(self, next_token, t_idx, past_key_values): + """One autoregressive token step. Constructs single-token `indexes` explicitly.""" + indexes = torch.tensor([[t_idx + 1], [0], [0]], dtype=torch.long, device=self.device) + outputs = self.language_model( + input_ids=next_token.unsqueeze(0), + indexes=indexes, + past_key_values=past_key_values, + use_cache=True, + ) + return outputs + + def _generate_think(self, prefix_outputs, past_key_values, t_idx, max_think_tokens=1024): + eos_token_id = self.tokenizer.convert_tokens_to_ids("<|im_end|>") + think_end_token_id = self.tokenizer.convert_tokens_to_ids("") + think_token_ids = [] + next_token = torch.argmax(prefix_outputs.logits[:, -1, :], dim=-1) + + for _ in range(max_think_tokens): + token_item = next_token.item() + if token_item == eos_token_id: + break + if token_item == think_end_token_id: + outputs = self._ar_step(next_token, t_idx, past_key_values) + past_key_values = outputs.past_key_values + t_idx += 1 + think_token_ids.append(token_item) + break + + think_token_ids.append(token_item) + outputs = self._ar_step(next_token, t_idx, past_key_values) + past_key_values = outputs.past_key_values + t_idx += 1 + next_token = torch.argmax(outputs.logits[:, -1, :], dim=-1) + + # Append "\n\n" tokens to cache + append_ids = self.tokenizer( + "\n\n", + return_tensors="pt", + add_special_tokens=False, + )["input_ids"].to(self.device) + t_idx = self._append_text_tokens_to_cache(past_key_values, t_idx, append_ids) + + think_text = self.tokenizer.decode(think_token_ids, skip_special_tokens=False) + return past_key_values, t_idx, think_text + + def _append_text_tokens_to_cache(self, cache, t_idx, input_ids): + if input_ids.shape[1] == 0: + return t_idx + seq_len = input_ids.shape[1] + # Route through ``language_model.__call__`` so CPU-offload hooks fire. + inputs_embeds = self.language_model(input_ids=input_ids, embed_only=True).inputs_embeds + t_indexes = torch.arange(t_idx + 1, t_idx + 1 + seq_len, dtype=torch.long, device=self.device) + h_indexes = torch.zeros(seq_len, dtype=torch.long, device=self.device) + w_indexes = torch.zeros(seq_len, dtype=torch.long, device=self.device) + indexes = torch.stack([t_indexes, h_indexes, w_indexes], dim=0) + + past_len = cache.get_seq_length() + mask = torch.zeros(1, 1, seq_len, past_len + seq_len, device=self.device) + causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=self.device)) + causal_mask = torch.where(causal_mask == 1, 0.0, float("-inf")) + mask[:, :, :, past_len:] = causal_mask + attention_mask_dict = {"full_attention": mask} + + self.language_model( + inputs_embeds=inputs_embeds, + indexes=indexes, + attention_mask=attention_mask_dict, + past_key_values=cache, + use_cache=True, + ) + return t_idx + seq_len + + # ----------------------------------------------------------------------- + # Text generation (T2T / I2T) + # ----------------------------------------------------------------------- + + def _generate_text( + self, + prefix_logits, + past_key_values, + t_idx, + max_tokens=512, + do_sample=False, + temperature=0.7, + ): + """Autoregressive text decoding seeded from prefix logits.""" + eos_token_id = self.tokenizer.convert_tokens_to_ids("<|im_end|>") + generated_ids: list[int] = [] + + logits = prefix_logits[:, -1, :] + for _ in range(max_tokens): + if do_sample and temperature > 0: + probs = torch.softmax(logits / temperature, dim=-1) + next_token = torch.multinomial(probs, num_samples=1).squeeze(-1) + else: + next_token = torch.argmax(logits, dim=-1) + + token_id = next_token.item() + if token_id == eos_token_id: + break + generated_ids.append(token_id) + outputs = self._ar_step(next_token, t_idx, past_key_values) + past_key_values = outputs.past_key_values + logits = outputs.logits[:, -1, :] + t_idx += 1 + + return self.tokenizer.decode(generated_ids, skip_special_tokens=True) + + def _build_chat_query(self, prompt, has_images=False): + """Build a chat-style query (understanding mode, no ```` appended).""" + sys_msg = "" + system_prompt = f"<|im_start|>system\n{sys_msg}<|im_end|>\n" if sys_msg else "" + if has_images and "" not in prompt: + prompt = "\n" + prompt + return system_prompt + f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n" + + @torch.inference_mode() + def _forward_text(self, p, input_images) -> DiffusionOutput: + """Text output path for text2text and img2text.""" + extra_args = p.extra_args + max_tokens = int(extra_args.get("max_tokens", 512)) + do_sample = bool(extra_args.get("do_sample", False)) + temperature = float(extra_args.get("temperature", 0.7)) + + if input_images is not None: + pixel_values, grid_hw = self._prepare_input_images(input_images) + query = self._build_chat_query(p.prompt, has_images=True) + # Expand placeholders + for i in range(grid_hw.shape[0]): + h, w = grid_hw[i] + num_patch = int(h * w * self.downsample_ratio**2) + tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * num_patch + IMG_END_TOKEN + query = query.replace("", tokens, 1) + embeds, indexes, mask = self._build_it2i_inputs(query, pixel_values, grid_hw) + outputs = self.language_model( + inputs_embeds=embeds, + indexes=indexes, + attention_mask=mask, + use_cache=True, + ) + else: + query = self._build_chat_query(p.prompt, has_images=False) + input_ids, indexes, mask = self._build_t2i_text_inputs(query) + outputs = self.language_model( + input_ids=input_ids, + indexes=indexes, + attention_mask=mask, + use_cache=True, + ) + + past_kv = outputs.past_key_values + prefix_logits = outputs.logits + t_idx = indexes[0].max().item() + + text_output = self._generate_text( + prefix_logits, + past_kv, + t_idx, + max_tokens=max_tokens, + do_sample=do_sample, + temperature=temperature, + ) + logger.info("Text generation: %d chars", len(text_output)) + return DiffusionOutput( + output=text_output, + custom_output={"text_output": text_output}, + ) + + # ----------------------------------------------------------------------- + # Main forward (T2I / IT2I / T2T / I2T) + # ----------------------------------------------------------------------- + + def _parse_request(self, req: OmniDiffusionRequest): + """Extract common parameters from the request.""" + first_prompt = req.prompts[0] + prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") + extra_args = getattr(req.sampling_params, "extra_args", {}) or {} + height = int(req.sampling_params.height) if req.sampling_params.height else 2048 + width = int(req.sampling_params.width) if req.sampling_params.width else 2048 + grid_factor = self.patch_size * self.merge_size + if height % grid_factor or width % grid_factor: + old_h, old_w = height, width + height = _round_by_factor(height, grid_factor) + width = _round_by_factor(width, grid_factor) + logger.warning( + "Output size %dx%d not divisible by %d, rounded to %dx%d", + old_w, + old_h, + grid_factor, + width, + height, + ) + return SimpleNamespace( + first_prompt=first_prompt, + prompt=prompt, + extra_args=extra_args, + image_size=(width, height), + num_steps=int(req.sampling_params.num_inference_steps or 50), + cfg_scale=float(extra_args.get("cfg_scale", 4.0)), + img_cfg_scale=float(extra_args.get("img_cfg_scale", 1.0)), + cfg_norm=str(extra_args.get("cfg_norm", "none")), + timestep_shift=float(extra_args.get("timestep_shift", 3.0)), + cfg_interval=tuple(extra_args.get("cfg_interval", (0.0, 1.0))), + batch_size=int(extra_args.get("batch_size", 1)), + seed=int(req.sampling_params.seed) if req.sampling_params.seed is not None else 42, + think_mode=bool(extra_args.get("think", False)), + t_eps=float(extra_args.get("t_eps", 0.02)), + ) + + def _extract_input_images(self, first_prompt): + """Return a list of PIL images from ``multi_modal_data``, or *None*.""" + if isinstance(first_prompt, str): + return None + mm_data = first_prompt.get("multi_modal_data") or {} + raw = mm_data.get("image") or mm_data.get("img2img") + if raw is None: + return None + if not isinstance(raw, list): + raw = [raw] + images = [] + for item in raw: + if isinstance(item, str): + images.append(Image.open(item).convert("RGB")) + elif isinstance(item, Image.Image): + images.append(item.convert("RGB")) + else: + raise TypeError(f"Unsupported image type: {type(item)}") + return images if images else None + + def _prepare_input_images(self, images: list[Image.Image]): + """Encode input images: load, smart-resize, patchify → (pixel_values, grid_hw).""" + pixel_values_list = [] + grid_hw_list = [] + num_images = len(images) + max_pixels_per_image = min(2048 * 2048, (4096 * 4096) // num_images) + for img in images: + pv, ghw = _load_image_native( + img, + patch_size=self.vis_cfg.patch_size, + downsample_ratio=self.downsample_ratio, + min_pixels=512 * 512, + max_pixels=max_pixels_per_image, + ) + pixel_values_list.append(pv.to(self.device, dtype=torch.bfloat16)) + grid_hw_list.append(ghw.to(self.device)) + return torch.cat(pixel_values_list), torch.cat(grid_hw_list) + + def _build_it2i_query(self, prompt, images, think_mode): + """Build the tokenizer query string for img2img, expanding ```` + placeholders into ``...`` with the + correct number of patches per image.""" + image_count = prompt.count("") + n_images = len(images["grid_hw"]) + if n_images > image_count: + if image_count == 0 and n_images > 1: + prefix = "".join(f"Image-{i + 1}:\n" for i in range(n_images)) + else: + prefix = "\n" * (n_images - image_count) + prompt = prefix + prompt + + think_content = "\n" if think_mode else "\n\n\n\n" + IMG_START_TOKEN + query = _build_t2i_query(prompt, system_message=SYSTEM_MESSAGE_FOR_GEN, append_text=think_content) + + for i in range(images["grid_hw"].shape[0]): + h, w = images["grid_hw"][i] + num_patch_token = int(h * w * self.downsample_ratio**2) + image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * num_patch_token + IMG_END_TOKEN + query = query.replace("", image_tokens, 1) + return query + + def _init_noise_and_schedule(self, p): + """Init Gaussian noise and compute the flow-matching timestep schedule.""" + merge_size = self.merge_size + grid_h = p.image_size[1] // self.patch_size + grid_w = p.image_size[0] // self.patch_size + token_h = p.image_size[1] // (self.patch_size * merge_size) + token_w = p.image_size[0] // (self.patch_size * merge_size) + grid_hw = torch.tensor([[grid_h, grid_w]] * p.batch_size, device=self.device) + + noise_scale = self.top_cfg.noise_scale + if self.top_cfg.noise_scale_mode in ("resolution", "dynamic", "dynamic_sqrt"): + base = float(self.top_cfg.noise_scale_base_image_seq_len) + scale = math.sqrt((grid_h * grid_w) / (merge_size**2) / base) + noise_scale = scale * float(self.top_cfg.noise_scale) + if self.top_cfg.noise_scale_mode == "dynamic_sqrt": + noise_scale = math.sqrt(noise_scale) + noise_scale = min(noise_scale, self.top_cfg.noise_scale_max_value) + + generator = torch.Generator(self.device).manual_seed(p.seed) + image_prediction = noise_scale * torch.randn( + (p.batch_size, 3, p.image_size[1], p.image_size[0]), + device=self.device, + dtype=self.od_config.dtype, + generator=generator, + ) + + timesteps = torch.linspace(0.0, 1.0, p.num_steps + 1, device=self.device) + timesteps = self._apply_time_schedule(timesteps, token_h * token_w, p.timestep_shift) + + return SimpleNamespace( + merge_size=merge_size, + grid_h=grid_h, + grid_w=grid_w, + token_h=token_h, + token_w=token_w, + grid_hw=grid_hw, + noise_scale=noise_scale, + image_prediction=image_prediction, + timesteps=timesteps, + ) + + def _denoise_step(self, image_prediction, ns, t, z, image_embeds, caches, p, step_i): + """Execute one denoising step with the appropriate CFG strategy.""" + kv_cond = caches["cond"] + idx_cond = caches["idx_cond"] + mask_cond = caches["mask_cond"] + + out_cond = self._t2i_predict_v( + image_embeds, + idx_cond, + mask_cond, + kv_cond, + t, + z, + image_token_num=ns.token_h * ns.token_w, + image_size=p.image_size, + ) + + is_it2i = "img_cond" in caches + if is_it2i: + return self._denoise_step_it2i(out_cond, image_embeds, ns, t, z, caches, p, step_i) + return self._denoise_step_t2i(out_cond, image_embeds, ns, t, z, caches, p, step_i) + + def _denoise_step_t2i(self, out_cond, image_embeds, ns, t, z, caches, p, step_i): + """T2I: two-way CFG (condition vs uncondition), supports cfg_zero_star.""" + if t >= p.cfg_interval[0] and t <= p.cfg_interval[1] and p.cfg_scale > 1: + out_uncond = self._t2i_predict_v( + image_embeds, + caches["idx_uncond"], + caches["mask_uncond"], + caches["uncond"], + t, + z, + image_token_num=ns.token_h * ns.token_w, + image_size=p.image_size, + ) + if p.cfg_norm == "cfg_zero_star": + pos_flat = out_cond.view(p.batch_size, -1) + neg_flat = out_uncond.view(p.batch_size, -1) + alpha = _optimized_scale(pos_flat, neg_flat) + alpha = alpha.view(p.batch_size, *([1] * (len(out_cond.shape) - 1))).to(pos_flat.dtype) + v_pred = ( + out_cond * 0.0 + if step_i <= 0 + else (out_uncond * alpha + p.cfg_scale * (out_cond - out_uncond * alpha)) + ) + else: + v_pred = out_uncond + p.cfg_scale * (out_cond - out_uncond) + v_pred = self._apply_cfg_norm(v_pred, out_cond, p.cfg_norm) + else: + v_pred = out_cond + return v_pred + + def _denoise_step_it2i(self, out_cond, image_embeds, ns, t, z, caches, p, step_i): + """IT2I: dual CFG with ``cfg_scale`` (text) and ``img_cfg_scale`` (image).""" + use_cfg = (t > p.cfg_interval[0] and t < p.cfg_interval[1]) or p.cfg_interval[0] == 0 + needs_cfg = not (p.cfg_scale == 1 and p.img_cfg_scale == 1) + + if not use_cfg or not needs_cfg: + return out_cond + + predict_kw = dict(image_token_num=ns.token_h * ns.token_w, image_size=p.image_size) + + if p.img_cfg_scale == 1: + out_img_cond = self._t2i_predict_v( + image_embeds, + caches["idx_img_cond"], + caches["mask_img_cond"], + caches["img_cond"], + t, + z, + **predict_kw, + ) + v_pred = out_img_cond + p.cfg_scale * (out_cond - out_img_cond) + elif p.cfg_scale == p.img_cfg_scale: + out_uncond = self._t2i_predict_v( + image_embeds, + caches["idx_uncond"], + caches["mask_uncond"], + caches["uncond"], + t, + z, + **predict_kw, + ) + v_pred = out_uncond + p.cfg_scale * (out_cond - out_uncond) + else: + out_img_cond = self._t2i_predict_v( + image_embeds, + caches["idx_img_cond"], + caches["mask_img_cond"], + caches["img_cond"], + t, + z, + **predict_kw, + ) + out_uncond = self._t2i_predict_v( + image_embeds, + caches["idx_uncond"], + caches["mask_uncond"], + caches["uncond"], + t, + z, + **predict_kw, + ) + v_pred = ( + out_uncond + p.cfg_scale * (out_cond - out_img_cond) + p.img_cfg_scale * (out_img_cond - out_uncond) + ) + + if p.cfg_scale > 1 or p.img_cfg_scale > 1: + v_pred = self._apply_cfg_norm(v_pred, out_cond, p.cfg_norm) + return v_pred + + @staticmethod + def _apply_cfg_norm(v_pred, out_cond, cfg_norm): + if cfg_norm == "global": + norm_c = torch.norm(out_cond, dim=(1, 2), keepdim=True) + norm_v = torch.norm(v_pred, dim=(1, 2), keepdim=True) + v_pred = v_pred * (norm_c / (norm_v + 1e-8)).clamp(0, 1.0) + elif cfg_norm == "channel": + norm_c = torch.norm(out_cond, dim=-1, keepdim=True) + norm_v = torch.norm(v_pred, dim=-1, keepdim=True) + v_pred = v_pred * (norm_c / (norm_v + 1e-8)).clamp(0, 1.0) + return v_pred + + def _expand_and_prepare_kv(self, kv, token_hw, batch_size): + """Expand KV cache for batch and prepare flash attention buffers.""" + for layer in kv.layers: + layer.keys = layer.keys.expand(batch_size, *layer.keys.shape[1:]) + layer.values = layer.values.expand(batch_size, *layer.values.shape[1:]) + prepare_flash_kv_cache(kv, current_len=token_hw, batch_size=batch_size) + + @torch.inference_mode() + def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + p = self._parse_request(req) + self.top_cfg.t_eps = p.t_eps + + input_images = self._extract_input_images(p.first_prompt) + modalities = p.first_prompt.get("modalities", []) if isinstance(p.first_prompt, dict) else [] + is_text_output = "text" in modalities + + if is_text_output: + return self._forward_text(p, input_images) + if input_images is not None: + return self._forward_it2i(p, input_images) + return self._forward_t2i(p) + + def _forward_t2i(self, p) -> DiffusionOutput: + """Text-to-image generation path.""" + ns = self._init_noise_and_schedule(p) + + think_content = "\n" if p.think_mode else "\n\n\n\n" + IMG_START_TOKEN + query_cond = _build_t2i_query(p.prompt, system_message=SYSTEM_MESSAGE_FOR_GEN, append_text=think_content) + query_uncond = _build_t2i_query("", append_text=IMG_START_TOKEN) + + input_ids_cond, indexes_cond, mask_cond = self._build_t2i_text_inputs(query_cond) + input_ids_uncond, indexes_uncond, mask_uncond = self._build_t2i_text_inputs(query_uncond) + + indexes_image_cond = self._build_t2i_image_indexes(ns.token_h, ns.token_w, indexes_cond.shape[1], self.device) + indexes_image_uncond = self._build_t2i_image_indexes( + ns.token_h, ns.token_w, indexes_uncond.shape[1], self.device + ) + + think_text = "" + if p.think_mode: + outputs_cond = self.language_model( + input_ids=input_ids_cond, + indexes=indexes_cond, + attention_mask=mask_cond, + use_cache=True, + ) + past_kv_cond = outputs_cond.past_key_values + t_index_cond = indexes_cond[0].max().item() + past_kv_cond, t_index_cond, think_text = self._generate_think(outputs_cond, past_kv_cond, t_index_cond) + indexes_image_cond = self._build_t2i_image_indexes(ns.token_h, ns.token_w, t_index_cond + 1, self.device) + else: + past_kv_cond, _ = self._t2i_prefix_forward(input_ids_cond, indexes_cond, mask_cond) + + past_kv_uncond, _ = self._t2i_prefix_forward(input_ids_uncond, indexes_uncond, mask_uncond) + + self._expand_and_prepare_kv(past_kv_cond, ns.token_h * ns.token_w, p.batch_size) + self._expand_and_prepare_kv(past_kv_uncond, ns.token_h * ns.token_w, p.batch_size) + + caches = { + "cond": past_kv_cond, + "idx_cond": indexes_image_cond, + "mask_cond": {"full_attention": None}, + "uncond": past_kv_uncond, + "idx_uncond": indexes_image_uncond, + "mask_uncond": {"full_attention": None}, + } + return self._run_denoising_loop(ns, caches, p, think_text) + + def _forward_it2i(self, p, input_images: list[Image.Image]) -> DiffusionOutput: + """Image-to-image (editing) generation path with dual CFG.""" + ns = self._init_noise_and_schedule(p) + + pixel_values, grid_hw = self._prepare_input_images(input_images) + images_info = {"grid_hw": grid_hw, "pixel_values": pixel_values} + logger.info("img2img: %d input image(s), grid_hw=%s", len(input_images), grid_hw.tolist()) + + needs_cfg = not (p.cfg_scale == 1 and p.img_cfg_scale == 1) + needs_img_cond = needs_cfg and (p.img_cfg_scale == 1 or p.cfg_scale != p.img_cfg_scale) + needs_uncond = needs_cfg and p.img_cfg_scale != 1 + + # --- Condition query: full prompt + images --- + query_cond = self._build_it2i_query(p.prompt, images_info, p.think_mode) + embeds_cond, idx_cond, mask_cond = self._build_it2i_inputs(query_cond, pixel_values, grid_hw) + + # --- Image-only condition: just "" placeholders + pixel values --- + if needs_img_cond: + img_only_prompt = "" * grid_hw.shape[0] + think_off = "\n\n\n\n" + IMG_START_TOKEN + query_img_cond = _build_t2i_query(img_only_prompt, append_text=think_off) + for i in range(grid_hw.shape[0]): + h, w = grid_hw[i] + num_patch = int(h * w * self.downsample_ratio**2) + tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * num_patch + IMG_END_TOKEN + query_img_cond = query_img_cond.replace("", tokens, 1) + embeds_img_cond, idx_img_cond, mask_img_cond = self._build_it2i_inputs( + query_img_cond, + pixel_values, + grid_hw, + ) + else: + embeds_img_cond = idx_img_cond = mask_img_cond = None + + # --- Uncondition: empty prompt, no images --- + if needs_uncond: + query_uncond = _build_t2i_query("", append_text=IMG_START_TOKEN) + embeds_uncond, idx_uncond, mask_uncond = self._build_it2i_inputs(query_uncond) + else: + embeds_uncond = idx_uncond = mask_uncond = None + + # --- Prefix forwards to build KV caches --- + think_text = "" + if p.think_mode: + outputs_cond = self.language_model( + inputs_embeds=embeds_cond, + indexes=idx_cond, + attention_mask=mask_cond, + use_cache=True, + ) + past_kv_cond = outputs_cond.past_key_values + t_index_cond = idx_cond[0].max().item() + past_kv_cond, t_index_cond, think_text = self._generate_think( + outputs_cond, + past_kv_cond, + t_index_cond, + ) + idx_image_cond = self._build_t2i_image_indexes( + ns.token_h, + ns.token_w, + t_index_cond + 1, + self.device, + ) + else: + past_kv_cond, _ = self._it2i_prefix_forward(embeds_cond, idx_cond, mask_cond) + idx_image_cond = self._build_t2i_image_indexes( + ns.token_h, + ns.token_w, + idx_cond[0].max().item() + 1, + self.device, + ) + + caches = { + "cond": past_kv_cond, + "idx_cond": idx_image_cond, + "mask_cond": {"full_attention": None}, + } + + if needs_img_cond: + past_kv_img_cond, _ = self._it2i_prefix_forward(embeds_img_cond, idx_img_cond, mask_img_cond) + idx_image_img_cond = self._build_t2i_image_indexes( + ns.token_h, + ns.token_w, + idx_img_cond[0].max().item() + 1, + self.device, + ) + caches["img_cond"] = past_kv_img_cond + caches["idx_img_cond"] = idx_image_img_cond + caches["mask_img_cond"] = {"full_attention": None} + + if needs_uncond: + past_kv_uncond, _ = self._it2i_prefix_forward(embeds_uncond, idx_uncond, mask_uncond) + idx_image_uncond = self._build_t2i_image_indexes( + ns.token_h, + ns.token_w, + idx_uncond[0].max().item() + 1, + self.device, + ) + caches["uncond"] = past_kv_uncond + caches["idx_uncond"] = idx_image_uncond + caches["mask_uncond"] = {"full_attention": None} + + # Free vision tensors + del pixel_values, grid_hw, embeds_cond, idx_cond, mask_cond + if embeds_img_cond is not None: + del embeds_img_cond, idx_img_cond, mask_img_cond + if embeds_uncond is not None: + del embeds_uncond, idx_uncond, mask_uncond + + # Expand all KV caches for batch + for key in ("cond", "img_cond", "uncond"): + if key in caches and not isinstance(caches[key], dict): + self._expand_and_prepare_kv(caches[key], ns.token_h * ns.token_w, p.batch_size) + + return self._run_denoising_loop(ns, caches, p, think_text) + + def _run_denoising_loop(self, ns, caches, p, think_text="") -> DiffusionOutput: + """Shared denoising loop for both T2I and IT2I.""" + merge_size = self.merge_size + image_prediction = ns.image_prediction + + for step_i in range(p.num_steps): + t = ns.timesteps[step_i] + t_next = ns.timesteps[step_i + 1] + + z = _patchify(image_prediction, self.patch_size * merge_size) + image_input = _patchify(image_prediction, self.patch_size, channel_first=True) + image_embeds = self._extract_feature( + image_input.view(p.batch_size * ns.grid_h * ns.grid_w, -1), + gen_model=True, + grid_hw=ns.grid_hw, + ).view(p.batch_size, ns.token_h * ns.token_w, -1) + + t_expanded = t.expand(p.batch_size * ns.token_h * ns.token_w) + timestep_embeddings = self.fm_modules["timestep_embedder"](t_expanded).view( + p.batch_size, + ns.token_h * ns.token_w, + -1, + ) + if self.top_cfg.add_noise_scale_embedding: + ns_tensor = torch.full_like(t_expanded, ns.noise_scale / self.top_cfg.noise_scale_max_value) + ns_emb = self.fm_modules["noise_scale_embedder"](ns_tensor).view( + p.batch_size, + ns.token_h * ns.token_w, + -1, + ) + timestep_embeddings = timestep_embeddings + ns_emb + image_embeds = image_embeds + timestep_embeddings + + v_pred = self._denoise_step(image_prediction, ns, t, z, image_embeds, caches, p, step_i) + z = z + (t_next - t) * v_pred + image_prediction = _unpatchify(z, self.patch_size * merge_size, p.image_size[1], p.image_size[0]) + + # Cleanup KV caches + for key in ("cond", "uncond", "img_cond"): + if key in caches and not isinstance(caches[key], dict): + clear_flash_kv_cache(caches[key]) + + images = _to_pil(image_prediction) + img = images[0] if images else None + custom = {} + if think_text: + custom["think_text"] = think_text + return DiffusionOutput(output=img, custom_output=custom) + + # ----------------------------------------------------------------------- + # Weight loading + # ----------------------------------------------------------------------- + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # More specific _mot_gen patterns FIRST to avoid substring + # ambiguity (e.g. `.q_proj` is a substring of `.q_proj_mot_gen`). + (".qkv_proj_mot_gen", ".q_proj_mot_gen", "q"), + (".qkv_proj_mot_gen", ".k_proj_mot_gen", "k"), + (".qkv_proj_mot_gen", ".v_proj_mot_gen", "v"), + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + # MLP gate/up fused into MergedColumnParallelLinear + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + loaded = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + stacked_name = name.replace(weight_name, param_name) + if stacked_name not in params_dict: + break + param = params_dict[stacked_name] + weight_loader = getattr(param, "weight_loader", None) + if weight_loader is not None: + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(stacked_name) + loaded = True + break + + if loaded: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", None) + if weight_loader is not None: + weight_loader(param, loaded_weight) + else: + assert param.shape == loaded_weight.shape, ( + f"Shape mismatch for {name}: param={param.shape} vs loaded={loaded_weight.shape}" + ) + param.data.copy_(loaded_weight) + loaded_params.add(name) + + return loaded_params diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py new file mode 100644 index 000000000..da43f55df --- /dev/null +++ b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py @@ -0,0 +1,741 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Qwen3 LLM with Mixture-of-Tokenizers (MoT) for SenseNova-U1. + +Ported from the sensenova_u1 package with vllm tensor-parallel support: +- QKVParallelLinear for fused q/k/v projections (both und and gen paths) +- MergedColumnParallelLinear for fused gate+up projections +- RowParallelLinear for o_proj and down_proj +- VocabParallelEmbedding for token embeddings +""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass + +import torch +import torch.nn as nn +from transformers.cache_utils import DynamicCache +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention + + +@dataclass +class SenseNovaU1ModelOutput: + last_hidden_state: torch.Tensor + past_key_values: DynamicCache | None = None + + +@dataclass +class SenseNovaU1CausalLMOutput: + logits: torch.Tensor | None = None + past_key_values: DynamicCache | None = None + hidden_states: torch.Tensor | None = None + inputs_embeds: torch.Tensor | None = None + + +# --------------------------------------------------------------------------- +# Mask / Cache utilities +# --------------------------------------------------------------------------- + + +def create_block_causal_mask(index: torch.Tensor): + """Block-wise causal mask from 1D time-index. Returns (1, 1, L, L).""" + L = index.size(0) + idx_i = index.unsqueeze(1).expand(L, L) + idx_j = index.unsqueeze(0).expand(L, L) + arange = torch.arange(L, device=index.device) + mask = (idx_j == idx_i) | (arange.unsqueeze(0) <= arange.unsqueeze(1)) + return torch.where(mask[None, None], 0.0, float("-inf")) + + +def prepare_flash_kv_cache(past_key_values, current_len: int, batch_size: int): + """Convert prefix cache [B,H,S,D] → flash layout [B,S,H,D] and + preallocate buffers for [prefix + current] tokens.""" + if past_key_values is None: + return + for layer in past_key_values.layers: + past_k, past_v = layer.keys, layer.values + if past_k is None or past_v is None: + layer.flash_prefix_len = 0 + layer.flash_total_len = current_len + layer.flash_k_cache = None + layer.flash_v_cache = None + continue + past_k_flash = past_k.transpose(1, 2).contiguous() + past_v_flash = past_v.transpose(1, 2).contiguous() + prefix_len = past_k_flash.shape[1] + total_len = prefix_len + current_len + k_cache = torch.empty( + (batch_size, total_len, past_k_flash.shape[2], past_k_flash.shape[3]), + device=past_k_flash.device, + dtype=past_k_flash.dtype, + ) + v_cache = torch.empty_like(k_cache) + k_cache[:, :prefix_len].copy_(past_k_flash) + v_cache[:, :prefix_len].copy_(past_v_flash) + layer.flash_prefix_len = prefix_len + layer.flash_total_len = total_len + layer.flash_k_cache = k_cache + layer.flash_v_cache = v_cache + + +def clear_flash_kv_cache(past_key_values): + if past_key_values is None: + return + for layer in past_key_values.layers: + for attr in ("flash_prefix_len", "flash_total_len", "flash_k_cache", "flash_v_cache"): + if hasattr(layer, attr): + delattr(layer, attr) + + +# --------------------------------------------------------------------------- +# RoPE +# --------------------------------------------------------------------------- + + +def rotate_half(x): + 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, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) + + +def _compute_default_rope_parameters(config, device=None, **_kw): + base = config.rope_theta + partial_rotary_factor = getattr(config, "partial_rotary_factor", 1.0) + head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + dim = int(head_dim * partial_rotary_factor) + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float().to(device) / dim)) + return inv_freq, 1.0 + + +class Qwen3RotaryEmbedding(nn.Module): + """Custom RoPE that doubles head_dim then subsamples every-other freq.""" + + def __init__(self, config, device=None): + super().__init__() + self.config = config + + def _rope_init_fn(cfg, dev=None): + inv_freq, attn_scaling = _compute_default_rope_parameters(cfg, dev) + cfg2 = copy.deepcopy(cfg) + head_dim = getattr(cfg2, "head_dim", cfg2.hidden_size // cfg2.num_attention_heads) + cfg2.head_dim = int(head_dim) * 2 + inv_freq_full, _ = _compute_default_rope_parameters(cfg2, dev) + return inv_freq_full[::2], attn_scaling + + inv_freq, self.attention_scaling = _rope_init_fn(config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + device_type = x.device.type if x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +# --------------------------------------------------------------------------- +# Norms +# --------------------------------------------------------------------------- + + +class Qwen3RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +# --------------------------------------------------------------------------- +# MLP +# --------------------------------------------------------------------------- + + +class SenseNovaU1MLP(nn.Module): + def __init__( + self, hidden_size: int, intermediate_size: int, hidden_act: str = "silu", quant_config=None, prefix: str = "" + ): + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size, intermediate_size], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + from transformers.activations import ACT2FN + + self.act_fn = ACT2FN[hidden_act] + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + gate, up = gate_up.chunk(2, dim=-1) + x = self.act_fn(gate) * up + x, _ = self.down_proj(x) + return x + + +# --------------------------------------------------------------------------- +# Attention with MoT +# --------------------------------------------------------------------------- + + +class SenseNovaU1Attention(nn.Module): + """Multi-head attention with Mixture-of-Tokenizers (MoT) routing. + + Two parallel sets of Q/K/V/O projections: + - qkv_proj / o_proj: understanding (und) path + - qkv_proj_mot_gen / o_proj_mot_gen: generation (gen) path + + 3D RoPE: time (t), height (h), width (w) with separate frequencies. + """ + + def __init__(self, config, layer_idx: int, quant_config=None, prefix: str = ""): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.total_num_heads = config.num_attention_heads + self.total_num_kv_heads = getattr(config, "num_key_value_heads", config.num_attention_heads) + self.num_kv_groups = self.total_num_heads // self.total_num_kv_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = getattr(config, "attention_dropout", 0.0) + bias = getattr(config, "attention_bias", True) + + # Understanding path + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=bias, + input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + + # Generation (MoT) path + self.qkv_proj_mot_gen = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj_mot_gen", + ) + self.o_proj_mot_gen = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=bias, + input_is_parallel=True, + quant_config=quant_config, + prefix=f"{prefix}.o_proj_mot_gen", + ) + + # Local head counts (TP-sharded) + self.num_heads = self.qkv_proj.num_heads + self.num_kv_heads = self.qkv_proj.num_kv_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + # QK norms — understanding + self.q_norm = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps) + self.k_norm = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps) + self.q_norm_hw = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps) + self.k_norm_hw = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps) + + # QK norms — generation + self.q_norm_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps) + self.k_norm_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps) + self.q_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps) + self.k_norm_hw_mot_gen = Qwen3RMSNorm(self.head_dim // 2, eps=config.rms_norm_eps) + + # 3D RoPE: t uses half head_dim, h/w each use quarter head_dim + t_config = copy.deepcopy(config) + t_config.head_dim = config.head_dim // 2 + self.rotary_emb = Qwen3RotaryEmbedding(config=t_config) + + hw_config = copy.deepcopy(config) + hw_config.head_dim = config.head_dim // 4 + hw_config.rope_theta = config.rope_theta_hw + hw_config.max_position_embeddings = config.max_position_embeddings_hw + self.rotary_emb_hw = Qwen3RotaryEmbedding(config=hw_config) + + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=self.scaling, + num_kv_heads=self.num_heads, + prefix=f"{prefix}.attn", + ) + self.attn.attention = self.attn.sdpa_fallback + + @staticmethod + def _align_mask_dtype(mask: torch.Tensor | None, query: torch.Tensor) -> torch.Tensor | None: + """SDPA requires float ``attn_mask`` to match query's dtype.""" + if mask is None or not mask.is_floating_point() or mask.dtype == query.dtype: + return mask + return mask.to(query.dtype) + + def _run_attn( + self, + query_bhsd: torch.Tensor, + key_bhsd: torch.Tensor, + value_bhsd: torch.Tensor, + attention_mask: torch.Tensor | None, + ) -> torch.Tensor: + """Run unified attention with [B, H, S, D] inputs. Returns [B, S, H, D].""" + if self.num_kv_groups > 1: + n = self.num_kv_groups + key_bhsd = key_bhsd.repeat_interleave(n, dim=1) + value_bhsd = value_bhsd.repeat_interleave(n, dim=1) + q = query_bhsd.transpose(1, 2).contiguous() + k = key_bhsd.transpose(1, 2).contiguous() + v = value_bhsd.transpose(1, 2).contiguous() + attention_mask = self._align_mask_dtype(attention_mask, q) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + return self.attn(q, k, v, attn_metadata) + + def _run_attn_bshd( + self, + query_bshd: torch.Tensor, + key_bshd: torch.Tensor, + value_bshd: torch.Tensor, + attention_mask: torch.Tensor | None, + ) -> torch.Tensor: + """Run unified attention with [B, S, H, D] inputs. Returns [B, S, H, D].""" + if self.num_kv_groups > 1: + n = self.num_kv_groups + key_bshd = key_bshd.repeat_interleave(n, dim=2) + value_bshd = value_bshd.repeat_interleave(n, dim=2) + attention_mask = self._align_mask_dtype(attention_mask, query_bshd) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + return self.attn(query_bshd, key_bshd, value_bshd, attn_metadata) + + def _project_and_rope(self, hidden_states, indexes, qkv_proj, q_norm, k_norm, q_norm_hw, k_norm_hw): + """Project Q/K/V via the given QKVParallelLinear and apply 3D RoPE.""" + input_shape = hidden_states.shape[:-1] # (B, S) + qkv, _ = qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q = q.view(*input_shape, self.num_heads, self.head_dim) + k = k.view(*input_shape, self.num_kv_heads, self.head_dim) + v = v.view(*input_shape, self.num_kv_heads, self.head_dim) + + # Split head_dim into t and hw halves + q_t, q_hw = q.chunk(2, dim=-1) + k_t, k_hw = k.chunk(2, dim=-1) + + # Norms then transpose to [B, H, S, D/2] + q_t = q_norm(q_t).transpose(1, 2) + k_t = k_norm(k_t).transpose(1, 2) + q_hw = q_norm_hw(q_hw).transpose(1, 2) + k_hw = k_norm_hw(k_hw).transpose(1, 2) + + # Split hw into h and w quarters + q_h, q_w = q_hw.chunk(2, dim=-1) + k_h, k_w = k_hw.chunk(2, dim=-1) + + # RoPE for each dimension + cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0)) + q_t, k_t = apply_rotary_pos_emb(q_t, k_t, cos_t, sin_t) + cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0)) + q_h, k_h = apply_rotary_pos_emb(q_h, k_h, cos_h, sin_h) + cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0)) + q_w, k_w = apply_rotary_pos_emb(q_w, k_w, cos_w, sin_w) + + # Reassemble: [B, H, S, head_dim] + query_states = torch.cat([q_t, q_h, q_w], dim=-1) + key_states = torch.cat([k_t, k_h, k_w], dim=-1) + value_states = v.transpose(1, 2) # [B, H, S, D] + return query_states, key_states, value_states + + def forward_und(self, hidden_states, indexes, attention_mask, past_key_values=None, **kwargs): + """Understanding path — unified Attention with explicit 4D mask.""" + input_shape = hidden_states.shape[:-1] + query_states, key_states, value_states = self._project_and_rope( + hidden_states, + indexes, + self.qkv_proj, + self.q_norm, + self.k_norm, + self.q_norm_hw, + self.k_norm_hw, + ) + update_cache = kwargs.get("update_cache", True) + if past_key_values is not None: + if update_cache: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + else: + layer = past_key_values.layers[self.layer_idx] + if layer.keys is not None: + key_states = torch.cat([layer.keys, key_states], dim=2) + value_states = torch.cat([layer.values, value_states], dim=2) + + attn_output = self._run_attn(query_states, key_states, value_states, attention_mask) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output, _ = self.o_proj(attn_output) + return attn_output + + def forward_gen(self, hidden_states, indexes, attention_mask, past_key_values=None, **kwargs): + """Generation path — unified Attention, bidirectional with optional KV cache.""" + input_shape = hidden_states.shape[:-1] + query_states, key_states, value_states = self._project_and_rope( + hidden_states, + indexes, + self.qkv_proj_mot_gen, + self.q_norm_mot_gen, + self.k_norm_mot_gen, + self.q_norm_hw_mot_gen, + self.k_norm_hw_mot_gen, + ) + update_cache = kwargs.get("update_cache", True) + + if attention_mask is None: + # Bidirectional path: no causal mask, optionally attend to a prefix. + q = query_states.transpose(1, 2).contiguous() # [B,S,H,D] + k_cur = key_states.transpose(1, 2).contiguous() + v_cur = value_states.transpose(1, 2).contiguous() + + if past_key_values is not None: + if update_cache: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + k = key_states.transpose(1, 2).contiguous() + v = value_states.transpose(1, 2).contiguous() + else: + layer = past_key_values.layers[self.layer_idx] + if hasattr(layer, "flash_k_cache") and layer.flash_k_cache is not None: + prefix_len = layer.flash_prefix_len + cur_len = k_cur.shape[1] + layer.flash_k_cache[:, prefix_len : prefix_len + cur_len].copy_(k_cur) + layer.flash_v_cache[:, prefix_len : prefix_len + cur_len].copy_(v_cur) + k = layer.flash_k_cache[:, : prefix_len + cur_len] + v = layer.flash_v_cache[:, : prefix_len + cur_len] + else: + past_k = past_key_values.layers[self.layer_idx].keys + past_v = past_key_values.layers[self.layer_idx].values + if past_k is not None: + k = torch.cat([past_k.transpose(1, 2).contiguous(), k_cur], dim=1) + v = torch.cat([past_v.transpose(1, 2).contiguous(), v_cur], dim=1) + else: + k, v = k_cur, v_cur + else: + k, v = k_cur, v_cur + + attn_output = self._run_attn_bshd(q, k, v, None) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output, _ = self.o_proj_mot_gen(attn_output) + return attn_output + + # Masked fallback with explicit 4D additive mask + if past_key_values is not None: + if update_cache: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + else: + layer = past_key_values.layers[self.layer_idx] + if layer.keys is not None: + key_states = torch.cat([layer.keys, key_states], dim=2) + value_states = torch.cat([layer.values, value_states], dim=2) + + attn_output = self._run_attn(query_states, key_states, value_states, attention_mask) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output, _ = self.o_proj_mot_gen(attn_output) + return attn_output + + def forward( + self, + hidden_states, + image_gen_indicators, + exist_und, + exist_gen, + indexes, + attention_mask, + past_key_values=None, + **kwargs, + ): + if exist_und and not exist_gen: + return self.forward_und(hidden_states, indexes, attention_mask, past_key_values, **kwargs) + if not exist_und and exist_gen: + return self.forward_gen(hidden_states, indexes, attention_mask, past_key_values, **kwargs) + raise NotImplementedError("Mixed und+gen tokens in a single forward not implemented for initial port") + + +# --------------------------------------------------------------------------- +# Decoder Layer +# --------------------------------------------------------------------------- + + +class SenseNovaU1DecoderLayer(nn.Module): + def __init__(self, config, layer_idx: int, quant_config=None, prefix: str = ""): + super().__init__() + self.self_attn = SenseNovaU1Attention( + config, layer_idx, quant_config=quant_config, prefix=f"{prefix}.self_attn" + ) + self.mlp = SenseNovaU1MLP( + config.hidden_size, + config.intermediate_size, + config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.mlp_mot_gen = SenseNovaU1MLP( + config.hidden_size, + config.intermediate_size, + config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp_mot_gen", + ) + self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_layernorm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.attention_type = config.layer_types[layer_idx] + + def _forward_und(self, hidden_states, indexes, attention_mask, past_key_values, **kwargs): + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states, + image_gen_indicators=None, + exist_und=True, + exist_gen=False, + indexes=indexes, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp(self.post_attention_layernorm(hidden_states)) + return residual + hidden_states + + def _forward_gen(self, hidden_states, indexes, attention_mask, past_key_values, **kwargs): + residual = hidden_states + hidden_states = self.input_layernorm_mot_gen(hidden_states) + hidden_states = self.self_attn( + hidden_states, + image_gen_indicators=None, + exist_und=False, + exist_gen=True, + indexes=indexes, + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.mlp_mot_gen(self.post_attention_layernorm_mot_gen(hidden_states)) + return residual + hidden_states + + def forward( + self, + hidden_states, + image_gen_indicators, + exist_und, + exist_gen, + indexes, + attention_mask, + past_key_values=None, + **kwargs, + ): + if exist_und and not exist_gen: + return self._forward_und(hidden_states, indexes, attention_mask, past_key_values, **kwargs) + if not exist_und and exist_gen: + return self._forward_gen(hidden_states, indexes, attention_mask, past_key_values, **kwargs) + raise NotImplementedError("Mixed und+gen tokens in a single forward not implemented for initial port") + + +# --------------------------------------------------------------------------- +# Model (decoder stack) +# --------------------------------------------------------------------------- + + +class SenseNovaU1Model(nn.Module): + def __init__(self, config, quant_config=None, prefix: str = ""): + super().__init__() + self.config = config + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=f"{prefix}.embed_tokens", + ) + self.layers = nn.ModuleList( + [ + SenseNovaU1DecoderLayer(config, i, quant_config=quant_config, prefix=f"{prefix}.layers.{i}") + for i in range(config.num_hidden_layers) + ] + ) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids=None, + image_gen_indicators=None, + indexes=None, + attention_mask=None, + past_key_values=None, + inputs_embeds=None, + use_cache=None, + **kwargs, + ): + if image_gen_indicators is None: + exist_und, exist_gen = True, False + else: + exist_und = (~image_gen_indicators).any().item() + exist_gen = image_gen_indicators.any().item() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache() + + # Resolve attention mask. Callers must always provide `indexes`; this + # keeps the model stateless and safe under concurrent requests. + assert indexes is not None, "SenseNovaU1Model.forward requires explicit `indexes`." + if not isinstance(attention_mask, dict): + past_len = past_key_values.get_seq_length() if past_key_values else 0 + seq_len = inputs_embeds.shape[1] + total_len = past_len + seq_len + mask = torch.zeros(1, 1, seq_len, total_len, device=inputs_embeds.device) + if seq_len > 1: + causal = torch.tril(torch.ones(seq_len, seq_len, device=inputs_embeds.device)) + mask[:, :, :, past_len:] = torch.where(causal == 1, 0.0, float("-inf")) + causal_mask_mapping = {"full_attention": mask} + else: + causal_mask_mapping = attention_mask + + hidden_states = inputs_embeds + for layer in self.layers: + hidden_states = layer( + hidden_states, + image_gen_indicators=image_gen_indicators, + exist_und=exist_und, + exist_gen=exist_gen, + indexes=indexes, + attention_mask=causal_mask_mapping.get(layer.attention_type, causal_mask_mapping.get("full_attention")), + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + if not exist_gen: + hidden_states = self.norm(hidden_states) + else: + hidden_states = self.norm_mot_gen(hidden_states) + + return SenseNovaU1ModelOutput( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + +# --------------------------------------------------------------------------- +# ForCausalLM wrapper +# --------------------------------------------------------------------------- + + +class SenseNovaU1ForCausalLM(nn.Module): + def __init__(self, config, quant_config=None, prefix: str = ""): + super().__init__() + self.config = config + self.model = SenseNovaU1Model(config, quant_config=quant_config, prefix=f"{prefix}.model") + self.vocab_size = config.vocab_size + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, prefix=f"{prefix}.lm_head") + # LogitsProcessor handles the TP all-gather of vocab-sharded ParallelLMHead + # outputs so callers see full-vocab logits regardless of tp_size. + self.logits_processor = LogitsProcessor(config.vocab_size) + + def forward( + self, + input_ids=None, + indexes=None, + attention_mask=None, + past_key_values=None, + inputs_embeds=None, + use_cache=None, + embed_only: bool = False, + compute_logits: bool = True, + **kwargs, + ): + # Routing every access through ``forward`` keeps any externally-attached + # hooks (e.g. CPU-offload swap) firing for sub-module accesses such as + # token embedding lookup or hidden-state-only model runs. Callers should + # prefer this entry point over reaching into ``self.model`` / + # ``self.lm_head`` directly so that offloaded weights get materialised + # on the right device first. + if embed_only: + if input_ids is None: + raise ValueError("embed_only=True requires input_ids") + return SenseNovaU1CausalLMOutput( + inputs_embeds=self.model.embed_tokens(input_ids), + ) + + outputs = self.model( + input_ids=input_ids, + indexes=indexes, + attention_mask=attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + logits = self.logits_processor(self.lm_head, outputs.last_hidden_state) if compute_logits else None + return SenseNovaU1CausalLMOutput( + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.last_hidden_state, + ) + + def get_input_embeddings(self): + return self.model.embed_tokens diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index ad3e1b7c2..8fb675b69 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -222,6 +222,13 @@ class ModelLevelOffloadBackend(OffloadBackend): except Exception as exc: logger.debug("Failed to move VAE to GPU: %s", exc) + # Pin resident modules on GPU (small hot submodules called inside the DiT loop). + for res, name in zip(modules.resident_modules, modules.resident_names): + try: + res.to(self.device) + except Exception as exc: + logger.warning("Failed to move resident module '%s' to GPU: %s", name, exc) + # Apply sequential offloading hooks apply_sequential_offload( dit_modules=modules.dits, @@ -237,9 +244,10 @@ class ModelLevelOffloadBackend(OffloadBackend): self.enabled = True logger.info( - "Model-level offloading enabled: %s <-> %s (mutual exclusion)", + "Model-level offloading enabled: %s <-> %s (mutual exclusion)%s", ", ".join(modules.dit_names), ", ".join(modules.encoder_names), + f"; resident on GPU: {', '.join(modules.resident_names)}" if modules.resident_names else "", ) def disable(self) -> None: diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index f530115d7..b6a5a3700 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -221,6 +221,11 @@ _DIFFUSION_MODELS = { "pipeline_dreamid_omni", "DreamIDOmniPipeline", ), + "SenseNovaU1Pipeline": ( + "sensenova_u1", + "pipeline_sensenova_u1", + "SenseNovaU1Pipeline", + ), "AudioXPipeline": ( "audiox", "pipeline_audiox", @@ -272,6 +277,7 @@ DiffusionModelRegistry = _ModelRegistry( _NO_CACHE_ACCELERATION = { # Pipelines that do not support cache acceleration (cache_dit / tea_cache). "NextStep11Pipeline", + "SenseNovaU1Pipeline", "AudioXPipeline", } @@ -475,6 +481,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "MagiHumanPipeline": "get_magi_human_post_process_func", "OmniVoicePipeline": "get_omnivoice_post_process_func", "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", + "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 531c62d26..4e84d6200 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -450,11 +450,12 @@ class OmniBase(PDDisaggregationMixin): except Exception: logger.exception("[%s] Finalize request handling error", self.__class__.__name__) - images = getattr(engine_outputs, "images", []) if stage_meta["final_output_type"] == "image" else [] + output_type = getattr(engine_outputs, "final_output_type", stage_meta["final_output_type"]) + images = getattr(engine_outputs, "images", []) if output_type == "image" else [] return OmniRequestOutput( request_id=req_id or "", stage_id=stage_id, - final_output_type=stage_meta["final_output_type"], + final_output_type=output_type, request_output=engine_outputs, images=images, trajectory_latents=getattr(engine_outputs, "trajectory_latents", None), diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 7558e85aa..24d2e6592 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -14,6 +14,7 @@ from fastapi import Request from PIL import Image from pydantic import TypeAdapter +from vllm_omni.diffusion.diffusion_engine import get_extra_body_params, get_extra_output_params from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse from vllm_omni.entrypoints.utils import coerce_param_message_types @@ -95,7 +96,9 @@ from vllm_omni.entrypoints.openai.stage_params import ( from vllm_omni.entrypoints.openai.utils import ( get_stage_type, get_supported_speakers_from_hf_config, + is_single_stage_diffusion, parse_lora_request, + resolve_diffusion_od_config, validate_requested_speaker, ) from vllm_omni.lora.request import LoRARequest @@ -119,6 +122,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): _diffusion_engine: AsyncOmni | None = None _diffusion_model_name: str = "" _supported_speakers: set[str] | None = None + _diffusion_extra_body_params: frozenset[str] | None = None + _diffusion_extra_output_params: frozenset[str] | None = None @classmethod def for_diffusion( @@ -145,6 +150,45 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): instance._diffusion_model_name = model_name return instance + def _get_diffusion_extra_body_params(self) -> frozenset[str]: + """Return the pipeline-declared ``EXTRA_BODY_PARAMS`` (cached).""" + if self._diffusion_extra_body_params is not None: + return self._diffusion_extra_body_params + + params: frozenset[str] = frozenset() + try: + od_config = resolve_diffusion_od_config(self.engine_client, self._diffusion_engine) + if od_config is not None and getattr(od_config, "model_class_name", None): + params = get_extra_body_params(od_config.model_class_name) + except Exception as e: + logger.warning("Failed to read EXTRA_BODY_PARAMS from pipeline: %s", e) + + self._diffusion_extra_body_params = params + return params + + def _get_diffusion_extra_output_params( + self, + custom_output: dict[str, Any] | None, + ) -> dict[str, Any] | None: + """Pick pipeline-declared ``EXTRA_OUTPUT_PARAMS`` keys from *custom_output*.""" + if not custom_output: + return None + + if self._diffusion_extra_output_params is None: + params: frozenset[str] = frozenset() + try: + od_config = resolve_diffusion_od_config(self.engine_client, self._diffusion_engine) + if od_config is not None and getattr(od_config, "model_class_name", None): + params = get_extra_output_params(od_config.model_class_name) + except Exception as e: + logger.warning("Failed to read EXTRA_OUTPUT_PARAMS from pipeline: %s", e) + self._diffusion_extra_output_params = params + + if not self._diffusion_extra_output_params: + return None + out = {k: custom_output[k] for k in self._diffusion_extra_output_params if k in custom_output} + return out or None + def _get_supported_speakers(self) -> set[str]: """Load supported speakers from model config (cached).""" if self._supported_speakers is not None: @@ -386,6 +430,34 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): logger.warning("Failed to build image-generation prompt for omni multistage: %s", e) _image_gen_height = None _image_gen_width = None + elif request.modalities and ("text" in request.modalities) and is_single_stage_diffusion(self.engine_client): + # Single-stage diffusion text output (img2text / text2text). + # Build a diffusion-style prompt with modalities=["text"] so the + # pipeline routes to its text generation path. + try: + extracted_prompt, reference_images = self._extract_diffusion_prompt_and_images_from_messages( + request.messages + ) + if not extracted_prompt: + return self.create_error_response("No text prompt found in messages") + + tprompt: OmniTextPrompt = {"prompt": extracted_prompt} + tprompt["modalities"] = ["text"] + + if reference_images: + try: + img_bytes = base64.b64decode(reference_images[0]) + img = Image.open(BytesIO(img_bytes)) + tprompt["multi_modal_data"] = {"image": img} + tprompt["multi_modal_uuids"] = {"image": [f"{request_id}-image-0"]} + except Exception: + pass + + engine_prompts = [tprompt] + except Exception as e: + logger.warning("Failed to build text-output prompt for single-stage diffusion: %s", e) + _image_gen_height = None + _image_gen_width = None else: _image_gen_height = None _image_gen_width = None @@ -418,6 +490,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): sp.extra_args["cfg_text_scale"] = cfg_text_scale if cfg_img_scale is not None: sp.extra_args["cfg_img_scale"] = cfg_img_scale + for _key in self._get_diffusion_extra_body_params(): + _val = extra_body.get(_key) + if _val is not None: + sp.extra_args[_key] = _val self._log_inputs( request_id, @@ -1610,21 +1686,37 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): continue if omni_outputs.final_output_type == "text": - ( - choices_data, - usage, - prompt_logprobs, - prompt_token_ids, - kv_transfer_params, - ) = self._create_text_choice( - request, - omni_outputs, - tokenizer, - conversation, - role, - reasoning_parser, - ) - final_res = omni_outputs.request_output + if omni_outputs.request_output is not None: + ( + choices_data, + usage, + prompt_logprobs, + prompt_token_ids, + kv_transfer_params, + ) = self._create_text_choice( + request, + omni_outputs, + tokenizer, + conversation, + role, + reasoning_parser, + ) + final_res = omni_outputs.request_output + else: + # Diffusion pipeline text output (e.g. single-stage + # img2text / text2text) — no AR request_output, so build + # a simple text choice from custom_output. + text_body = (omni_outputs.custom_output or {}).get("text_output", "") + message = ChatMessage(role=role, content=text_body) + choices_data = [ + ChatCompletionResponseChoice( + index=0, + message=message, + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + ] elif omni_outputs.final_output_type == "audio": choices_data = self._create_audio_choice(omni_outputs, role, request, stream=False) elif omni_outputs.final_output_type == "image": @@ -1640,6 +1732,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): response_metrics = {} response_metrics.setdefault("stage_durations", omni_outputs.stage_durations or {}) response_metrics.setdefault("peak_memory_mb", float(omni_outputs.peak_memory_mb or 0.0)) + extra = self._get_diffusion_extra_output_params(omni_outputs.custom_output) + if extra: + response_metrics.update(extra) choices.extend(choices_data) response = OmniChatCompletionResponse( @@ -2511,6 +2606,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): gen_params.extra_args["cfg_text_scale"] = cfg_text_scale if cfg_img_scale is not None: gen_params.extra_args["cfg_img_scale"] = cfg_img_scale + for _key in self._get_diffusion_extra_body_params(): + _val = extra_body.get(_key) + if _val is not None: + gen_params.extra_args[_key] = _val if num_frames is not None: gen_params.num_frames = num_frames if guidance_scale_2 is not None: @@ -2539,6 +2638,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): except Exception as e: # pragma: no cover - safeguard logger.warning("Failed to parse LoRA request: %s", e) + # Route text modality for single-stage diffusion (img2text / text2text) + requested_modalities = extra_body.get("modalities") or [] + is_text_request = "text" in requested_modalities + + if is_text_request: + gen_prompt["modalities"] = ["text"] + # Add reference image if provided (from messages content) if pil_images: if len(pil_images) == 1: @@ -2588,6 +2694,38 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): result = output if result is None: return self._create_error_response("No output generated from AsyncOmni") + + # Text output path (img2text / text2text) + if is_text_request and result.final_output_type == "text": + text_body = (result.custom_output or {}).get("text_output", "") + message = ChatMessage(role="assistant", content=text_body) + choice = ChatCompletionResponseChoice( + index=0, + message=message, + finish_reason="stop", + logprobs=None, + stop_reason=None, + ) + response = OmniChatCompletionResponse( + id=request_id, + created=created_time, + model=self._diffusion_model_name, + choices=[choice], + usage=UsageInfo( + prompt_tokens=len(prompt.split()), + completion_tokens=len(text_body.split()), + total_tokens=len(prompt.split()) + len(text_body.split()), + ), + metrics=self._get_diffusion_extra_output_params(result.custom_output), + ) + logger.info( + "Diffusion chat completed for request %s: text output (%d chars)", + request_id, + len(text_body), + ) + return response + + # Image output path (text2img / img2img) final_output_type = getattr(result, "final_output_type", "image") # Handle nested OmniRequestOutput structure where images might be in request_output images = getattr(result.request_output, "images", []) @@ -2702,7 +2840,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): stop_reason=None, ) - response = ChatCompletionResponse( + response = OmniChatCompletionResponse( id=request_id, created=created_time, model=self._diffusion_model_name, @@ -2712,6 +2850,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): completion_tokens=1, total_tokens=len(prompt.split()) + 1, ), + metrics=self._get_diffusion_extra_output_params(result.custom_output), ) logger.info( diff --git a/vllm_omni/entrypoints/openai/utils.py b/vllm_omni/entrypoints/openai/utils.py index f411526fd..5cca49129 100644 --- a/vllm_omni/entrypoints/openai/utils.py +++ b/vllm_omni/entrypoints/openai/utils.py @@ -70,6 +70,27 @@ def get_supported_speakers_from_hf_config(hf_config: Any) -> set[str]: return set() +def resolve_diffusion_od_config(engine_client: Any, diffusion_engine: Any = None) -> Any: + """Resolve the OmniDiffusionConfig from the engine or diffusion engine.""" + od_config = None + if hasattr(engine_client, "get_diffusion_od_config"): + od_config = engine_client.get_diffusion_od_config() + if od_config is None and diffusion_engine is not None: + if hasattr(diffusion_engine, "get_diffusion_od_config"): + od_config = diffusion_engine.get_diffusion_od_config() + else: + od_config = getattr(diffusion_engine, "od_config", None) + return od_config + + +def is_single_stage_diffusion(engine_client: Any) -> bool: + """Return True if the engine is a single-stage diffusion pipeline.""" + stage_configs = getattr(engine_client, "stage_configs", None) or [] + if len(stage_configs) != 1: + return False + return getattr(stage_configs[0], "stage_type", None) in ("diffusion", "DIFFUSION") + + def validate_requested_speaker(speaker: str | None, supported_speakers: set[str]) -> str | None: """Normalize and validate an optional speaker value.