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 2df248e03..1c1a5ff3a 100644 --- a/examples/offline_inference/image_to_image/image_to_image.md +++ b/examples/offline_inference/image_to_image/image_to_image.md @@ -51,5 +51,6 @@ Key arguments: - `--vae-use-tiling`: enable VAE tiling for memory optimization. - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel). - `--enable-cpu-offload`: enable CPU offloading for diffusion models. +- `--strength`: **Z-Image only** - controls the denoising start timestep for I2I (default: 0.6). Range: [0.0, 1.0]. Lower values preserve more of the original image; higher values allow more creative changes. > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage. diff --git a/examples/online_serving/image_to_image/README.md b/examples/online_serving/image_to_image/README.md index 789258473..59b1f0e2c 100644 --- a/examples/online_serving/image_to_image/README.md +++ b/examples/online_serving/image_to_image/README.md @@ -314,6 +314,7 @@ count, use `size` and `n` rather than `height`, `width`, or | `seed` | int | None | Random seed (reproducible) | | `negative_prompt` | str | None | Negative prompt | | `num_outputs_per_prompt` | int | 1 | Number of images to generate | +| `strength` | float | 0.6 | **Z-Image only** - Denoising start timestep for I2I. Range: [0.0, 1.0]. Lower preserves more of original image. | | `layers` | int | 4 | Number of layers (Qwen-Image-Layered) | | `resolution` | int | 640 | Resolution, 640 or 1024 (Qwen-Image-Layered) | diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index b9aceed2e..5bea59a20 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -21,9 +21,10 @@ import os from collections.abc import Callable, Iterable from typing import Any +import PIL.Image import torch import torch.nn as nn -from diffusers.image_processor import VaeImageProcessor +from diffusers.image_processor import PipelineImageInput, VaeImageProcessor from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils import logging from diffusers.utils.torch_utils import randn_tensor @@ -59,7 +60,7 @@ def get_post_process_func( vae_config = json.load(f) vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8 - image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2, do_convert_rgb=True) def post_process_func( images: torch.Tensor, @@ -83,6 +84,20 @@ def calculate_shift( return mu +# Copied from diffusers +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, @@ -187,6 +202,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2, do_convert_rgb=True) + def encode_prompt( self, prompt: str | list[str], @@ -282,12 +299,45 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): device, generator, latents=None, + image=None, + timestep=None, ): height = 2 * (int(height) // (self.vae_scale_factor * 2)) width = 2 * (int(width) // (self.vae_scale_factor * 2)) shape = (batch_size, num_channels_latents, height, width) + if image is not None: + if latents is not None: + return latents.to(device=device, dtype=dtype) + + image = image.to(device=device, dtype=dtype) + if image.shape[1] != num_channels_latents: + if isinstance(generator, list): + image_latents = [ + retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i]) + 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) + + image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor + else: + image_latents = image + + if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0: + 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." + ) + + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self.scheduler.scale_noise(image_latents, timestep, noise) + return latents + if latents is None: latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) else: @@ -296,6 +346,14 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): latents = latents.to(device) return latents + def get_timesteps(self, num_inference_steps, strength, device): + init_timestep = min(num_inference_steps * strength, num_inference_steps) + t_start = int(max(num_inference_steps - init_timestep, 0)) + timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] + if hasattr(self.scheduler, "set_begin_index"): + self.scheduler.set_begin_index(t_start * self.scheduler.order) + return timesteps, num_inference_steps - t_start + @property def guidance_scale(self): return self._guidance_scale @@ -320,6 +378,8 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): self, req: OmniDiffusionRequest, prompt: str | list[str] | None = None, + image: PipelineImageInput = None, + strength: float = 0.6, height: int = 1024, width: int = 1024, num_inference_steps: int = 50, @@ -347,6 +407,11 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): prompt (`str` or `list[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. + image (`PipelineImageInput`, *optional*): + The image to use for img2img generation. If provided, the pipeline + will perform img2img instead of text-to-image. + strength (`float`, *optional*, defaults to 0.6): + Indicates extent to transform the reference `image`. Must be between 0 and 1. height (`int`, *optional*, defaults to 1024): The height in pixels of the generated image. width (`int`, *optional*, defaults to 1024): @@ -425,6 +490,34 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] + # Handle img2img: extract image from request + if image is None and req.prompts: + if len(req.prompts) > 1: + logger.warning( + "This model only supports a single prompt for img2img, not a batched request. " + "Taking only the first image for now." + ) + first_prompt = req.prompts[0] + if not isinstance(first_prompt, str): + raw_image = first_prompt.get("multi_modal_data", {}).get("image") + if raw_image is not None: + if isinstance(raw_image, list): + image = [PIL.Image.open(im) if isinstance(im, str) else raw_image[0] for im in raw_image[:1]] + else: + image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else raw_image + + # strength is currently only applicable for Z-Image I2I; other pipelines ignore this parameter + strength = req.sampling_params.strength if req.sampling_params.strength is not None else strength + if strength is not None and image is None: + logger.warning( + "strength parameter (%.2f) is only applicable for image-to-image (I2I) generation. " + "It will be ignored for text-to-image (T2I) generation.", + strength, + ) + strength = None + if image is not None and strength is not None and (strength < 0 or strength > 1): + raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}") + height = req.sampling_params.height or height width = req.sampling_params.width or width num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps @@ -491,16 +584,71 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): # 4. Prepare latent variables num_channels_latents = self.transformer.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - torch.float32, - device, - generator, - latents, - ) + # img2img mode: prepare latents from input image + if image is not None: + # Handle image list - take first image + if isinstance(image, list): + image = image[0] + + # Prepare image for VAE encoding using image_processor + if not isinstance(image, torch.Tensor): + init_image = self.image_processor.preprocess(image, height, width) + image = init_image.to(dtype=torch.float32, device=device) + + # Initialize scheduler kwargs for img2img + mu = calculate_shift( + (height // self.vae_scale_factor // 2) * (width // self.vae_scale_factor // 2), + 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), + ) + self.scheduler.sigma_min = 0.0 + scheduler_kwargs = {"mu": mu} + + # First initialize timesteps in scheduler + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + **scheduler_kwargs, + ) + + # Then adjust timesteps based on strength + timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) + + if num_inference_steps < 1: + raise ValueError( + f"After adjusting the num_inference_steps by strength parameter: " + f"{strength}, the number of pipeline steps is {num_inference_steps} " + f"which is < 1 and not appropriate for this pipeline." + ) + latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt) + + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds[0].dtype, + device, + generator, + latents, + image, + latent_timestep, + ) + else: + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + torch.float32, + device, + generator, + latents, + ) # Repeat prompt_embeds for num_images_per_prompt if num_images_per_prompt > 1: @@ -509,25 +657,28 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): negative_prompt_embeds = [npe for npe in negative_prompt_embeds for _ in range(num_images_per_prompt)] actual_batch_size = batch_size * num_images_per_prompt - image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) # 5. Prepare timesteps - 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), - ) - self.scheduler.sigma_min = 0.0 - scheduler_kwargs = {"mu": mu} - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - sigmas=sigmas, - **scheduler_kwargs, - ) + if image is None: + image_seq_len = (latents.shape[2] // 2) * (latents.shape[3] // 2) + 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), + ) + self.scheduler.sigma_min = 0.0 + scheduler_kwargs = {"mu": mu} + + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + sigmas=sigmas, + **scheduler_kwargs, + ) + num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) self._num_timesteps = len(timesteps) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 38d32f719..ebe4cf30b 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1421,6 +1421,7 @@ async def edit_images( negative_prompt: str | None = Form(None), num_inference_steps: int | None = Form(None), guidance_scale: float | None = Form(None), + strength: float | None = Form(None), true_cfg_scale: float | None = Form(None), seed: int | None = Form(None), generator_device: str | None = Form(None), @@ -1551,6 +1552,7 @@ async def edit_images( # 3.4 Add optional parameters ONLY if provided _update_if_not_none(gen_params, "num_inference_steps", num_inference_steps) _update_if_not_none(gen_params, "guidance_scale", guidance_scale) + _update_if_not_none(gen_params, "strength", strength) _update_if_not_none(gen_params, "true_cfg_scale", true_cfg_scale) # If seed is not provided, generate a random one to ensure # a proper generator is initialized in the backend. diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 7824e7092..9cb6c4433 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -241,6 +241,7 @@ class OmniDiffusionSamplingParams: guidance_scale_provided: bool = False guidance_scale_2: float | None = None guidance_rescale: float = 0.0 + strength: float | None = None # I2I: Z-Image specific now, uses to control denoising start timestep decode_timestep: float | list[float] | None = None decode_noise_scale: float | list[float] | None = None eta: float = 0.0