Flux2KleinInpaintPipeline: padding_mask_crop pastes a center crop of the result instead of resizing it #15131

Describe the bug

With padding_mask_crop, Flux2KleinInpaintPipeline crops the masked region, generates it at the full output size, and then calls image_processor.apply_overlay to paste it back. The generated image should be resized to fit the crop region first, but it is center cropped instead, so the pasted result has the wrong scale.

Expected behavior: the generated image is resized to fit the crop region before being pasted back into the original image.

Actual behavior: it shows the center of the generated image, magnified.

Cause: VaeImageProcessor.apply_overlay fits the image into the crop region with self.resize(..., resize_mode="crop"). Flux2ImageProcessor overrides _resize_and_crop to only center crop, without resizing. The FLUX.2 pipelines rely on that override to preprocess reference images, so it should stay as is. The overlay path came in with the Klein inpaint pipeline (#13050). This is a different problem from #9166, which is about the base apply_overlay when the output size differs from the original image.

Proposed fix: I'd like to open a PR, once the approach is confirmed here, that overrides Flux2ImageProcessor.apply_overlay to resize the image to the crop region with VaeImageProcessor._resize_and_crop before pasting. The reference image path stays unchanged.

Reproduction

This adds padding_mask_crop to the docs example. At a low strength the masked area should still look like the source. With padding_mask_crop=32 it shows the image at the wrong scale (see the images below). The two runs use different image context, so the printed differences are only a rough indication.

import numpy as np
import torch

from diffusers import Flux2KleinInpaintPipeline
from diffusers.utils import load_image


device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = Flux2KleinInpaintPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-base-4B", torch_dtype=torch.bfloat16)
pipe.to(device)

img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
source = load_image(img_url)
mask = load_image(mask_url)
in_mask = np.array(mask.convert("L")) > 127

# At a low strength the masked area should still look like the source. Compare the saved images: with
# padding_mask_crop=32 the masked area shows the image at the wrong scale.
for padding_mask_crop in [None, 32]:
    image = pipe(
        prompt="a dog sitting on a park bench",
        image=source,
        mask_image=mask,
        padding_mask_crop=padding_mask_crop,
        strength=0.1,
        generator=torch.Generator(device).manual_seed(0),
    ).images[0]
    image.save(f"padding_mask_crop_{padding_mask_crop}.png")
    diff = np.abs(np.array(image, dtype=np.float32) - np.array(source, dtype=np.float32))[in_mask].mean()
    print(f"padding_mask_crop={padding_mask_crop}: mean abs diff to source inside mask = {diff:.2f}")

Output:

padding_mask_crop=None: mean abs diff to source inside mask = 6.72
padding_mask_crop=32: mean abs diff to source inside mask = 46.03

From left: source, padding_mask_crop=None, padding_mask_crop=32. With padding_mask_crop=32, the dog inside the mask is pasted at a larger scale and the bench shows a seam.

Logs

System Info

• 🤗 Diffusers version: 0.42.0.dev0
• Platform: Linux-6.18.48-x86_64-with-glibc2.39
• Running on Google Colab?: No
• Python version: 3.12.3
• PyTorch version (GPU?): 2.14.1+cu126 (True)
• Huggingface_hub version: 2.2.0
• Transformers version: 5.20.2
• Accelerate version: 1.15.0
• PEFT version: not installed
• Safetensors version: 0.8.0
• xFormers version: not installed
• Accelerator: NVIDIA L4, 23034 MiB
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

Source revision: 537887041938fc6fe61ad583c3ae77477c9ebac9

Who can help?

No response