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.

[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.

[adi776borate]

Yes, I also got sub-optimal results, sometimes very poor.
I used 9B model with strength 1.0.

1. "Remove the person" (non-rectangular mask):

2. "Remove the person" (rectangular mask):

3. "Change the person's clothes to Blue" (non-rectangular mask):

I'll take a look at this again in this week to identify the bug.

[asomoza]

left some more comments, don't forget to run:

make style 
make quality

so we can run the tests.

[asomoza]

here we're concatenating the latents and the condition_image_latents, but when called with the methods, you're passing a dtype that can be different.

latents uses prompt_embeds.dtype while condition_image_latents uses self.vae.dtype, this means that if the user changes the vae dtype (for example to use a higher precision for encoding or decoding) this will fail.

This is probably something we can also fix in the other PR

[adi776borate]

@asomoza I have made the suggested changes and also modified the batching logic as this specific test was failing with earlier code : test_inference_batch_single_identical. The previous code was merging all images in a batch into a single sequence. I changed it to have 1-to-1 mapping for batched inputs and broadcasting for single inputs.

[adi776borate]

Hi! If everything appears to be in order, shall we proceed with the tests?

[asomoza]

thanks, left a few comments, we're close to merge now

[asomoza]

you skip the preprocessing but in 1017 you set the processed_image_reference as None, so if the user passes the image_reference as tensors it skips this and never set the processed_image_reference

[adi776borate]

Added an else block to handle it

[asomoza]

Suggested change

	else:
	image_latents = torch.cat([image_latents], dim=0)

the image_latents are already the same

[adi776borate]

Agreed, removing it

[asomoza]

is this path ever reached?

[adi776borate]

Yes it is reachable when we pass multiple reference images for a single sample.

[asomoza]

Suggested change

	batch_params = frozenset(["prompt", "image", "mask_image"])
	batch_params = frozenset(["prompt", "image", "image_reference", "mask_image"])

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

@claude can you help to do a review here?

[github-actions]

Claude Code is working…

I'll analyze this and get back to you.

View job run

[yiyixuxu]

@bot /style

[github-actions]

Style bot fixed some files and pushed the changes.

[yiyixuxu]

@adi776borate can you run make fix-copies?

[adi776borate]

@adi776borate can you run make fix-copies?

I intentionally held off on that. Could you please review #13299 first? I discussed this with @asomoza here.

[adi776borate]

@asomoza I have addressed your last comments and also removed some redundant dtype casts while I was at it. This should now correctly handle pre-encoded latents as source or ref image input. I think we should also update the docs to tell the user that, we expect patchified latents as input.