[Feat] image2image for Z-Image (#1580)

Signed-off-by: Lancer <maruixiang6688@gmail.com>
Signed-off-by: Lancer <402430575@qq.com>

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.

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) |
 

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)
 

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.

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