diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 391cab0b1..8b2cf6aa0 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -21,6 +21,7 @@ th { | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | +| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | | `WanPipeline` | Wan2.2 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 220b930d6..819b0918e 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -31,11 +31,13 @@ vLLM-Omni also supports the sequence parallelism (SP) for the diffusion model, t The following table shows which models are currently supported by each acceleration method: + | Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | |-------|-----------------|----------|-----------|-----------| | **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ✅ |❌ | | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ❌ | ✅ |✅ | +| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ❌ | ✅ |✅ | ## Performance Benchmarks diff --git a/docs/user_guide/examples/offline_inference/image_to_image.md b/docs/user_guide/examples/offline_inference/image_to_image.md index dbbcdaf8f..b0373e002 100644 --- a/docs/user_guide/examples/offline_inference/image_to_image.md +++ b/docs/user_guide/examples/offline_inference/image_to_image.md @@ -7,20 +7,40 @@ This example edits an input image with `Qwen/Qwen-Image-Edit` using the `image_e ## Local CLI Usage +### Single Image Editing + ```bash python image_edit.py \ --image qwen_bear.png \ --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \ --output output_image_edit.png \ --num_inference_steps 50 \ - --cfg_scale 4.0 + --cfg_scale 4.0 \ + --guidance_scale 1.0 +``` + +### Multiple Image Editing (Qwen-Image-Edit-2509+) + +For multiple image inputs, use `Qwen/Qwen-Image-Edit-2509` or later version: + +```bash +python image_edit.py \ + --model Qwen/Qwen-Image-Edit-2509 \ + --image img1.png img2.png img3.png \ + --prompt "Combine these images into a single scene" \ + --output output_image_edit.png \ + --num_inference_steps 50 \ + --cfg_scale 4.0 \ + --guidance_scale 1.0 ``` Key arguments: -- `--image`: path to the source image (PNG/JPG, converted to RGB). +- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support. +- `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images. - `--prompt` / `--negative_prompt`: text description (string). -- `--cfg_scale`: true CFG scale for Qwen-Image-Edit (quality vs. fidelity). +- `--cfg_scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality. +- `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models. - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower). - `--output`: path to save the generated PNG. diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 9547a3cd0..212809ad9 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -4,13 +4,23 @@ """ Example script for image editing with Qwen-Image-Edit. -Usage: +Usage (single image): python image_edit.py \ --image input.png \ --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \ --output output_image_edit.png \ --num_inference_steps 50 \ - --cfg_scale 4.0 + --cfg_scale 4.0 \ + --guidance_scale 1.0 + +Usage (multiple images): + python image_edit.py \ + --image input1.png input2.png input3.png \ + --prompt "Combine these images into a single scene" \ + --output output_image_edit.png \ + --num_inference_steps 50 \ + --cfg_scale 4.0 \ + --guidance_scale 1.0 For more options, run: python image_edit.py --help @@ -34,13 +44,18 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--model", default="Qwen/Qwen-Image-Edit", - help="Diffusion model name or local path.", + help=( + "Diffusion model name or local path. " + "For multiple image inputs, use Qwen/Qwen-Image-Edit-2509 or later version " + "which supports QwenImageEditPlusPipeline." + ), ) parser.add_argument( "--image", type=str, + nargs="+", required=True, - help="Path to input image file (PNG, JPG, etc.).", + help="Path(s) to input image file(s) (PNG, JPG, etc.). Can specify multiple images.", ) parser.add_argument( "--prompt", @@ -64,7 +79,22 @@ def parse_args() -> argparse.Namespace: "--cfg_scale", type=float, default=4.0, - help="True classifier-free guidance scale specific to Qwen-Image-Edit.", + help=( + "True classifier-free guidance scale (default: 4.0). Guidance scale as defined in Classifier-Free " + "Diffusion Guidance. Classifier-free guidance is enabled by setting cfg_scale > 1 and providing " + "a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, " + "usually at the expense of lower image quality." + ), + ) + parser.add_argument( + "--guidance_scale", + type=float, + default=1.0, + help=( + "Guidance scale for guidance-distilled models (default: 1.0, disabled). " + "Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale " + "directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models." + ), ) parser.add_argument( "--output", @@ -108,13 +138,19 @@ def parse_args() -> argparse.Namespace: def main(): args = parse_args() - # Validate input image exists - if not os.path.exists(args.image): - raise FileNotFoundError(f"Input image not found: {args.image}") + # Validate input images exist and load them + input_images = [] + for image_path in args.image: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Input image not found: {image_path}") + img = Image.open(image_path).convert("RGB") + input_images.append(img) - # Load input image - input_image = Image.open(args.image).convert("RGB") - print(f"Loaded input image from {args.image} (size: {input_image.size})") + # Use single image or list based on number of inputs + if len(input_images) == 1: + input_image = input_images[0] + else: + input_image = input_images device = detect_device_type() generator = torch.Generator(device=device).manual_seed(args.seed) @@ -153,10 +189,9 @@ def main(): # (e.g., QwenImagePipeline or FluxPipeline) } - # Initialize Omni with QwenImageEditPipeline + # Initialize Omni with appropriate pipeline omni = Omni( model=args.model, - model_class_name="QwenImageEditPipeline", vae_use_slicing=vae_use_slicing, vae_use_tiling=vae_use_tiling, cache_backend=args.cache_backend, @@ -171,6 +206,12 @@ def main(): print(f" Model: {args.model}") print(f" Inference steps: {args.num_inference_steps}") print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}") + if isinstance(input_image, list): + print(f" Number of input images: {len(input_image)}") + for idx, img in enumerate(input_image): + print(f" Image {idx + 1} size: {img.size}") + else: + print(f" Input image size: {input_image.size}") print(f" Parallel configuration: ulysses_degree={args.ulysses_degree}") print(f" Input image size: {input_image.size}") print(f"{'=' * 60}\n") @@ -183,6 +224,7 @@ def main(): negative_prompt=args.negative_prompt, generator=generator, true_cfg_scale=args.cfg_scale, + guidance_scale=args.guidance_scale, num_inference_steps=args.num_inference_steps, num_outputs_per_prompt=args.num_outputs_per_prompt, ) diff --git a/examples/offline_inference/image_to_image/image_to_image.md b/examples/offline_inference/image_to_image/image_to_image.md index 2bd6d84e1..b2c40dd85 100644 --- a/examples/offline_inference/image_to_image/image_to_image.md +++ b/examples/offline_inference/image_to_image/image_to_image.md @@ -4,19 +4,38 @@ This example edits an input image with `Qwen/Qwen-Image-Edit` using the `image_e ## Local CLI Usage +### Single Image Editing + ```bash python image_edit.py \ --image qwen_bear.png \ --prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \ --output output_image_edit.png \ --num_inference_steps 50 \ - --cfg_scale 4.0 + --cfg_scale 4.0 \ +``` + +### Multiple Image Editing (Qwen-Image-Edit-2509+) + +For multiple image inputs, use `Qwen/Qwen-Image-Edit-2509` or later version: + +```bash +python image_edit.py \ + --model Qwen/Qwen-Image-Edit-2509 \ + --image img1.png img2.png \ + --prompt "Combine these images into a single scene" \ + --output output_image_edit.png \ + --num_inference_steps 50 \ + --cfg_scale 4.0 \ + --guidance_scale 1.0 ``` Key arguments: -- `--image`: path to the source image (PNG/JPG, converted to RGB). +- `--model`: model name or path. Use `Qwen/Qwen-Image-Edit-2509` or later for multiple image support. +- `--image`: path(s) to the source image(s) (PNG/JPG, converted to RGB). Can specify multiple images. - `--prompt` / `--negative_prompt`: text description (string). -- `--cfg_scale`: true CFG scale for Qwen-Image-Edit (quality vs. fidelity). +- `--cfg_scale`: true classifier-free guidance scale (default: 4.0). Classifier-free guidance is enabled by setting cfg_scale > 1 and providing a negative_prompt. Higher guidance scale encourages images closely linked to the text prompt, usually at the expense of lower image quality. +- `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models. - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower). - `--output`: path to save the generated PNG. diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py new file mode 100644 index 000000000..d3b9bfe94 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -0,0 +1,793 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import logging +import os +from collections.abc import Iterable +from typing import Any + +import numpy as np +import PIL.Image +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.models.autoencoders.autoencoder_kl_qwenimage import ( + AutoencoderKLQwenImage, +) +from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import calculate_shift +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import ( + calculate_dimensions, + retrieve_latents, + retrieve_timesteps, +) +from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( + QwenImageTransformer2DModel, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, +) + +logger = logging.getLogger(__name__) + +CONDITION_IMAGE_SIZE = 384 * 384 +VAE_IMAGE_SIZE = 1024 * 1024 + + +def get_qwen_image_edit_plus_pre_process_func( + od_config: OmniDiffusionConfig, +): + """Pre-processing function for QwenImageEditPlusPipeline.""" + model_name = od_config.model + if os.path.exists(model_name): + model_path = model_name + else: + model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + vae_config_path = os.path.join(model_path, "vae/config.json") + with open(vae_config_path) as f: + vae_config = json.load(f) + vae_scale_factor = 2 ** len(vae_config["temporal_downsample"]) if "temporal_downsample" in vae_config else 8 + + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) + latent_channels = vae_config.get("z_dim", 16) + + def pre_process_func( + requests: list[OmniDiffusionRequest], + ): + """Pre-process requests for QwenImageEditPlusPipeline.""" + for req in requests: + image = req.pil_image + + # Handle single image or list of images + if image is None: + continue + + if not isinstance(image, list): + image = [image] + + # Calculate dimensions based on first image + image_size = image[0].size + calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1]) + height = req.height or calculated_height + width = req.width or calculated_width + + # Ensure dimensions are multiples of vae_scale_factor * 2 + multiple_of = vae_scale_factor * 2 + height = height // multiple_of * multiple_of + width = width // multiple_of * multiple_of + + # Store calculated dimensions in request + req.calculated_height = calculated_height + req.calculated_width = calculated_width + req.height = height + req.width = width + + # Preprocess images into condition_images (for prompt encoding) and vae_images (for VAE encoding) + condition_images = [] + vae_images = [] + condition_image_sizes = [] + vae_image_sizes = [] + + for img in image: + if isinstance(img, torch.Tensor) and len(img.shape) > 1 and img.shape[1] == latent_channels: + # Already a latent tensor + continue + + image_width, image_height = img.size + condition_width, condition_height = calculate_dimensions( + CONDITION_IMAGE_SIZE, image_width / image_height + ) + vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height) + + condition_image_sizes.append((condition_width, condition_height)) + vae_image_sizes.append((vae_width, vae_height)) + + condition_images.append(image_processor.resize(img, condition_height, condition_width)) + vae_images.append(image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2)) + + # Store preprocessed images in request + req.condition_images = condition_images + req.vae_images = vae_images + req.condition_image_sizes = condition_image_sizes + req.vae_image_sizes = vae_image_sizes + + return requests + + return pre_process_func + + +def get_qwen_image_edit_plus_post_process_func( + od_config: OmniDiffusionConfig, +): + """Post-processing function for QwenImageEditPlusPipeline.""" + model_name = od_config.model + if os.path.exists(model_name): + model_path = model_name + else: + model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + vae_config_path = os.path.join(model_path, "vae/config.json") + with open(vae_config_path) as f: + vae_config = json.load(f) + vae_scale_factor = 2 ** len(vae_config["temporal_downsample"]) if "temporal_downsample" in vae_config else 8 + + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) + + def post_process_func( + images: torch.Tensor, + ): + return image_processor.postprocess(images) + + return post_process_func + + +class QwenImageEditPlusPipeline( + nn.Module, +): + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + self.device = get_local_device() + model = od_config.model + + # Check if model is a local path + local_files_only = os.path.exists(model) + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( + model, subfolder="text_encoder", local_files_only=local_files_only + ) + + self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( + self.device + ) + self.transformer = QwenImageTransformer2DModel(od_config=od_config) + self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.processor = Qwen2VLProcessor.from_pretrained( + model, subfolder="processor", local_files_only=local_files_only + ) + + self.stage = None + + self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 + self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.tokenizer_max_length = 1024 + # Edit prompt template - different from generation template, supports multiple images + self.prompt_template_encode = ( + "<|im_start|>system\nDescribe the key features of the input image " + "(color, shape, size, texture, objects, background), then explain how the user's " + "text instruction should alter or modify the image. Generate a new image that meets " + "the user's requirements while maintaining consistency with the original input where " + "appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" + ) + self.prompt_template_encode_start_idx = 64 + self.default_sample_size = 128 + + def check_inputs( + self, + prompt, + height, + width, + image=None, + negative_prompt=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_embeds_mask=None, + negative_prompt_embeds_mask=None, + callback_on_step_end_tensor_inputs=None, + max_sequence_length=None, + ): + if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0: + logger.warning( + f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} " + f"but are {height} and {width}. Dimensions will be resized accordingly" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if negative_prompt is not None and negative_prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:" + f" {negative_prompt_embeds}. Make sure to only forward one of the two." + ) + + if prompt_embeds is not None and prompt_embeds_mask is None: + raise ValueError( + "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. " + "Make sure to generate `prompt_embeds_mask` from the same text encoder " + "that was used to generate `prompt_embeds`." + ) + if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None: + raise ValueError( + "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. " + "Make sure to generate `negative_prompt_embeds_mask` from the same text encoder " + "that was used to generate `negative_prompt_embeds`." + ) + + if max_sequence_length is not None and max_sequence_length > 1024: + raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + + def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_result = torch.split(selected, valid_lengths.tolist(), dim=0) + + return split_result + + def _get_qwen_prompt_embeds( + self, + prompt: str | list[str] = None, + image: list[torch.Tensor] | torch.Tensor | None = None, + dtype: torch.dtype | None = None, + ): + """Get prompt embeddings with support for multiple images.""" + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + + # Build image prompt template for multiple images + img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>" + if isinstance(image, list): + base_img_prompt = "" + for i, img in enumerate(image): + base_img_prompt += img_prompt_template.format(i + 1) + elif image is not None: + base_img_prompt = img_prompt_template.format(1) + else: + base_img_prompt = "" + + template = self.prompt_template_encode + drop_idx = self.prompt_template_encode_start_idx + txt = [template.format(base_img_prompt + e) for e in prompt] + + # Use processor to handle both text and image inputs + model_inputs = self.processor( + text=txt, + images=image, + padding=True, + return_tensors="pt", + ).to(self.device) + + outputs = self.text_encoder( + input_ids=model_inputs.input_ids, + attention_mask=model_inputs.attention_mask, + pixel_values=model_inputs.pixel_values, + image_grid_thw=model_inputs.image_grid_thw, + output_hidden_states=True, + ) + + hidden_states = outputs.hidden_states[-1] + split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states] + max_seq_len = max([e.size(0) for e in split_hidden_states]) + prompt_embeds = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states] + ) + encoder_attention_mask = torch.stack( + [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list] + ) + + prompt_embeds = prompt_embeds.to(dtype=dtype) + + return prompt_embeds, encoder_attention_mask + + def encode_prompt( + self, + prompt: str | list[str], + image: list[torch.Tensor] | torch.Tensor | None = None, + num_images_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + max_sequence_length: int = 1024, + ): + r""" + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + image (`torch.Tensor` or `list[torch.Tensor]`, *optional*): + image(s) to be encoded + num_images_per_prompt (`int`): + number of images that should be generated per prompt + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + """ + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image) + + _, 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) + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + return prompt_embeds, prompt_embeds_mask + + @staticmethod + def _pack_latents(latents, batch_size, num_channels_latents, height, width): + latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4) + + return latents + + @staticmethod + def _unpack_latents(latents, height, width, vae_scale_factor): + batch_size, num_patches, channels = latents.shape + + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (vae_scale_factor * 2)) + width = 2 * (int(width) // (vae_scale_factor * 2)) + + latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + + latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width) + + return latents + + def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator): + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax") + for i in range(image.shape[0]) + ] + image_latents = torch.cat(image_latents, dim=0) + else: + image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax") + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.latent_channels, 1, 1, 1) + .to(image_latents.device, image_latents.dtype) + ) + image_latents = (image_latents - latents_mean) / latents_std + + return image_latents + + def prepare_latents( + self, + images, + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents=None, + ): + # VAE applies 8x compression on images but we must also account for packing which requires + # latent height and width to be divisible by 2. + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + + shape = (batch_size, 1, num_channels_latents, height, width) + + image_latents = None + if images is not None: + if not isinstance(images, list): + images = [images] + all_image_latents = [] + for image in images: + image = image.to(device=device, dtype=dtype) + if image.shape[1] != self.latent_channels: + image_latents = self._encode_vae_image(image=image, generator=generator) + else: + image_latents = image + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + # expand init_latents for batch_size + additional_image_per_prompt = batch_size // image_latents.shape[0] + image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0) + elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0: + raise ValueError( + f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts." + ) + else: + image_latents = torch.cat([image_latents], dim=0) + + image_latent_height, image_latent_width = image_latents.shape[3:] + image_latents = self._pack_latents( + image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width + ) + all_image_latents.append(image_latents) + # Concatenate all image latents along dimension 1 + image_latents = torch.cat(all_image_latents, dim=1) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + if latents is None: + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width) + else: + latents = latents.to(device=device, dtype=dtype) + + return latents, image_latents + + def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len): + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + sigmas=sigmas, + mu=mu, + ) + return timesteps, num_inference_steps + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + def diffuse( + self, + prompt_embeds, + prompt_embeds_mask, + negative_prompt_embeds, + negative_prompt_embeds_mask, + latents, + image_latents, + img_shapes, + txt_seq_lens, + negative_txt_seq_lens, + timesteps, + do_true_cfg, + guidance, + true_cfg_scale, + ): + """Diffusion loop with optional image conditioning.""" + self.scheduler.set_begin_index(0) + for i, t in enumerate(timesteps): + if self.interrupt: + continue + self._current_timestep = t + # broadcast to batch dimension and place on same device/dtype as latents + timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) + + # Concatenate image latents with noise latents if available + latent_model_input = latents + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1) + + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + encoder_hidden_states_mask=prompt_embeds_mask, + encoder_hidden_states=prompt_embeds, + img_shapes=img_shapes, + txt_seq_lens=txt_seq_lens, + attention_kwargs=self.attention_kwargs, + return_dict=False, + )[0] + noise_pred = noise_pred[:, : latents.size(1)] + + if do_true_cfg: + neg_noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep / 1000, + guidance=guidance, + encoder_hidden_states_mask=negative_prompt_embeds_mask, + encoder_hidden_states=negative_prompt_embeds, + img_shapes=img_shapes, + txt_seq_lens=negative_txt_seq_lens, + attention_kwargs=self.attention_kwargs, + return_dict=False, + )[0] + neg_noise_pred = neg_noise_pred[:, : latents.size(1)] + comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) + cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True) + noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True) + noise_pred = comb_pred * (cond_norm / noise_norm) + # compute the previous noisy sample x_t -> x_t-1 + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + return latents + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] = "", + negative_prompt: str | list[str] = "", + image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor | None = None, + true_cfg_scale: float = 4.0, + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + sigmas: list[float] | None = None, + guidance_scale: float = 1.0, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds_mask: torch.Tensor | None = None, + output_type: str | None = "pil", + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 512, + ) -> DiffusionOutput: + """Forward pass for image editing with support for multiple images.""" + prompt = req.prompt if req.prompt is not None else prompt + negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt + + # Get preprocessed images from request (pre-processing is done in DiffusionEngine) + if hasattr(req, "vae_images") and hasattr(req, "condition_images"): + condition_images = req.condition_images + vae_images = req.vae_images + condition_image_sizes = req.condition_image_sizes + vae_image_sizes = req.vae_image_sizes + calculated_height = req.calculated_height + calculated_width = req.calculated_width + height = req.height + width = req.width + else: + # fallback to run pre-processing in pipeline (debug only) + if image is None: + raise ValueError("Image is required for QwenImageEditPlusPipeline") + + if not isinstance(image, list): + image = [image] + + image_size = image[0].size + calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1]) + height = height or calculated_height + width = width or calculated_width + + multiple_of = self.vae_scale_factor * 2 + width = width // multiple_of * multiple_of + height = height // multiple_of * multiple_of + + condition_images = [] + vae_images = [] + condition_image_sizes = [] + vae_image_sizes = [] + + for img in image: + image_width, image_height = img.size + condition_width, condition_height = calculate_dimensions( + CONDITION_IMAGE_SIZE, image_width / image_height + ) + vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height) + condition_image_sizes.append((condition_width, condition_height)) + vae_image_sizes.append((vae_width, vae_height)) + condition_images.append(self.image_processor.resize(img, condition_height, condition_width)) + vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2)) + + num_inference_steps = req.num_inference_steps or num_inference_steps + generator = req.generator or generator + true_cfg_scale = req.true_cfg_scale or true_cfg_scale + req_num_outputs = getattr(req, "num_outputs_per_prompt", None) + if req_num_outputs and req_num_outputs > 0: + num_images_per_prompt = req_num_outputs + + # 1. check inputs + # 2. encode prompts + # 3. prepare latents and timesteps + # 4. diffusion process + # 5. decode latents + # 6. post-process outputs + self.check_inputs( + prompt, + height, + width, + image, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + prompt_embeds_mask, + negative_prompt_embeds_mask, + callback_on_step_end_tensor_inputs, + max_sequence_length, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._current_timestep = None + self._interrupt = False + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + has_neg_prompt = negative_prompt is not None or ( + negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None + ) + + do_true_cfg = true_cfg_scale > 1 and has_neg_prompt + prompt_embeds, prompt_embeds_mask = self.encode_prompt( + prompt=prompt, + image=condition_images, # Use condition images for prompt encoding + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + + if do_true_cfg: + negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( + prompt=negative_prompt, + image=condition_images, # Use same condition images for negative prompt encoding + prompt_embeds=negative_prompt_embeds, + prompt_embeds_mask=negative_prompt_embeds_mask, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + + num_channels_latents = self.transformer.in_channels // 4 + # random noise latents, and image latents encoded by vae + latents, image_latents = self.prepare_latents( + vae_images, # Use VAE images for latent preparation + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + self.device, + generator, + latents, + ) + # img_shapes includes shapes for output image and all input images + img_shapes = [ + [ + (1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2), + *[ + (1, vae_height // self.vae_scale_factor // 2, vae_width // self.vae_scale_factor // 2) + for vae_width, vae_height in vae_image_sizes + ], + ] + ] * batch_size + + timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1]) + self._num_timesteps = len(timesteps) + + # handle guidance + if self.transformer.guidance_embeds: + guidance = torch.full([1], guidance_scale, dtype=torch.float32) + guidance = guidance.expand(latents.shape[0]) + else: + guidance = None + + if self.attention_kwargs is None: + self._attention_kwargs = {} + + txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None + negative_txt_seq_lens = ( + negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None + ) + + latents = self.diffuse( + prompt_embeds, + prompt_embeds_mask, + negative_prompt_embeds, + negative_prompt_embeds_mask, + latents, + image_latents, + img_shapes, + txt_seq_lens, + negative_txt_seq_lens, + timesteps, + do_true_cfg, + guidance, + true_cfg_scale, + ) + + self._current_timestep = None + if output_type == "latent": + image = latents + else: + latents = self._unpack_latents(latents, height, width, self.vae_scale_factor) + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] + + return DiffusionOutput(output=image) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index f47333209..9f9a18e86 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -19,6 +19,11 @@ _DIFFUSION_MODELS = { "pipeline_qwen_image_edit", "QwenImageEditPipeline", ), + "QwenImageEditPlusPipeline": ( + "qwen_image", + "pipeline_qwen_image_edit_plus", + "QwenImageEditPlusPipeline", + ), "ZImagePipeline": ( "z_image", "pipeline_z_image", @@ -65,6 +70,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key "QwenImagePipeline": "get_qwen_image_post_process_func", "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func", + "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func", "ZImagePipeline": "get_post_process_func", "WanPipeline": "get_wan22_post_process_func", } @@ -74,6 +80,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { # `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py, # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func", + "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func", }