[Modular]: Add Z-Image Inpaint and Controlnet Inpaint modular pipeline

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

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