Add FLUX.2 Klein Inpaint Pipeline by adi776borate · Pull Request #13050 · huggingface/diffusers

adi776borate · 2026-01-28T19:22:14Z

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

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).
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 · 2026-01-28T20:13:38Z

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.

src/diffusers/pipelines/flux2/pipeline_flux2_klein_inpaint.py

Copilot · 2026-01-29T02:50:11Z

tests/pipelines/flux2/test_pipeline_flux2_klein_inpaint.py

+
+class Flux2KleinInpaintPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+    pipeline_class = Flux2KleinInpaintPipeline
+    params = frozenset(["prompt", "height", "width", "guidance_scale", "prompt_embeds"])


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 · 2026-01-29T02:50:11Z

tests/pipelines/flux2/test_pipeline_flux2_klein_inpaint.py

+    @unittest.skip("Needs to be revisited")
+    def test_encode_prompt_works_in_isolation(self):
+        pass


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 · 2026-01-29T09:47:11Z

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 · 2026-02-03T16:25:13Z

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

asomoza · 2026-02-03T16:27:39Z

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 added 6 commits January 26, 2026 23:53

Add Flux2KleinInpaintPipeline

2d22e4c

Fixed mask channel mismatch and a bit of cleaning

d213e59

Added tests and minor refactors

738ac43

Added support for reference images for inpainting

6fd76dd

Style fixes

2516f06

Fixed the example docstring

c44b69a

sayakpaul requested a review from yiyixuxu January 29, 2026 02:39

sayakpaul assigned asomoza and unassigned asomoza Jan 29, 2026

sayakpaul requested review from asomoza and Copilot January 29, 2026 02:39

Copilot started reviewing on behalf of sayakpaul January 29, 2026 02:40 View session

Copilot AI reviewed Jan 29, 2026

View reviewed changes

Corrected mask latent preparation for correct dimensional alignment

9502d77

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add FLUX.2 Klein Inpaint Pipeline#13050

Add FLUX.2 Klein Inpaint Pipeline#13050
adi776borate wants to merge 7 commits intohuggingface:mainfrom
adi776borate:feature/flux2-klein-inpaint

adi776borate commented Jan 28, 2026

Uh oh!

Natans8 commented Jan 28, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI Jan 29, 2026

Uh oh!

Copilot AI Jan 29, 2026

Uh oh!

adi776borate commented Jan 29, 2026

Uh oh!

adi776borate commented Feb 3, 2026

Uh oh!

asomoza commented Feb 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

adi776borate commented Jan 28, 2026

What does this PR do?

Examples

Basic Inpainting

Inpainting with Reference Image

Known Limitations

Before submitting

Who can review?

Uh oh!

Natans8 commented Jan 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

adi776borate commented Jan 29, 2026

Uh oh!

adi776borate commented Feb 3, 2026

Uh oh!

asomoza commented Feb 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Natans8 commented Jan 28, 2026 •

edited

Loading