Add FLUX.2 Klein Inpaint Pipeline

[adi776borate]

What does this PR do?

Fixes #13005
This PR adds the Flux2KleinInpaintPipeline for image inpainting using the FLUX.2 [Klein] model with optional reference image conditioning.

Examples

Basic Inpainting

Image	Mask	Prompt	Result
		Face of a yellow cat, high resolution, sitting on a park bench
		A young college boy, high resolution, sitting on a park bench

Inpainting with Reference Image

Image	Mask	Reference	Prompt	Result
			Replace this ball

Known Limitations

1. Generation quality may vary - Some outputs may contain artifacts. This can often be mitigated with better prompts and tuning hyperparameters (strength, guidance_scale, num_inference_steps).

2. Reference image conditioning is experimental - Inpainting with image_reference may not consistently produce desired results.

Image	Mask	Reference	Prompt	Result
			Replace puppy with panda

These limitations may stem either from a bug in the pipeline implementation by me or from inherent constraints of the model. Feedback is appreciated.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@asomoza @sayakpaul
Anyone in the community is free to review the PR once the tests have passed.

[Natans8]

I have to be honest, I am not getting good results at all so far, especially when working with bounding box masks, even without a reference.
With black-forest-labs/FLUX.2-klein-4B, since it's a distilled model, guidance_scale is ignored, so I essentially have only strength and num_inference_steps to tweak.

At strength=1.0 as it often goes with editing models, I get the entire inpainting area overridden edge to edge instead of just the prompted subject, usually with a lot of artefacts.
At lower strengths I get my prompt plain ignored.

On the other hand, I don't have the same issues with the regular Flux2KleinPipeline with the same parameter values and prompt, it changes the prompted subject and leaves the rest of the image untouched. Same goes for other model pipelines like FluxKontextInpaintPipeline.

If you ask, I can provide examples of the results I get once my GPU frees up.

[Copilot]

Pull request overview

This PR adds the Flux2KleinInpaintPipeline to enable image inpainting capabilities for the FLUX.2 Klein model. The pipeline supports both basic text-guided inpainting and experimental reference image conditioning, addressing issue #13005.

Changes:

• Implements a new inpainting pipeline for FLUX.2 Klein with masking support
• Adds optional reference image conditioning for more controlled inpainting
• Extends Flux2ImageProcessor with do_binarize and do_convert_grayscale parameters for mask processing

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py	Main pipeline implementation with inpainting logic, mask handling, and reference image support
tests/pipelines/flux2/test_pipeline_flux2_klein_inpaint.py	Test suite covering basic inpainting functionality including different prompts, output shapes, and strength variations
src/diffusers/pipelines/flux2/image_processor.py	Enhanced image processor with binarization and grayscale conversion options
src/diffusers/pipelines/flux2/__init__.py	Export declarations for the new pipeline
src/diffusers/pipelines/__init__.py	Top-level export declarations
src/diffusers/__init__.py	Main package export declarations
src/diffusers/utils/dummy_torch_and_transformers_objects.py	Dummy object for missing dependencies

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The params frozenset should include "image" and "mask_image" since these are required parameters for an inpainting pipeline. Following the pattern in other inpainting pipelines (see tests/pipelines/pipeline_params.py TEXT_GUIDED_IMAGE_INPAINTING_PARAMS), the params should be:

params = frozenset(["prompt", "image", "mask_image", "height", "width", "guidance_scale", "prompt_embeds"])

Without including these parameters, the standard pipeline tests from PipelineTesterMixin may not properly validate these critical inputs.

[Copilot]

The PR description highlights the reference image conditioning feature (image_reference parameter), which is a key differentiating feature of this pipeline. However, there are no tests for this functionality. Consider adding a test like:

def test_flux2_klein_inpaint_with_reference_image(self):
    pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
    inputs = self.get_dummy_inputs(torch_device)
    
    # Add a reference image
    image_reference = floats_tensor((1, 3, 32, 32), rng=random.Random(0)).to(torch_device)
    inputs["image_reference"] = image_reference
    
    output_with_ref = pipe(**inputs).images[0]
    
    # Test without reference
    inputs_no_ref = self.get_dummy_inputs(torch_device)
    output_no_ref = pipe(**inputs_no_ref).images[0]
    
    # Outputs should be different
    max_diff = np.abs(output_with_ref - output_no_ref).max()
    assert max_diff > 1e-6

This would help ensure the reference image functionality works as intended.

[adi776borate]

I'd missed patchifying the mask latents earlier. With that corrected, here are some new example generations:

Image	Mask	Prompt	Result
		Face of a yellow cat, high resolution, sitting on a park bench	Using 4B model
			Using 9B model
		A young college boy, high resolution, sitting on a park bench	Using 4B model
			Using 9B model
		A young woman standing with bouquet, natural pose, realistic proportions, consistent lighting and environment	Using 4B model
		A young woman standing with bouquet, natural pose, realistic proportions, consistent lighting and environment	Using 4B model

Below all are using 9B model with strength=1.0:

Image	Mask	Reference	Prompt	Result
			Replace this ball
			Replace puppy with panda

@Natans8 this might improve what you observed but I'll suggest wait till maintainers' review.
Thanks!

[adi776borate]

@asomoza @yiyixuxu
Sorry for tagging again, but can you please review this?

[asomoza]

oh I was waiting to see if the your changes fixed the issue that @Natans8 was having, I'll give it a test today review the PR

[adi776borate]

Hi, @Natans8 can you please try my last commit? Also, it would be helpful, if you provide some examples.
@asomoza ,would appreciate a review when you get time.

[Natans8]

Hi, @Natans8 can you please try my last commit? Also, it would be helpful, if you provide some examples. @asomoza ,would appreciate a review when you get time.

It is still working a little weird to me. These are my results on the latest branch:
For example for the following code on

pipe = Flux2KleinInpaintPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-4B", torch_dtype=torch.bfloat16)
inpainted_img = pipe(
    prompt="Remove the person, keep the background unchanged",
    image=image,
    mask_image=mask_pil, #rectangular mask
    num_inference_steps=4,
    guidance_scale=1.0,
    generator=generator,
    strength=1.0).images[0]

And this image

I get the following:

Or if I try to change the prompt to "remove all people" I can get something like this:

"Change the person's cloths to blue":

Same prompt with strength 0.8:

For comparison this is Flux2KleinPipeline with the prompt "Remove the boy in red":

"Change the clothes of the boy in red to blue, keep the others unchanged":

I had to fiddle a little in the prompt to specify which person I was referring to, not sure if that comparison is fair, but the point is that there is no artifacts.

When trying to tweak the strength, or the prompt of the inpaint pipeline I don't get much better results, sometimes it's random trees or rooms insensitive to the contents of the image, sometimes it's a black box, a white box, a slice of a wall. Of course I have to compare how it fares with 9B as well, ideally.

There is also the possibility that I'm just doing something incorrectly, I do not have a lot of experience with Diffusers. But i hope this feedback will be useful.

[asomoza]

For some reason, the demo examples for inpainting that people always use are very simple which will never fail, for example the dog sitting on the bench from the original SD, any inpainting where there's a clear separation from the subject and the background is very simple and any modern model can do it without any issues, you can even use generative fill from photoshop before the AI era and it will also be good.

For what I've tested, Klein almost never changes the rest of the image if you are specific in the edit, the VAE is also very good which makes the quality loss minimal.

The use case for this pipeline (IMO) is to be able to separate subjects when they're very similar and you can't just prompt it, in @Natans8 you can prompt to for a specific person but its enough of a good example because it's not a simple image to inpaint and we can prompt for people, also I wouldn't test this with the base model since it's not worth the 50 steps and CFG.

I did myself a test pipeline with differential diffusion to see if the model is capable of doing a good edit with restrictions, I will use for all the test just the prompt remove the people and using @Natans8 source image, I got this:

So the model is capable of doing it, so now to test this pipeline:

With the default strength:

base	distilled

I'll use just the distilled one since the results are similar, so varying the strenght I got these:

0.8	0.9

With 0.9 you could see it's somewhat decent and it's on pair for what I've seen with other inpainting pipelines.

I also tried with the prompt change the the color of the shirt to blue while maintaining everything else intact but it changed the person too.

@adi776borate does this match your results? if this looks usable to you, we can continue and review it.