diff --git a/docs/source/en/api/pipelines/z_image.md b/docs/source/en/api/pipelines/z_image.md index ea7f010b51d6..3516cbfef3f7 100644 --- a/docs/source/en/api/pipelines/z_image.md +++ b/docs/source/en/api/pipelines/z_image.md @@ -88,6 +88,85 @@ image = pipe( image.save("zimage_inpaint.png") ``` +## Modular inpainting + +[`ModularPipeline`] automatically selects the Z-Image inpainting workflow when both `image` and `mask_image` are provided. White mask regions are regenerated and black regions are preserved. +Use `padding_mask_crop` to generate only around the masked region; it requires the default PIL output so the result can be overlaid onto the original image. + +```python +import torch +from diffusers import ModularPipeline +from diffusers.utils import load_image + +pipe = ModularPipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo") +pipe.load_components(dtype=torch.bfloat16) +pipe.to("cuda") + +image = load_image("path/to/image.png").convert("RGB") +mask_image = load_image("path/to/mask.png").convert("L") + +output = pipe( + prompt="A beautiful lake with mountains in the background", + image=image, + mask_image=mask_image, + height=image.height, + width=image.width, + strength=1.0, + num_inference_steps=8, + generator=torch.Generator(device="cuda").manual_seed(42), + output="images", +)[0] +output.save("zimage_modular_inpaint.png") +``` + +To add a ControlNet inpaint condition, load a compatible [`ZImageControlNetModel`] and update the modular pipeline. The control image is used together with the source image and mask. `control_guidance_start` and `control_guidance_end` specify the normalized denoising interval in which ControlNet is active. + +```python +import torch +from huggingface_hub import hf_hub_download +from diffusers import ModularPipeline, ZImageControlNetModel +from diffusers.utils import load_image + +controlnet = ZImageControlNetModel.from_single_file( + hf_hub_download( + "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.0", + filename="Z-Image-Turbo-Fun-Controlnet-Union-2.1.safetensors", + ), + torch_dtype=torch.bfloat16, +) + +pipe = ModularPipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo") +pipe.load_components(dtype=torch.bfloat16) +pipe.update_components(controlnet=controlnet) +pipe.to("cuda") + +image = load_image( + "https://huggingface.co/alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.0/resolve/main/asset/inpaint.jpg?download=true" +).convert("RGB") +mask_image = load_image( + "https://huggingface.co/alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.0/resolve/main/asset/mask.jpg?download=true" +).convert("L") +control_image = load_image( + "https://huggingface.co/alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union-2.0/resolve/main/asset/pose.jpg?download=true" +).convert("RGB") + +output = pipe( + prompt="A woman standing on a sunny coast, full-body portrait", + image=image, + mask_image=mask_image, + control_image=control_image, + controlnet_conditioning_scale=0.75, + control_guidance_start=0.0, + control_guidance_end=1.0, + height=image.height, + width=image.width, + num_inference_steps=25, + generator=torch.Generator(device="cuda").manual_seed(43), + output="images", +)[0] +output.save("zimage_modular_controlnet_inpaint.png") +``` + ## ZImagePipeline [[autodoc]] ZImagePipeline diff --git a/src/diffusers/modular_pipelines/z_image/before_denoise.py b/src/diffusers/modular_pipelines/z_image/before_denoise.py index 9b6c1976a1bf..1864868e3983 100644 --- a/src/diffusers/modular_pipelines/z_image/before_denoise.py +++ b/src/diffusers/modular_pipelines/z_image/before_denoise.py @@ -16,7 +16,7 @@ import torch -from ...models import ZImageTransformer2DModel +from ...models import ZImageControlNetModel, ZImageTransformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging from ...utils.torch_utils import randn_tensor @@ -423,14 +423,20 @@ def inputs(self) -> list[InputParam]: type_hint=int, description="Number of prompts, the final batch size of model inputs should be `batch_size * num_images_per_prompt`. Can be generated in input step.", ), - InputParam("dtype", type_hint=torch.dtype, description="The dtype of the model inputs"), + InputParam( + "dtype", + type_hint=torch.dtype, + description="The dtype of the model inputs", + ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( - "latents", type_hint=torch.Tensor, description="The initial latents to use for the denoising process" + "latents", + type_hint=torch.Tensor, + description="The initial latents to use for the denoising process", ) ] @@ -521,7 +527,9 @@ def inputs(self) -> list[InputParam]: def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( - "timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process" + "timesteps", + type_hint=torch.Tensor, + description="The timesteps to use for the denoising process", ), ] @@ -530,7 +538,10 @@ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> P block_state = self.get_block_state(state) device = components._execution_device - latent_height, latent_width = block_state.latents.shape[2], block_state.latents.shape[3] + latent_height, latent_width = ( + block_state.latents.shape[2], + block_state.latents.shape[3], + ) image_seq_len = (latent_height // 2) * (latent_width // 2) # sequence length after patchify mu = calculate_shift( @@ -586,7 +597,10 @@ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> P block_state = self.get_block_state(state) self.check_inputs(components, block_state) - init_timestep = min(block_state.num_inference_steps * block_state.strength, block_state.num_inference_steps) + init_timestep = min( + block_state.num_inference_steps * block_state.strength, + block_state.num_inference_steps, + ) t_start = int(max(block_state.num_inference_steps - init_timestep, 0)) timesteps = components.scheduler.timesteps[t_start * components.scheduler.order :] @@ -625,3 +639,171 @@ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> P self.set_block_state(state, block_state) return components, state + + +class ZImageInpaintInputStep(ModularPipelineBlocks): + model_name = "z-image" + + @property + def description(self) -> str: + return "Expands source image latents and the inpaint mask to the denoising batch." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("image_latents", required=True, type_hint=torch.Tensor), + InputParam("mask", required=True, type_hint=torch.Tensor), + InputParam("batch_size", required=True, type_hint=int), + InputParam("num_images_per_prompt", default=1, type_hint=int), + InputParam("height"), + InputParam("width"), + ] + + @torch.no_grad() + def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + height, width = calculate_dimension_from_latents( + block_state.image_latents, components.vae_scale_factor_spatial + ) + block_state.height = block_state.height or height + block_state.width = block_state.width or width + block_state.image_latents = repeat_tensor_to_batch_size( + "image_latents", + block_state.image_latents, + block_state.batch_size, + block_state.num_images_per_prompt, + ) + block_state.mask = torch.nn.functional.interpolate( + block_state.mask.to( + device=components._execution_device, + dtype=block_state.image_latents.dtype, + ), + size=block_state.image_latents.shape[-2:], + mode="nearest", + ) + block_state.mask = repeat_tensor_to_batch_size( + "mask", + block_state.mask, + block_state.batch_size, + block_state.num_images_per_prompt, + ) + self.set_block_state(state, block_state) + return components, state + + +class ZImagePrepareInpaintLatentsStep(ModularPipelineBlocks): + model_name = "z-image" + + @property + def description(self) -> str: + return "Adds noise to source-image latents and preserves that noise for inpaint blending." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latents", required=True, type_hint=torch.Tensor), + InputParam("image_latents", required=True, type_hint=torch.Tensor), + InputParam("timesteps", required=True, type_hint=torch.Tensor), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "image_noise", + type_hint=torch.Tensor, + description="Noise used for inpaint blending.", + ) + ] + + @torch.no_grad() + def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.image_noise = block_state.latents + timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0]) + block_state.latents = components.scheduler.scale_noise( + block_state.image_latents, timestep, block_state.image_noise + ) + self.set_block_state(state, block_state) + return components, state + + +class ZImageControlNetInpaintInputStep(ModularPipelineBlocks): + model_name = "z-image" + + @property + def description(self) -> str: + return "Expands the latent ControlNet inpaint condition to the denoising batch." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("control_image_latents", required=True, type_hint=torch.Tensor), + InputParam("batch_size", required=True, type_hint=int), + InputParam("num_images_per_prompt", default=1, type_hint=int), + InputParam("height"), + InputParam("width"), + ] + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", ZImageTransformer2DModel), + ComponentSpec("controlnet", ZImageControlNetModel), + ] + + @torch.no_grad() + def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + height = block_state.control_image_latents.shape[-2] * components.vae_scale_factor_spatial // 2 + width = block_state.control_image_latents.shape[-1] * components.vae_scale_factor_spatial // 2 + block_state.height = block_state.height or height + block_state.width = block_state.width or width + block_state.control_image_latents = repeat_tensor_to_batch_size( + "control_image_latents", + block_state.control_image_latents, + block_state.batch_size, + block_state.num_images_per_prompt, + ) + self.set_block_state(state, block_state) + return components, state + + +class ZImageControlNetBeforeDenoiserStep(ModularPipelineBlocks): + model_name = "z-image" + + @property + def description(self) -> str: + return "Prepares the per-step ControlNet conditioning schedule." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("control_guidance_start"), + InputParam.template("control_guidance_end"), + InputParam("timesteps", required=True, type_hint=torch.Tensor), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "controlnet_keep", + type_hint=list[float], + description="Per-step ControlNet conditioning multipliers.", + ) + ] + + @torch.no_grad() + def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.controlnet_keep = [ + 1.0 + - float( + i / len(block_state.timesteps) < block_state.control_guidance_start + or (i + 1) / len(block_state.timesteps) > block_state.control_guidance_end + ) + for i in range(len(block_state.timesteps)) + ] + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/z_image/decoders.py b/src/diffusers/modular_pipelines/z_image/decoders.py index 353253102376..1f249e4b815a 100644 --- a/src/diffusers/modular_pipelines/z_image/decoders.py +++ b/src/diffusers/modular_pipelines/z_image/decoders.py @@ -24,6 +24,7 @@ from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .modular_pipeline import ZImageModularPipeline logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -79,7 +80,10 @@ def __call__(self, components, state: PipelineState) -> PipelineState: vae_dtype = components.vae.dtype latents = block_state.latents.to(vae_dtype) - latents = latents / components.vae.config.scaling_factor + components.vae.config.shift_factor + latents = ( + latents / components.vae.config.scaling_factor + + components.vae.config.shift_factor + ) block_state.images = components.vae.decode(latents, return_dict=False)[0] block_state.images = components.image_processor.postprocess( @@ -89,3 +93,55 @@ def __call__(self, components, state: PipelineState) -> PipelineState: self.set_block_state(state, block_state) return components, state + + +class ZImageInpaintOverlayMaskStep(ModularPipelineBlocks): + model_name = "z-image" + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 8 * 2}), + default_creation_method="from_config", + ), + ] + + @property + def description(self) -> str: + return "Overlays a cropped inpaint result onto the original image." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("images", required=True), + InputParam("image", required=True, type_hint=PIL.Image.Image), + InputParam("mask_image", required=True, type_hint=PIL.Image.Image), + InputParam.template("padding_mask_crop"), + InputParam("crops_coords", type_hint=tuple[int, int, int, int] | None), + InputParam("output_type", default="pil", type_hint=str), + ] + + @torch.no_grad() + def __call__( + self, components: ZImageModularPipeline, state: PipelineState + ) -> PipelineState: + block_state = self.get_block_state(state) + if block_state.padding_mask_crop is not None: + if block_state.output_type != "pil": + raise ValueError( + "`output_type` must be 'pil' when `padding_mask_crop` is provided." + ) + block_state.images = [ + components.image_processor.apply_overlay( + block_state.mask_image, + block_state.image, + image, + block_state.crops_coords, + ) + for image in block_state.images + ] + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/z_image/denoise.py b/src/diffusers/modular_pipelines/z_image/denoise.py index 863df312389a..6123614eac10 100644 --- a/src/diffusers/modular_pipelines/z_image/denoise.py +++ b/src/diffusers/modular_pipelines/z_image/denoise.py @@ -18,7 +18,7 @@ from ...configuration_utils import FrozenDict from ...guiders import ClassifierFreeGuidance -from ...models import ZImageTransformer2DModel +from ...models import ZImageControlNetModel, ZImageTransformer2DModel from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging from ..modular_pipeline import ( @@ -63,7 +63,13 @@ def inputs(self) -> list[InputParam]: ] @torch.no_grad() - def __call__(self, components: ZImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + def __call__( + self, + components: ZImageModularPipeline, + block_state: BlockState, + i: int, + t: torch.Tensor, + ): latents = block_state.latents.unsqueeze(2).to( block_state.dtype ) # [batch_size, num_channels, 1, height, width] @@ -151,7 +157,11 @@ def inputs(self) -> list[tuple[str, Any]]: @torch.no_grad() def __call__( - self, components: ZImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor + self, + components: ZImageModularPipeline, + block_state: BlockState, + i: int, + t: torch.Tensor, ) -> PipelineState: components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) @@ -219,7 +229,13 @@ def description(self) -> str: ) @torch.no_grad() - def __call__(self, components: ZImageModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + def __call__( + self, + components: ZImageModularPipeline, + block_state: BlockState, + i: int, + t: torch.Tensor, + ): # Perform scheduler step using the predicted output latents_dtype = block_state.latents.dtype block_state.latents = components.scheduler.step( @@ -235,6 +251,94 @@ def __call__(self, components: ZImageModularPipeline, block_state: BlockState, i return components, block_state +class ZImageInpaintLoopAfterDenoiser(ZImageLoopAfterDenoiser): + model_name = "z-image" + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("image_latents", required=True, type_hint=torch.Tensor), + InputParam("image_noise", required=True, type_hint=torch.Tensor), + InputParam("mask", required=True, type_hint=torch.Tensor), + ] + + @torch.no_grad() + def __call__( + self, + components: ZImageModularPipeline, + block_state: BlockState, + i: int, + t: torch.Tensor, + ): + components, block_state = super().__call__(components, block_state, i, t) + image_latents = block_state.image_latents + if i < len(block_state.timesteps) - 1: + image_latents = components.scheduler.scale_noise( + image_latents, + block_state.timesteps[i + 1].unsqueeze(0), + block_state.image_noise, + ) + block_state.latents = (1 - block_state.mask) * image_latents + block_state.mask * block_state.latents + return components, block_state + + +class ZImageControlNetLoopDenoiser(ZImageLoopDenoiser): + model_name = "z-image" + + @property + def expected_components(self) -> list[ComponentSpec]: + return super().expected_components + [ComponentSpec("controlnet", ZImageControlNetModel)] + + @property + def inputs(self) -> list[tuple[str, Any]]: + return super().inputs + [ + InputParam("control_image_latents", required=True, type_hint=torch.Tensor), + InputParam("controlnet_conditioning_scale", default=0.75, type_hint=float), + InputParam("controlnet_keep", required=True, type_hint=list[float]), + ] + + @torch.no_grad() + def __call__( + self, + components: ZImageModularPipeline, + block_state: BlockState, + i: int, + t: torch.Tensor, + ) -> PipelineState: + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields) + + for guider_state_batch in guider_state: + components.guider.prepare_models(components.transformer) + cond_kwargs = guider_state_batch.as_dict() + cond_kwargs = { + key: [value.to(block_state.dtype) for value in values] + if isinstance(values, list) + else values.to(block_state.dtype) + for key, values in cond_kwargs.items() + if key in self._guider_input_fields + } + controlnet_block_samples = components.controlnet( + x=block_state.latent_model_input, + t=block_state.timestep, + cap_feats=cond_kwargs["cap_feats"], + control_context=list(block_state.control_image_latents.to(block_state.dtype).unbind(dim=0)), + conditioning_scale=block_state.controlnet_conditioning_scale * block_state.controlnet_keep[i], + ) + model_out_list = components.transformer( + x=block_state.latent_model_input, + t=block_state.timestep, + controlnet_block_samples=controlnet_block_samples, + return_dict=False, + **cond_kwargs, + )[0] + guider_state_batch.noise_pred = -torch.stack(model_out_list, dim=0).squeeze(2) + components.guider.cleanup_models(components.transformer) + + block_state.noise_pred = components.guider(guider_state)[0] + return components, block_state + + class ZImageDenoiseLoopWrapper(LoopSequentialPipelineBlocks): model_name = "z-image" @@ -273,7 +377,8 @@ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> P block_state = self.get_block_state(state) block_state.num_warmup_steps = max( - len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0 + len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, + 0, ) with self.progress_bar(total=block_state.num_inference_steps) as progress_bar: @@ -312,3 +417,21 @@ def description(self) -> str: " - `ZImageLoopAfterDenoiser`\n" "This block supports text-to-image and image-to-image tasks for Z-Image." ) + + +class ZImageInpaintDenoiseStep(ZImageDenoiseLoopWrapper): + block_classes = [ + ZImageLoopBeforeDenoiser, + ZImageLoopDenoiser(guider_input_fields={"cap_feats": ("prompt_embeds", "negative_prompt_embeds")}), + ZImageInpaintLoopAfterDenoiser, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + +class ZImageControlNetDenoiseStep(ZImageDenoiseLoopWrapper): + block_classes = [ + ZImageLoopBeforeDenoiser, + ZImageControlNetLoopDenoiser(guider_input_fields={"cap_feats": ("prompt_embeds", "negative_prompt_embeds")}), + ZImageLoopAfterDenoiser, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] diff --git a/src/diffusers/modular_pipelines/z_image/encoders.py b/src/diffusers/modular_pipelines/z_image/encoders.py index 06deb8236893..ffffb9d7fca6 100644 --- a/src/diffusers/modular_pipelines/z_image/encoders.py +++ b/src/diffusers/modular_pipelines/z_image/encoders.py @@ -15,6 +15,7 @@ import PIL import torch +import torch.nn.functional as F from transformers import Qwen2Tokenizer, Qwen3Model from ...configuration_utils import FrozenDict @@ -81,7 +82,9 @@ def get_qwen_prompt_embeds( # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents def retrieve_latents( - encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" + 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) @@ -100,6 +103,7 @@ def encode_vae_image( device: torch.device, dtype: torch.dtype, latent_channels: int = 16, + sample_mode: str = "sample", ): if not isinstance(image_tensor, torch.Tensor): raise ValueError(f"Expected image_tensor to be a tensor, got {type(image_tensor)}.") @@ -113,12 +117,16 @@ def encode_vae_image( if isinstance(generator, list): image_latents = [ - retrieve_latents(vae.encode(image_tensor[i : i + 1]), generator=generator[i]) + retrieve_latents( + vae.encode(image_tensor[i : i + 1]), + generator=generator[i], + sample_mode=sample_mode, + ) for i in range(image_tensor.shape[0]) ] image_latents = torch.cat(image_latents, dim=0) else: - image_latents = retrieve_latents(vae.encode(image_tensor), generator=generator) + image_latents = retrieve_latents(vae.encode(image_tensor), generator=generator, sample_mode=sample_mode) image_latents = (image_latents - vae.config.shift_factor) * vae.config.scaling_factor @@ -341,3 +349,208 @@ def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> P self.set_block_state(state, block_state) return components, state + + +class ZImageInpaintVaeImageEncoderStep(ModularPipelineBlocks): + model_name = "z-image" + + @property + def description(self) -> str: + return "Encodes an inpaint image and preprocesses its binary mask." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKL), + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 8 * 2}), + default_creation_method="from_config", + ), + ComponentSpec( + "mask_processor", + VaeImageProcessor, + config=FrozenDict( + { + "vae_scale_factor": 8 * 2, + "do_normalize": False, + "do_binarize": True, + "do_convert_grayscale": True, + } + ), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("image", type_hint=PIL.Image.Image, required=True), + InputParam("mask_image", type_hint=PIL.Image.Image, required=True), + InputParam("height"), + InputParam("width"), + InputParam.template("padding_mask_crop"), + InputParam("control_image"), + InputParam("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "image_latents", + type_hint=torch.Tensor, + description="Latents of the source image.", + ), + OutputParam("mask", type_hint=torch.Tensor, description="Binary inpainting mask."), + OutputParam( + "crops_coords", + type_hint=tuple[int, int, int, int] | None, + description="Crop coordinates used to process and overlay the inpaint result.", + ), + ] + + @torch.no_grad() + def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + height, width = components.image_processor.get_default_height_width( + block_state.image, block_state.height, block_state.width + ) + block_state.height = height + block_state.width = width + if block_state.padding_mask_crop is not None: + block_state.crops_coords = components.mask_processor.get_crop_region( + block_state.mask_image, width, height, pad=block_state.padding_mask_crop + ) + resize_mode = "fill" + else: + block_state.crops_coords = None + resize_mode = "default" + image = components.image_processor.preprocess( + block_state.image, + height=height, + width=width, + crops_coords=block_state.crops_coords, + resize_mode=resize_mode, + ).to(device=components._execution_device, dtype=torch.float32) + block_state.mask = components.mask_processor.preprocess( + block_state.mask_image, + height=image.shape[-2], + width=image.shape[-1], + crops_coords=block_state.crops_coords, + resize_mode=resize_mode, + ) + block_state.image_latents = encode_vae_image( + image_tensor=image, + vae=components.vae, + generator=block_state.generator, + device=components._execution_device, + dtype=components.vae.dtype, + sample_mode="argmax" if block_state.control_image is not None else "sample", + ) + self.set_block_state(state, block_state) + return components, state + + +class ZImageControlNetInpaintVaeEncoderStep(ModularPipelineBlocks): + model_name = "z-image" + + @property + def description(self) -> str: + return "Encodes the ControlNet inpaint condition from the control image, source image, and mask." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKL), + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 8 * 2}), + default_creation_method="from_config", + ), + ComponentSpec( + "mask_processor", + VaeImageProcessor, + config=FrozenDict( + { + "vae_scale_factor": 8, + "do_normalize": False, + "do_binarize": True, + "do_convert_grayscale": True, + } + ), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("image", type_hint=PIL.Image.Image, required=True), + InputParam("mask_image", type_hint=PIL.Image.Image, required=True), + InputParam("control_image", type_hint=PIL.Image.Image, required=True), + InputParam("height"), + InputParam("width"), + InputParam("crops_coords", type_hint=tuple[int, int, int, int] | None), + InputParam("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "control_image_latents", + type_hint=torch.Tensor, + description="Latent ControlNet inpaint condition.", + ) + ] + + @torch.no_grad() + def __call__(self, components: ZImageModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + resize_mode = "fill" if block_state.crops_coords is not None else "default" + image = components.image_processor.preprocess( + block_state.image, + height=block_state.height, + width=block_state.width, + crops_coords=block_state.crops_coords, + resize_mode=resize_mode, + ).to(device=components._execution_device, dtype=components.vae.dtype) + control_image = components.image_processor.preprocess( + block_state.control_image, + height=image.shape[-2], + width=image.shape[-1], + crops_coords=block_state.crops_coords, + resize_mode=resize_mode, + ).to(device=components._execution_device, dtype=components.vae.dtype) + mask = components.mask_processor.preprocess( + block_state.mask_image, + height=image.shape[-2], + width=image.shape[-1], + crops_coords=block_state.crops_coords, + resize_mode=resize_mode, + ).to(device=components._execution_device, dtype=components.vae.dtype) + mask = torch.tile(mask, [1, 3, 1, 1]) + + control_image_latents = encode_vae_image( + image_tensor=control_image, + vae=components.vae, + generator=block_state.generator, + device=components._execution_device, + dtype=components.vae.dtype, + sample_mode="argmax", + ).unsqueeze(2) + image_latents = encode_vae_image( + image_tensor=image * (mask < 0.5), + vae=components.vae, + generator=block_state.generator, + device=components._execution_device, + dtype=components.vae.dtype, + sample_mode="argmax", + ).unsqueeze(2) + mask = F.interpolate(1 - mask[:, :1], size=image_latents.shape[-2:], mode="nearest").unsqueeze(2) + block_state.control_image_latents = torch.cat([control_image_latents, mask, image_latents], dim=1) + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py b/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py index 1440128a3c55..378b64fc8aad 100644 --- a/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py +++ b/src/diffusers/modular_pipelines/z_image/modular_blocks_z_image.py @@ -13,21 +13,33 @@ # limitations under the License. from ...utils import logging -from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline import ( + AutoPipelineBlocks, + ConditionalPipelineBlocks, + SequentialPipelineBlocks, +) from ..modular_pipeline_utils import OutputParam from .before_denoise import ( ZImageAdditionalInputsStep, + ZImageControlNetBeforeDenoiserStep, + ZImageControlNetInpaintInputStep, + ZImageInpaintInputStep, + ZImagePrepareInpaintLatentsStep, ZImagePrepareLatentsStep, ZImagePrepareLatentswithImageStep, ZImageSetTimestepsStep, ZImageSetTimestepsWithStrengthStep, ZImageTextInputStep, ) -from .decoders import ZImageVaeDecoderStep +from .decoders import ZImageInpaintOverlayMaskStep, ZImageVaeDecoderStep from .denoise import ( + ZImageControlNetDenoiseStep, ZImageDenoiseStep, + ZImageInpaintDenoiseStep, ) from .encoders import ( + ZImageControlNetInpaintVaeEncoderStep, + ZImageInpaintVaeImageEncoderStep, ZImageTextEncoderStep, ZImageVaeImageEncoderStep, ) @@ -164,8 +176,64 @@ def outputs(self): return [OutputParam.template("latents")] +class ZImageInpaintCoreDenoiseStep(SequentialPipelineBlocks): + block_classes = [ + ZImageTextInputStep, + ZImageInpaintInputStep, + ZImagePrepareLatentsStep, + ZImageSetTimestepsStep, + ZImageSetTimestepsWithStrengthStep, + ZImagePrepareInpaintLatentsStep, + ZImageInpaintDenoiseStep, + ] + block_names = [ + "input", + "inpaint_input", + "prepare_latents", + "set_timesteps", + "set_timesteps_with_strength", + "prepare_inpaint_latents", + "denoise", + ] + + @property + def description(self): + return "Denoises source-image latents while preserving the unmasked region." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + +class ZImageControlNetInpaintCoreDenoiseStep(SequentialPipelineBlocks): + block_classes = [ + ZImageTextInputStep, + ZImageControlNetInpaintInputStep, + ZImagePrepareLatentsStep, + ZImageSetTimestepsStep, + ZImageControlNetBeforeDenoiserStep, + ZImageControlNetDenoiseStep, + ] + block_names = [ + "input", + "controlnet_input", + "prepare_latents", + "set_timesteps", + "controlnet_before_denoiser", + "denoise", + ] + + @property + def description(self): + return "Denoises latents using a ControlNet inpaint condition." + + @property + def outputs(self): + return [OutputParam.template("latents")] + + # auto_docstring -class ZImageAutoDenoiseStep(AutoPipelineBlocks): +class ZImageAutoDenoiseStep(ConditionalPipelineBlocks): """ Denoise step that iteratively denoise the latents. This is a auto pipeline block that works for text2image and image2image tasks. - `ZImageCoreDenoiseStep` (text2image) for text2image tasks. - @@ -209,19 +277,33 @@ class ZImageAutoDenoiseStep(AutoPipelineBlocks): """ block_classes = [ + ZImageControlNetInpaintCoreDenoiseStep, + ZImageInpaintCoreDenoiseStep, ZImageImage2ImageCoreDenoiseStep, ZImageCoreDenoiseStep, ] - block_names = ["image2image", "text2image"] - block_trigger_inputs = ["image_latents", None] + block_names = ["controlnet_inpaint", "inpaint", "image2image", "text2image"] + block_trigger_inputs = ["control_image_latents", "mask", "image_latents"] + default_block_name = "text2image" + + def select_block(self, control_image_latents=None, mask=None, image_latents=None): + if control_image_latents is not None: + return "controlnet_inpaint" + if mask is not None: + return "inpaint" + if image_latents is not None: + return "image2image" + return "text2image" @property def description(self) -> str: return ( "Denoise step that iteratively denoise the latents. " - "This is a auto pipeline block that works for text2image and image2image tasks." + "This is an auto pipeline block that works for text2image, image2image, inpaint, and ControlNet inpaint tasks." " - `ZImageCoreDenoiseStep` (text2image) for text2image tasks." " - `ZImageImage2ImageCoreDenoiseStep` (image2image) for image2image tasks." + " - `ZImageInpaintCoreDenoiseStep` (inpaint) for inpaint tasks." + " - `ZImageControlNetInpaintCoreDenoiseStep` (controlnet_inpaint) for ControlNet inpaint tasks." + " - if `image_latents` is provided, `ZImageImage2ImageCoreDenoiseStep` will be used.\n" + " - if `image_latents` is not provided, `ZImageCoreDenoiseStep` will be used.\n" ) @@ -250,31 +332,51 @@ class ZImageAutoVaeImageEncoderStep(AutoPipelineBlocks): video latent representation with the first frame image condition """ - block_classes = [ZImageVaeImageEncoderStep] - block_names = ["vae_encoder"] - block_trigger_inputs = ["image"] + block_classes = [ZImageInpaintVaeImageEncoderStep, ZImageVaeImageEncoderStep] + block_names = ["inpaint", "image2image"] + block_trigger_inputs = ["mask_image", "image"] @property def description(self) -> str: return "Vae Image Encoder step that encode the image to generate the image latents" - +"This is an auto pipeline block that works for image2image tasks." + +"This is an auto pipeline block that works for image2image and inpaint tasks." + +" - `ZImageInpaintVaeImageEncoderStep` is used when `mask_image` is provided." +" - `ZImageVaeImageEncoderStep` is used when `image` is provided." +" - if `image` is not provided, step will be skipped." +class ZImageOptionalControlNetInpaintVaeEncoderStep(AutoPipelineBlocks): + block_classes = [ZImageControlNetInpaintVaeEncoderStep] + block_names = ["controlnet_inpaint"] + block_trigger_inputs = ["control_image"] + + +class ZImageInpaintDecodeStep(SequentialPipelineBlocks): + block_classes = [ZImageVaeDecoderStep(), ZImageInpaintOverlayMaskStep()] + block_names = ["decode", "mask_overlay"] + + +class ZImageAutoDecodeStep(AutoPipelineBlocks): + block_classes = [ZImageInpaintDecodeStep, ZImageVaeDecoderStep] + block_names = ["inpaint", "default"] + block_trigger_inputs = ["mask", None] + + # auto_docstring class ZImageAutoBlocks(SequentialPipelineBlocks): """ - Auto Modular pipeline for text-to-image and image-to-image using ZImage. + Auto Modular pipeline for text-to-image, image-to-image, inpainting, and ControlNet inpainting using ZImage. Supported workflows: - `text2image`: requires `prompt` - `image2image`: requires `image`, `prompt` + - `inpainting`: requires `image`, `mask_image`, `prompt` + - `controlnet_inpainting`: requires `image`, `mask_image`, `control_image`, `prompt` Components: text_encoder (`Qwen3Model`) tokenizer (`Qwen2Tokenizer`) guider (`ClassifierFreeGuidance`) vae - (`AutoencoderKL`) image_processor (`VaeImageProcessor`) transformer (`ZImageTransformer2DModel`) scheduler - (`FlowMatchEulerDiscreteScheduler`) + (`AutoencoderKL`) image_processor (`VaeImageProcessor`) mask_processor (`VaeImageProcessor`) transformer + (`ZImageTransformer2DModel`) controlnet (`ZImageControlNetModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) Inputs: prompt (`None`, *optional*): @@ -285,6 +387,12 @@ class ZImageAutoBlocks(SequentialPipelineBlocks): TODO: Add description. image (`Image`, *optional*): TODO: Add description. + mask_image (`Image`, *optional*): + Mask image for inpainting. + control_image (`Image`, *optional*): + Control image for ControlNet conditioning. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. height (`None`, *optional*): TODO: Add description. width (`None`, *optional*): @@ -303,6 +411,12 @@ class ZImageAutoBlocks(SequentialPipelineBlocks): TODO: Add description. strength (`None`, *optional*, defaults to 0.6): TODO: Add description. + control_guidance_start (`float`, *optional*, defaults to 0.0): + When to start applying ControlNet. + control_guidance_end (`float`, *optional*, defaults to 1.0): + When to stop applying ControlNet. + controlnet_conditioning_scale (`float`, *optional*, defaults to 0.75): + Scale for ControlNet conditioning. **denoiser_input_fields (`None`, *optional*): The conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. output_type (`str`, *optional*, defaults to pil): @@ -316,18 +430,32 @@ class ZImageAutoBlocks(SequentialPipelineBlocks): block_classes = [ ZImageTextEncoderStep, ZImageAutoVaeImageEncoderStep, + ZImageOptionalControlNetInpaintVaeEncoderStep, ZImageAutoDenoiseStep, - ZImageVaeDecoderStep, + ZImageAutoDecodeStep, + ] + block_names = [ + "text_encoder", + "vae_encoder", + "controlnet_vae_encoder", + "denoise", + "decode", ] - block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] _workflow_map = { "text2image": {"prompt": True}, "image2image": {"image": True, "prompt": True}, + "inpainting": {"image": True, "mask_image": True, "prompt": True}, + "controlnet_inpainting": { + "image": True, + "mask_image": True, + "control_image": True, + "prompt": True, + }, } @property def description(self) -> str: - return "Auto Modular pipeline for text-to-image and image-to-image using ZImage." + return "Auto Modular pipeline for text-to-image, image-to-image, inpainting, and ControlNet inpainting using ZImage." @property def outputs(self): diff --git a/src/diffusers/modular_pipelines/z_image/modular_pipeline.py b/src/diffusers/modular_pipelines/z_image/modular_pipeline.py index f1d8e53a3639..fe1e303f9d8f 100644 --- a/src/diffusers/modular_pipelines/z_image/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/z_image/modular_pipeline.py @@ -14,6 +14,7 @@ from ...loaders import ZImageLoraLoaderMixin +from ...models import ZImageControlNetModel from ...utils import logging from ..modular_pipeline import ModularPipeline @@ -33,6 +34,20 @@ class ZImageModularPipeline( default_blocks_name = "ZImageAutoBlocks" + def _link_controlnet(self): + controlnet = self.components.get("controlnet") + transformer = self.components.get("transformer") + if controlnet is not None and transformer is not None: + ZImageControlNetModel.from_transformer(controlnet, transformer) + + def update_components(self, **kwargs): + super().update_components(**kwargs) + self._link_controlnet() + + def load_components(self, names: list[str] | str | None = None, workflow: str | None = None, **kwargs): + super().load_components(names=names, workflow=workflow, **kwargs) + self._link_controlnet() + @property def default_height(self): return 1024 diff --git a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py index b67be5bc88e9..c6158da9967d 100644 --- a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py +++ b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py @@ -14,7 +14,14 @@ # limitations under the License. +import PIL +import torch + +from diffusers import ZImageControlNetModel from diffusers.modular_pipelines import ZImageAutoBlocks, ZImageModularPipeline +from diffusers.modular_pipelines.z_image.before_denoise import ( + ZImageControlNetBeforeDenoiserStep, +) from ..testing_utils import ( BaseModularPipelineTesterConfig, @@ -46,6 +53,32 @@ ("denoise.denoise", "ZImageDenoiseStep"), ("decode", "ZImageVaeDecoderStep"), ], + "inpainting": [ + ("text_encoder", "ZImageTextEncoderStep"), + ("vae_encoder", "ZImageInpaintVaeImageEncoderStep"), + ("denoise.input", "ZImageTextInputStep"), + ("denoise.inpaint_input", "ZImageInpaintInputStep"), + ("denoise.prepare_latents", "ZImagePrepareLatentsStep"), + ("denoise.set_timesteps", "ZImageSetTimestepsStep"), + ("denoise.set_timesteps_with_strength", "ZImageSetTimestepsWithStrengthStep"), + ("denoise.prepare_inpaint_latents", "ZImagePrepareInpaintLatentsStep"), + ("denoise.denoise", "ZImageInpaintDenoiseStep"), + ("decode.decode", "ZImageVaeDecoderStep"), + ("decode.mask_overlay", "ZImageInpaintOverlayMaskStep"), + ], + "controlnet_inpainting": [ + ("text_encoder", "ZImageTextEncoderStep"), + ("vae_encoder", "ZImageInpaintVaeImageEncoderStep"), + ("controlnet_vae_encoder", "ZImageControlNetInpaintVaeEncoderStep"), + ("denoise.input", "ZImageTextInputStep"), + ("denoise.controlnet_input", "ZImageControlNetInpaintInputStep"), + ("denoise.prepare_latents", "ZImagePrepareLatentsStep"), + ("denoise.set_timesteps", "ZImageSetTimestepsStep"), + ("denoise.controlnet_before_denoiser", "ZImageControlNetBeforeDenoiserStep"), + ("denoise.denoise", "ZImageControlNetDenoiseStep"), + ("decode.decode", "ZImageVaeDecoderStep"), + ("decode.mask_overlay", "ZImageInpaintOverlayMaskStep"), + ], } @@ -75,6 +108,81 @@ class TestZImageModularPipelineFast(ZImageModularPipelineTesterConfig, ModularPi def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=5e-3) + def test_inpaint_inference(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs.update( + image=PIL.Image.new("RGB", (32, 32), 0), + mask_image=PIL.Image.new("L", (32, 32), 255), + strength=1.0, + ) + output = pipe(**inputs, output="images") + assert output.shape == (1, 3, 32, 32) + + def test_inpaint_padding_mask_crop(self): + pipe = self.get_pipeline() + inputs = self.get_dummy_inputs() + inputs.update( + image=PIL.Image.new("RGB", (64, 32), (255, 0, 0)), + mask_image=PIL.Image.new("L", (64, 32), 0), + height=32, + width=32, + padding_mask_crop=0, + strength=1.0, + output_type="pil", + ) + inputs["mask_image"].paste(255, (16, 0, 48, 32)) + output = pipe(**inputs, output="images") + assert output[0].size == (64, 32) + assert output[0].getpixel((0, 16)) == (255, 0, 0) + + def test_controlnet_guidance_window(self): + pipe = ZImageControlNetBeforeDenoiserStep().init_pipeline() + pipe.load_components() + output = pipe( + timesteps=torch.tensor([3.0, 2.0, 1.0, 0.0]), + control_guidance_start=0.25, + control_guidance_end=0.75, + output="controlnet_keep", + ) + assert output == [0.0, 1.0, 1.0, 0.0] + + def test_controlnet_inpaint_inference(self): + pipe = self.get_pipeline() + transformer_config = pipe.transformer.config + pipe.update_components( + controlnet=ZImageControlNetModel( + control_layers_places=[0], + control_refiner_layers_places=[], + control_in_dim=33, + all_patch_size=transformer_config.all_patch_size, + all_f_patch_size=transformer_config.all_f_patch_size, + dim=transformer_config.dim, + n_refiner_layers=transformer_config.n_refiner_layers, + n_heads=transformer_config.n_heads, + n_kv_heads=transformer_config.n_kv_heads, + norm_eps=transformer_config.norm_eps, + qk_norm=transformer_config.qk_norm, + ) + ) + assert pipe.controlnet.t_embedder is pipe.transformer.t_embedder + inputs = self.get_dummy_inputs() + inputs.update( + image=PIL.Image.new("RGB", (64, 32), (255, 0, 0)), + mask_image=PIL.Image.new("L", (64, 32), 0), + control_image=PIL.Image.new("RGB", (64, 32), 0), + height=32, + width=32, + padding_mask_crop=0, + control_guidance_start=0.25, + control_guidance_end=0.75, + output_type="pil", + ) + inputs["mask_image"].paste(255, (16, 0, 48, 32)) + output = pipe(**inputs, output="images") + assert output[0].size == (64, 32) + assert output[0].getpixel((0, 16)) == (255, 0, 0) + class TestZImageModularPipelineLoading(ZImageModularPipelineTesterConfig, ModularLoadingTesterMixin): pass