diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 023ad1932..119f3bf68 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -48,6 +48,7 @@ th { | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 39dc36648..526d2fd45 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -193,6 +193,7 @@ Factory function `get_offload_backend()` selects the appropriate backend based o | NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` | | OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` | | QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` | +| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | `"down_blocks"`, `"up_blocks"` | | StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` | | Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | | Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 234446c01..6c96015d4 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -130,6 +130,7 @@ The following tables show which models support each feature: | **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **SenseNova-U1** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | | **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ | | **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | diff --git a/recipes/StabilityAI/Stable-Diffusion-XL.md b/recipes/StabilityAI/Stable-Diffusion-XL.md new file mode 100644 index 000000000..94c7e9646 --- /dev/null +++ b/recipes/StabilityAI/Stable-Diffusion-XL.md @@ -0,0 +1,123 @@ +# Stable Diffusion XL + +> Text-to-image generation with Stability AI's SDXL Base 1.0 + +## Summary + +- Vendor: Stability AI +- Model: `stabilityai/stable-diffusion-xl-base-1.0` +- Task: Text-to-image generation +- Mode: Offline inference / Online serving with the OpenAI-compatible images API +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want a high-quality, open-weight text-to-image model +that generates 1024x1024 images. SDXL Base 1.0 is one of the most widely used +diffusion models and produces photorealistic outputs with strong prompt +adherence. It requires ~6.5 GiB for model weights and runs inference at +~7.5 seconds per image (30 steps) on a single Intel Data Center GPU. + +## References + +- Model: +- Architecture reference: +- Related example under `examples/`: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) + +## Hardware Support + +## GPU + +### 1x Intel Arc BMG GPU (XPU) + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: Intel XPU environment with Intel Arc BMG GPU +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +**Offline inference:** + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model stabilityai/stable-diffusion-xl-base-1.0 \ + --prompt "a beautiful sunset over the ocean, photorealistic" \ + --height 1024 --width 1024 \ + --num-inference-steps 30 \ + --guidance-scale 7.5 \ + --seed 42 \ + --output sdxl_output.png +``` + +**Online serving:** + +```bash +vllm serve stabilityai/stable-diffusion-xl-base-1.0 --omni --port 8091 +``` + +#### Verification + +**For offline inference:** Check the output image file exists and has valid content. + +**For online serving:** + +```bash +curl -X POST http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "stabilityai/stable-diffusion-xl-base-1.0", + "prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus", + "negative_prompt": "blurry, out of focus, low quality", + "guidance_scale": 7.5, + "n": 1, + "size": "1024x1024" + }' | python3 -c " +import json, base64, sys +data = json.load(sys.stdin) +with open('output.png', 'wb') as f: + f.write(base64.b64decode(data['data'][0]['b64_json'])) +print('saved output.png') +" +``` + +#### Notes + +- Memory usage: Model weights ~6.5 GiB, peak memory ~10.9 GiB at 1024x1024 with CFG. +- Generation time: ~7.5 seconds for 30 inference steps at 1024x1024. +- `negative_prompt` requires `guidance_scale > 1` to take effect. Recommended range: 5.0-9.0. +- Default resolution is 1024x1024. Other supported sizes: 512x512, 768x768. +- The model uses dual CLIP text encoders (ViT-L/14 + OpenCLIP ViT-bigG/14) for rich text understanding. + +### 1x NVIDIA GPU (24+ GB VRAM) + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: NVIDIA CUDA environment with an A100 +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model stabilityai/stable-diffusion-xl-base-1.0 \ + --prompt "a beautiful sunset over the ocean, photorealistic" \ + --height 1024 --width 1024 \ + --num-inference-steps 30 \ + --guidance-scale 7.5 \ + --seed 42 \ + --output sdxl_output.png +``` + +#### Notes + +- Memory usage: ~6.5 GiB model weights, ~10 GiB peak with CFG at 1024x1024. +- Fits on a single GPU with 24 GB VRAM (e.g., RTX 3090/4090, A5000). +- For GPUs with less VRAM, use `--enable-cpu-offload` to offload text encoders during diffusion. diff --git a/tests/e2e/online_serving/test_sdxl.py b/tests/e2e/online_serving/test_sdxl.py new file mode 100644 index 000000000..4d7fc34d5 --- /dev/null +++ b/tests/e2e/online_serving/test_sdxl.py @@ -0,0 +1,59 @@ +""" +Tests for Stable Diffusion XL Base 1.0 model. + +Validates single-GPU text-to-image generation with CFG (guidance_scale > 1). +""" + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import ( + OmniServer, + OmniServerParams, + OpenAIClientHandler, + dummy_messages_from_mix_data, +) + +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}) +POSITIVE_PROMPT = "A serene mountain landscape at sunset, photorealistic" +NEGATIVE_PROMPT = "blurry, low quality, distorted" + + +def _get_diffusion_feature_cases(model: str): + return [ + pytest.param( + OmniServerParams( + model=model, + ), + id="default", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + +@pytest.mark.advanced_model +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases( + model="stabilityai/stable-diffusion-xl-base-1.0", + ), + indirect=True, +) +def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): + messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT) + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 1024, + "width": 1024, + "num_inference_steps": 30, + "negative_prompt": NEGATIVE_PROMPT, + "guidance_scale": 7.5, + "seed": 42, + }, + } + + openai_client.send_diffusion_request(request_config) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 00f3b0b5a..a2dfb5a15 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -903,7 +903,12 @@ class OmniDiffusionConfig: ) else: tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) - self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) + if tf_config_dict is None: + tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model) + if tf_config_dict is not None: + self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) + else: + self.set_tf_model_config(TransformerConfig()) else: raise FileNotFoundError("model_index.json not found") except (AttributeError, OSError, ValueError, FileNotFoundError): diff --git a/vllm_omni/diffusion/models/sdxl/__init__.py b/vllm_omni/diffusion/models/sdxl/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/diffusion/models/sdxl/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py b/vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py new file mode 100644 index 000000000..d4a966cb0 --- /dev/null +++ b/vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import logging +import os +from collections.abc import Iterable + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.schedulers.scheduling_euler_discrete import EulerDiscreteScheduler +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery +from vllm_omni.diffusion.models.sdxl.sdxl_unet import SDXLUNet2DConditionModel +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest + +logger = logging.getLogger(__name__) + + +def get_sdxl_image_post_process_func(od_config: OmniDiffusionConfig): + if od_config.output_type == "latent": + return lambda x: x + image_processor = VaeImageProcessor(vae_scale_factor=8) + + def post_process_func(images: torch.Tensor): + if images.device.type != "cpu": + images = images.cpu() + return image_processor.postprocess(images) + + return post_process_func + + +class StableDiffusionXLPipeline( + nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): + _dit_modules: list[str] = ["unet"] + _encoder_modules: list[str] = ["text_encoder", "text_encoder_2"] + _vae_modules: list[str] = ["vae"] + _resident_modules: list[str] = [] + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="unet", + revision=None, + prefix="unet.", + fall_back_to_pt=True, + ) + ] + + self.device = get_local_device() + model = od_config.model + local_files_only = os.path.exists(model) + dtype = od_config.dtype + + sdxl_subfolders = [ + "scheduler", + "tokenizer", + "tokenizer_2", + "text_encoder", + "text_encoder_2", + "vae", + ] + prefetch_subfolders(model, sdxl_subfolders, local_files_only=local_files_only) + + self.scheduler = EulerDiscreteScheduler.from_pretrained( + model, + subfolder="scheduler", + local_files_only=local_files_only, + torch_dtype=torch.float32, + ) + self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.tokenizer_2 = CLIPTokenizer.from_pretrained( + model, subfolder="tokenizer_2", local_files_only=local_files_only + ) + self.text_encoder = from_pretrained_with_prefetch( + CLIPTextModel.from_pretrained, + model, + subfolder="text_encoder", + prefetch_list=sdxl_subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ) + self.text_encoder_2 = from_pretrained_with_prefetch( + CLIPTextModelWithProjection.from_pretrained, + model, + subfolder="text_encoder_2", + prefetch_list=sdxl_subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ) + self.unet = SDXLUNet2DConditionModel(od_config=od_config) + self.vae = from_pretrained_with_prefetch( + DistributedAutoencoderKL.from_pretrained, + model, + subfolder="vae", + prefetch_list=sdxl_subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ).to(self.device) + + self.vae_scale_factor = 8 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.tokenizer_max_length = 77 + self.default_sample_size = 128 # 1024 / 8 + self.output_type = self.od_config.output_type + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + def _get_clip_prompt_embeds( + self, + prompt: str | list[str], + tokenizer: CLIPTokenizer, + text_encoder: nn.Module, + num_images_per_prompt: int = 1, + return_pooled: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + + text_encoder_device = next(text_encoder.parameters()).device + outputs = text_encoder(text_input_ids.to(text_encoder_device), output_hidden_states=True) + prompt_embeds = outputs.hidden_states[-2].to(dtype=self.od_config.dtype, device=self.device) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + pooled_prompt_embeds = None + if return_pooled: + pooled_prompt_embeds = outputs[0].to(dtype=self.od_config.dtype, device=self.device) + pooled_prompt_embeds = pooled_prompt_embeds.repeat(num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + return prompt_embeds, pooled_prompt_embeds + + def encode_prompt( + self, + prompt: str | list[str], + num_images_per_prompt: int = 1, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Encode with first CLIP (ViT-L/14, 768-dim hidden states) + prompt_embeds_1, _ = self._get_clip_prompt_embeds( + prompt, + self.tokenizer, + self.text_encoder, + num_images_per_prompt, + return_pooled=False, + ) + # Encode with second CLIP (OpenCLIP ViT-bigG/14, 1280-dim hidden states + pooled) + prompt_embeds_2, pooled_prompt_embeds = self._get_clip_prompt_embeds( + prompt, + self.tokenizer_2, + self.text_encoder_2, + num_images_per_prompt, + return_pooled=True, + ) + + # Concatenate along feature dimension: 768 + 1280 = 2048 + prompt_embeds = torch.cat([prompt_embeds_1, prompt_embeds_2], dim=-1) + + return prompt_embeds, pooled_prompt_embeds + + def _get_add_time_ids( + self, + original_size: tuple[int, int], + crops_coords_top_left: tuple[int, int], + target_size: tuple[int, int], + batch_size: int, + ) -> torch.Tensor: + add_time_ids = torch.tensor( + [list(original_size) + list(crops_coords_top_left) + list(target_size)], + dtype=torch.float32, + device=self.device, + ) + return add_time_ids.repeat(batch_size, 1) + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=self.device, dtype=self.od_config.dtype) + + shape = ( + batch_size, + num_channels_latents, + height // self.vae_scale_factor, + width // self.vae_scale_factor, + ) + latents = randn_tensor(shape, generator=generator, device=self.device, dtype=self.od_config.dtype) + latents = latents * self.scheduler.init_noise_sigma + return latents + + def predict_noise(self, **kwargs): + result = self.unet(**kwargs) + return result[0] + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + def diffuse( + self, + latents: torch.Tensor, + timesteps: torch.Tensor, + prompt_embeds: torch.Tensor, + added_cond_kwargs: dict, + negative_prompt_embeds: torch.Tensor | None, + negative_added_cond_kwargs: dict | None, + do_cfg: bool, + guidance_scale: float, + ) -> torch.Tensor: + for _, t in enumerate(timesteps): + if self.interrupt: + continue + self._current_timestep = t + + # Scale model input + latent_model_input = self.scheduler.scale_model_input(latents, t) + timestep = t.expand(latent_model_input.shape[0]).to(device=self.device, dtype=self.od_config.dtype) + + positive_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": prompt_embeds, + "added_cond_kwargs": added_cond_kwargs, + "return_dict": False, + } + if do_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": negative_prompt_embeds, + "added_cond_kwargs": negative_added_cond_kwargs, + "return_dict": False, + } + else: + negative_kwargs = None + + noise_pred = self.predict_noise_maybe_with_cfg( + do_cfg, + guidance_scale, + positive_kwargs, + negative_kwargs, + cfg_normalize=False, + ) + + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_cfg) + + return latents + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] = "", + negative_prompt: str | list[str] = "", + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> DiffusionOutput: + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + negative_prompt = [ + "" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts + ] or negative_prompt + + height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor + width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor + num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + generator = req.sampling_params.generator or generator + num_images_per_prompt = ( + req.sampling_params.num_outputs_per_prompt + if req.sampling_params.num_outputs_per_prompt > 0 + else num_images_per_prompt + ) + + self._guidance_scale = req.sampling_params.guidance_scale + self._current_timestep = None + self._interrupt = False + + if isinstance(prompt, str): + batch_size = 1 + else: + batch_size = len(prompt) + + # Encode prompt + prompt_embeds, pooled_prompt_embeds = self.encode_prompt( + prompt=prompt, + num_images_per_prompt=num_images_per_prompt, + ) + + # Encode negative prompt for CFG + do_cfg = self.guidance_scale > 1.0 + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if do_cfg: + negative_prompt_embeds, negative_pooled_prompt_embeds = self.encode_prompt( + prompt=negative_prompt, + num_images_per_prompt=num_images_per_prompt, + ) + + # Prepare added conditioning + original_size = (height, width) + target_size = (height, width) + crops_coords_top_left = (0, 0) + + add_time_ids = self._get_add_time_ids( + original_size, + crops_coords_top_left, + target_size, + batch_size * num_images_per_prompt, + ) + added_cond_kwargs = { + "text_embeds": pooled_prompt_embeds, + "time_ids": add_time_ids, + } + negative_added_cond_kwargs = None + if do_cfg: + negative_added_cond_kwargs = { + "text_embeds": negative_pooled_prompt_embeds, + "time_ids": add_time_ids, + } + + # XPU workaround: set_timesteps uses numpy internally and needs CPU tensors. + for attr in ("alphas_cumprod", "sigmas", "timesteps"): + if hasattr(self.scheduler, attr): + val = getattr(self.scheduler, attr) + if hasattr(val, "device") and val.device.type != "cpu": + setattr(self.scheduler, attr, val.cpu()) + self.scheduler.set_timesteps(num_inference_steps) + self.scheduler.sigmas = self.scheduler.sigmas.to(self.device) + self.scheduler.timesteps = self.scheduler.timesteps.to(self.device) + timesteps = self.scheduler.timesteps + + # Prepare latents (uses scheduler.init_noise_sigma which depends on set_timesteps) + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + self.unet.in_channels, + height, + width, + generator, + latents, + ) + self._num_timesteps = len(timesteps) + + # Denoising loop + latents = self.diffuse( + latents=latents, + timesteps=timesteps, + prompt_embeds=prompt_embeds, + added_cond_kwargs=added_cond_kwargs, + negative_prompt_embeds=negative_prompt_embeds, + negative_added_cond_kwargs=negative_added_cond_kwargs, + do_cfg=do_cfg, + guidance_scale=self.guidance_scale, + ) + + self._current_timestep = None + + if self.output_type == "latent": + image = latents + else: + latents = latents.to(self.vae.dtype) + latents = latents / self.vae.config.scaling_factor + image = self.vae.decode(latents, return_dict=False)[0] + + return DiffusionOutput( + output=image, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/sdxl/sdxl_unet.py b/vllm_omni/diffusion/models/sdxl/sdxl_unet.py new file mode 100644 index 000000000..76382e645 --- /dev/null +++ b/vllm_omni/diffusion/models/sdxl/sdxl_unet.py @@ -0,0 +1,871 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) + +logger = init_logger(__name__) + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + max_period: int = 10000, +) -> torch.Tensor: + half_dim = embedding_dim // 2 + exponent = ( + -math.log(max_period) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / half_dim + ) + emb = timesteps[:, None].float() * torch.exp(exponent)[None, :] + emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1) + if embedding_dim % 2 == 1: + emb = F.pad(emb, (0, 1)) + return emb + + +class SDXLTimestepEmbedding(nn.Module): + def __init__(self, in_channels: int, time_embed_dim: int): + super().__init__() + self.linear_1 = nn.Linear(in_channels, time_embed_dim) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + sample = self.linear_1(sample) + sample = self.act(sample) + sample = self.linear_2(sample) + return sample + + +class SDXLAddTimestepEmbedding(nn.Module): + def __init__(self, addition_time_embed_dim: int, text_embed_dim: int, time_embed_dim: int): + super().__init__() + self.linear_1 = nn.Linear(text_embed_dim + addition_time_embed_dim * 6, time_embed_dim) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim) + + def forward(self, text_embeds: torch.Tensor, time_ids: torch.Tensor) -> torch.Tensor: + time_embeds = get_timestep_embedding(time_ids.flatten(), 256) + time_embeds = time_embeds.reshape(text_embeds.shape[0], -1) + add_embeds = torch.cat([text_embeds, time_embeds], dim=-1) + add_embeds = add_embeds.to(dtype=self.linear_1.weight.dtype) + add_embeds = self.linear_1(add_embeds) + add_embeds = self.act(add_embeds) + add_embeds = self.linear_2(add_embeds) + return add_embeds + + +class SDXLResnetBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + time_embed_dim: int, + groups: int = 32, + ): + super().__init__() + self.norm1 = nn.GroupNorm(groups, in_channels, eps=1e-5) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + self.time_emb_proj = nn.Linear(time_embed_dim, out_channels) + self.norm2 = nn.GroupNorm(groups, out_channels, eps=1e-5) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) + self.nonlinearity = nn.SiLU() + + self.conv_shortcut = None + if in_channels != out_channels: + self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1) + + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states) + + temb = self.nonlinearity(temb) + temb = self.time_emb_proj(temb) + hidden_states = hidden_states + temb[:, :, None, None] + + hidden_states = self.norm2(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + + return hidden_states + residual + + +class SDXLGEGLU(nn.Module): + def __init__(self, dim_in: int, dim_out: int): + super().__init__() + self.proj = MergedColumnParallelLinear(dim_in, [dim_out, dim_out], bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.proj(hidden_states) + hidden_states, gate = hidden_states.chunk(2, dim=-1) + return hidden_states * F.gelu(gate) + + +class SDXLFeedForward(nn.Module): + def __init__(self, dim: int, inner_dim: int): + super().__init__() + self.geglu = SDXLGEGLU(dim, inner_dim) + self.out_proj = RowParallelLinear(inner_dim, dim, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.geglu(hidden_states) + hidden_states, _ = self.out_proj(hidden_states) + return hidden_states + + +class SDXLSelfAttention(nn.Module): + def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, prefix: str = ""): + super().__init__() + inner_dim = num_attention_heads * attention_head_dim + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=attention_head_dim, + total_num_heads=num_attention_heads, + bias=False, + ) + self.to_out = nn.ModuleList( + [ + RowParallelLinear(inner_dim, dim, bias=True), + ] + ) + self.attention = Attention( + num_heads=num_attention_heads, + head_size=attention_head_dim, + causal=False, + softmax_scale=1.0 / math.sqrt(attention_head_dim), + role="self", + prefix=prefix, + ) + + +class SDXLCrossAttention(nn.Module): + def __init__( + self, dim: int, cross_attention_dim: int, num_attention_heads: int, attention_head_dim: int, prefix: str = "" + ): + super().__init__() + inner_dim = num_attention_heads * attention_head_dim + self.to_q = ColumnParallelLinear(dim, inner_dim, bias=False) + self.to_k = ColumnParallelLinear(cross_attention_dim, inner_dim, bias=False) + self.to_v = ColumnParallelLinear(cross_attention_dim, inner_dim, bias=False) + self.to_out = nn.ModuleList( + [ + RowParallelLinear(inner_dim, dim, bias=True), + ] + ) + self.attention = Attention( + num_heads=num_attention_heads, + head_size=attention_head_dim, + causal=False, + softmax_scale=1.0 / math.sqrt(attention_head_dim), + role="cross", + prefix=prefix, + skip_sequence_parallel=True, + ) + + +class SDXLBasicTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + cross_attention_dim: int, + prefix: str = "", + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + + self.norm1 = nn.LayerNorm(dim, elementwise_affine=True) + self.norm2 = nn.LayerNorm(dim, elementwise_affine=True) + self.norm3 = nn.LayerNorm(dim, elementwise_affine=True) + + self.attn1 = SDXLSelfAttention( + dim=dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + prefix=f"{prefix}.attn1", + ) + self.attn2 = SDXLCrossAttention( + dim=dim, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + prefix=f"{prefix}.attn2", + ) + + # Feed-forward + self.ff = SDXLFeedForward(dim=dim, inner_dim=dim * 4) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: None = None, + ) -> torch.Tensor: + local_num_heads = self.attn1.to_qkv.num_heads + head_dim = self.attention_head_dim + + # Self-attention + norm_hidden = self.norm1(hidden_states) + qkv, _ = self.attn1.to_qkv(norm_hidden) + q, k, v = qkv.chunk(3, dim=-1) + q = q.unflatten(-1, (local_num_heads, head_dim)) + k = k.unflatten(-1, (local_num_heads, head_dim)) + v = v.unflatten(-1, (local_num_heads, head_dim)) + attn_out = self.attn1.attention(q, k, v) + attn_out = attn_out.flatten(-2) + attn_out, _ = self.attn1.to_out[0](attn_out) + hidden_states = hidden_states + attn_out + + # Cross-attention + norm_hidden = self.norm2(hidden_states) + q, _ = self.attn2.to_q(norm_hidden) + k, _ = self.attn2.to_k(encoder_hidden_states) + v, _ = self.attn2.to_v(encoder_hidden_states) + q = q.unflatten(-1, (local_num_heads, head_dim)) + k = k.unflatten(-1, (local_num_heads, head_dim)) + v = v.unflatten(-1, (local_num_heads, head_dim)) + attn_out = self.attn2.attention(q, k, v) + attn_out = attn_out.flatten(-2) + attn_out, _ = self.attn2.to_out[0](attn_out) + hidden_states = hidden_states + attn_out + + # Feed-forward + norm_hidden = self.norm3(hidden_states) + ff_out = self.ff(norm_hidden) + hidden_states = hidden_states + ff_out + + return hidden_states + + +class SDXLTransformer2DModel(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + in_channels: int, + num_layers: int, + cross_attention_dim: int, + prefix: str = "", + ): + super().__init__() + inner_dim = num_attention_heads * attention_head_dim + + self.norm = nn.GroupNorm(32, in_channels, eps=1e-6, affine=True) + self.proj_in = nn.Linear(in_channels, inner_dim) + + self.transformer_blocks = nn.ModuleList( + [ + SDXLBasicTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + cross_attention_dim=cross_attention_dim, + prefix=f"{prefix}.transformer_blocks.{i}", + ) + for i in range(num_layers) + ] + ) + + self.proj_out = nn.Linear(inner_dim, in_channels) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> torch.Tensor: + batch, channels, height, width = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, channels) + hidden_states = self.proj_in(hidden_states) + + for block in self.transformer_blocks: + hidden_states = block(hidden_states, encoder_hidden_states) + + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.reshape(batch, height, width, channels).permute(0, 3, 1, 2) + + return hidden_states + residual + + +class SDXLDownsample2D(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.conv = nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.conv(hidden_states) + + +class SDXLUpsample2D(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest") + return self.conv(hidden_states) + + +class SDXLDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + time_embed_dim: int, + num_layers: int = 2, + add_downsample: bool = True, + ): + super().__init__() + self.resnets = nn.ModuleList() + for i in range(num_layers): + res_in = in_channels if i == 0 else out_channels + self.resnets.append(SDXLResnetBlock2D(res_in, out_channels, time_embed_dim)) + + self.downsamplers = None + if add_downsample: + self.downsamplers = nn.ModuleList([SDXLDownsample2D(out_channels)]) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + output_states = [] + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb) + output_states.append(hidden_states) + + if self.downsamplers is not None: + hidden_states = self.downsamplers[0](hidden_states) + output_states.append(hidden_states) + + return hidden_states, output_states + + +class SDXLCrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + time_embed_dim: int, + num_layers: int = 2, + num_attention_heads: int = 10, + cross_attention_dim: int = 2048, + transformer_layers_per_block: int = 1, + add_downsample: bool = True, + prefix: str = "", + ): + super().__init__() + attention_head_dim = out_channels // num_attention_heads + + self.resnets = nn.ModuleList() + self.attentions = nn.ModuleList() + for i in range(num_layers): + res_in = in_channels if i == 0 else out_channels + self.resnets.append(SDXLResnetBlock2D(res_in, out_channels, time_embed_dim)) + self.attentions.append( + SDXLTransformer2DModel( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + prefix=f"{prefix}.attentions.{i}", + ) + ) + + self.downsamplers = None + if add_downsample: + self.downsamplers = nn.ModuleList([SDXLDownsample2D(out_channels)]) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + output_states = [] + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb) + hidden_states = attn(hidden_states, encoder_hidden_states) + output_states.append(hidden_states) + + if self.downsamplers is not None: + hidden_states = self.downsamplers[0](hidden_states) + output_states.append(hidden_states) + + return hidden_states, output_states + + +class SDXLUNetMidBlock2DCrossAttn(nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + num_attention_heads: int = 20, + cross_attention_dim: int = 2048, + transformer_layers_per_block: int = 10, + prefix: str = "", + ): + super().__init__() + attention_head_dim = in_channels // num_attention_heads + + self.resnets = nn.ModuleList( + [ + SDXLResnetBlock2D(in_channels, in_channels, time_embed_dim), + SDXLResnetBlock2D(in_channels, in_channels, time_embed_dim), + ] + ) + self.attentions = nn.ModuleList( + [ + SDXLTransformer2DModel( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + in_channels=in_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + prefix=f"{prefix}.attentions.0", + ) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.resnets[0](hidden_states, temb) + hidden_states = self.attentions[0](hidden_states, encoder_hidden_states) + hidden_states = self.resnets[1](hidden_states, temb) + return hidden_states + + +class SDXLUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + time_embed_dim: int, + num_layers: int = 3, + add_upsample: bool = True, + resnet_in_channels_list: list[int] | None = None, + ): + super().__init__() + self.resnets = nn.ModuleList() + for i in range(num_layers): + if resnet_in_channels_list is not None: + res_in = resnet_in_channels_list[i] + else: + skip_channels = prev_output_channel if i == 0 else out_channels + res_in = (in_channels if i == 0 else out_channels) + skip_channels + self.resnets.append(SDXLResnetBlock2D(res_in, out_channels, time_embed_dim)) + + self.upsamplers = None + if add_upsample: + self.upsamplers = nn.ModuleList([SDXLUpsample2D(out_channels)]) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + for resnet in self.resnets: + res_hidden = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden], dim=1) + hidden_states = resnet(hidden_states, temb) + + if self.upsamplers is not None: + hidden_states = self.upsamplers[0](hidden_states) + + return hidden_states + + +class SDXLCrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + time_embed_dim: int, + num_layers: int = 3, + num_attention_heads: int = 10, + cross_attention_dim: int = 2048, + transformer_layers_per_block: int = 1, + add_upsample: bool = True, + prefix: str = "", + resnet_in_channels_list: list[int] | None = None, + ): + super().__init__() + attention_head_dim = out_channels // num_attention_heads + + self.resnets = nn.ModuleList() + self.attentions = nn.ModuleList() + for i in range(num_layers): + if resnet_in_channels_list is not None: + res_in = resnet_in_channels_list[i] + else: + skip_channels = prev_output_channel if i == 0 else out_channels + res_in = (in_channels if i == 0 else out_channels) + skip_channels + self.resnets.append(SDXLResnetBlock2D(res_in, out_channels, time_embed_dim)) + self.attentions.append( + SDXLTransformer2DModel( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + prefix=f"{prefix}.attentions.{i}", + ) + ) + + self.upsamplers = None + if add_upsample: + self.upsamplers = nn.ModuleList([SDXLUpsample2D(out_channels)]) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + for resnet, attn in zip(self.resnets, self.attentions): + res_hidden = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden], dim=1) + hidden_states = resnet(hidden_states, temb) + hidden_states = attn(hidden_states, encoder_hidden_states) + + if self.upsamplers is not None: + hidden_states = self.upsamplers[0](hidden_states) + + return hidden_states + + +def _build_sdxl_sp_plan() -> dict: + """Build _sp_plan for all attention modules in the SDXL UNet. + + Shards hidden_states at the first transformer block input and gathers + at proj_out for each SDXLTransformer2DModel instance. + """ + plan: dict = {} + attn_paths = [ + # down_blocks[1]: 2 attention modules + "down_blocks.1.attentions.0", + "down_blocks.1.attentions.1", + # down_blocks[2]: 2 attention modules + "down_blocks.2.attentions.0", + "down_blocks.2.attentions.1", + # mid_block: 1 attention module + "mid_block.attentions.0", + # up_blocks[0]: 3 attention modules + "up_blocks.0.attentions.0", + "up_blocks.0.attentions.1", + "up_blocks.0.attentions.2", + # up_blocks[1]: 3 attention modules + "up_blocks.1.attentions.0", + "up_blocks.1.attentions.1", + "up_blocks.1.attentions.2", + ] + for path in attn_paths: + plan[f"{path}.transformer_blocks.0"] = { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), + } + plan[f"{path}.proj_out"] = SequenceParallelOutput(gather_dim=1, expected_dims=3) + return plan + + +class SDXLUNet2DConditionModel(nn.Module): + _repeated_blocks = ["SDXLBasicTransformerBlock", "SDXLResnetBlock2D"] + _layerwise_offload_blocks_attrs = ["down_blocks", "up_blocks"] + _sp_plan = _build_sdxl_sp_plan() + + @staticmethod + def _is_shardable_block(name: str, module) -> bool: + """Match transformer blocks and ResNet blocks for HSDP sharding.""" + if not name.split(".")[-1].isdigit(): + return False + return "transformer_blocks" in name or "resnets" in name + + _hsdp_shard_conditions = [_is_shardable_block] + + def __init__(self, *, od_config: OmniDiffusionConfig): + super().__init__() + self.od_config = od_config + + # SDXL Base configuration + model_channels = 320 + time_embed_dim = model_channels * 4 # 1280 + block_out_channels = [320, 640, 1280] + layers_per_block = 2 + transformer_layers_per_block = [1, 2, 10] + cross_attention_dim = 2048 + addition_time_embed_dim = 256 + num_head_channels = 64 + + self.in_channels = 4 + self.out_channels = 4 + + # Input convolution + self.conv_in = nn.Conv2d(4, model_channels, kernel_size=3, padding=1) + + # Time embedding + self.time_proj_dim = model_channels + self.time_embedding = SDXLTimestepEmbedding(model_channels, time_embed_dim) + + # Addition embedding (for pooled text + crop/size conditioning) + self.add_embedding = SDXLAddTimestepEmbedding( + addition_time_embed_dim=addition_time_embed_dim, + text_embed_dim=1280, + time_embed_dim=time_embed_dim, + ) + + # Down blocks + self.down_blocks = nn.ModuleList() + + # Level 0: DownBlock2D (no attention) + self.down_blocks.append( + SDXLDownBlock2D( + in_channels=model_channels, + out_channels=block_out_channels[0], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block, + add_downsample=True, + ) + ) + # Level 1: CrossAttnDownBlock2D + self.down_blocks.append( + SDXLCrossAttnDownBlock2D( + in_channels=block_out_channels[0], + out_channels=block_out_channels[1], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block, + num_attention_heads=block_out_channels[1] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[1], + add_downsample=True, + prefix="down_blocks.1", + ) + ) + # Level 2: CrossAttnDownBlock2D (no downsampler) + self.down_blocks.append( + SDXLCrossAttnDownBlock2D( + in_channels=block_out_channels[1], + out_channels=block_out_channels[2], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block, + num_attention_heads=block_out_channels[2] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[2], + add_downsample=False, + prefix="down_blocks.2", + ) + ) + + # Mid block + self.mid_block = SDXLUNetMidBlock2DCrossAttn( + in_channels=block_out_channels[2], + time_embed_dim=time_embed_dim, + num_attention_heads=block_out_channels[2] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[2], + prefix="mid_block", + ) + + # Up blocks (reversed, with skip connections) + # Skip connection channels consumed from down_block_res_samples (right to left): + # down_block_res_samples channels = [320, 320, 320, 320, 640, 640, 640, 1280, 1280] + # up_blocks[0] consumes last 3: skips are [1280, 1280, 640] (consumed from right) + # up_blocks[1] consumes next 3: skips are [640, 640, 320] (consumed from right) + # up_blocks[2] consumes next 3: skips are [320, 320, 320] (consumed from right) + self.up_blocks = nn.ModuleList() + + # Level 0: CrossAttnUpBlock2D (mirrors down_blocks[2]) + # Input from mid: 1280, skips (consumed right-to-left): 1280, 1280, 640 + self.up_blocks.append( + SDXLCrossAttnUpBlock2D( + in_channels=block_out_channels[2], + prev_output_channel=block_out_channels[2], + out_channels=block_out_channels[2], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block + 1, + num_attention_heads=block_out_channels[2] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[2], + add_upsample=True, + prefix="up_blocks.0", + resnet_in_channels_list=[1280 + 1280, 1280 + 1280, 1280 + 640], + ) + ) + # Level 1: CrossAttnUpBlock2D (mirrors down_blocks[1]) + # Input from up_blocks[0]: 1280, skips (consumed right-to-left): 640, 640, 320 + self.up_blocks.append( + SDXLCrossAttnUpBlock2D( + in_channels=block_out_channels[2], + prev_output_channel=block_out_channels[1], + out_channels=block_out_channels[1], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block + 1, + num_attention_heads=block_out_channels[1] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[1], + add_upsample=True, + prefix="up_blocks.1", + resnet_in_channels_list=[1280 + 640, 640 + 640, 640 + 320], + ) + ) + # Level 2: UpBlock2D (mirrors down_blocks[0], no attention, no upsample) + # Input from up_blocks[1]: 640, skips (consumed right-to-left): 320, 320, 320 + self.up_blocks.append( + SDXLUpBlock2D( + in_channels=block_out_channels[1], + prev_output_channel=block_out_channels[0], + out_channels=block_out_channels[0], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block + 1, + add_upsample=False, + resnet_in_channels_list=[640 + 320, 320 + 320, 320 + 320], + ) + ) + + # Output + self.conv_norm_out = nn.GroupNorm(32, block_out_channels[0], eps=1e-5) + self.conv_act = nn.SiLU() + self.conv_out = nn.Conv2d(block_out_channels[0], 4, kernel_size=3, padding=1) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + added_cond_kwargs: dict, + return_dict: bool = False, + ) -> tuple[torch.Tensor]: + # Time embedding + dtype = hidden_states.dtype + t_emb = get_timestep_embedding(timestep, self.time_proj_dim).to(dtype=dtype) + emb = self.time_embedding(t_emb) + + # Addition embedding (pooled text + size/crop) + text_embeds = added_cond_kwargs["text_embeds"].to(dtype=dtype) + time_ids = added_cond_kwargs["time_ids"] + add_emb = self.add_embedding(text_embeds, time_ids) + emb = emb + add_emb + + # Input conv + hidden_states = self.conv_in(hidden_states) + + # Down blocks + down_block_res_samples = [hidden_states] + for down_block in self.down_blocks: + if isinstance(down_block, SDXLCrossAttnDownBlock2D): + hidden_states, res_samples = down_block(hidden_states, emb, encoder_hidden_states) + else: + hidden_states, res_samples = down_block(hidden_states, emb) + down_block_res_samples.extend(res_samples) + + # Mid block + hidden_states = self.mid_block(hidden_states, emb, encoder_hidden_states) + + # Up blocks + for up_block in self.up_blocks: + # Determine how many skip connections this block needs + n_resnets = len(up_block.resnets) + res_samples = tuple(down_block_res_samples[-n_resnets:]) + down_block_res_samples = down_block_res_samples[:-n_resnets] + + if isinstance(up_block, SDXLCrossAttnUpBlock2D): + hidden_states = up_block(hidden_states, emb, encoder_hidden_states, res_samples) + else: + hidden_states = up_block(hidden_states, emb, res_samples) + + # Output + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return (hidden_states,) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # Self-attention QKV fusion + (".attn1.to_qkv", ".attn1.to_q", "q"), + (".attn1.to_qkv", ".attn1.to_k", "k"), + (".attn1.to_qkv", ".attn1.to_v", "v"), + ] + self.stacked_params_mapping = stacked_params_mapping + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + # Handle GEGLU: diffusers stores as ff.net.0.proj.weight/bias + # MergedColumnParallelLinear needs the two halves loaded separately + if ".ff.net.0.proj." in name: + name = name.replace(".ff.net.0.proj.", ".ff.geglu.proj.") + if name in params_dict: + param = params_dict[name] + weight_loader = param.weight_loader + half = loaded_weight.shape[0] // 2 + weight_loader(param, loaded_weight[:half], 0) + weight_loader(param, loaded_weight[half:], 1) + loaded_params.add(name) + continue + elif ".ff.net.2." in name: + name = name.replace(".ff.net.2.", ".ff.out_proj.") + + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if name not in params_dict: + break + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + handled = True + break + + if not handled: + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 9cc65486a..996e74b32 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -286,6 +286,11 @@ _DIFFUSION_MODELS = { "pipeline_dreamzero", "DreamZeroPipeline", ), + "StableDiffusionXLPipeline": ( + "sdxl", + "pipeline_sdxl", + "StableDiffusionXLPipeline", + ), } @@ -510,6 +515,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func", "HiDreamImagePipeline": "get_hidream_image_post_process_func", + "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func", } _DIFFUSION_ACTION_POST_PROCESS_FUNCS = {