update: FluxKontextInpaintPipeline support #11820

vuongminh1907 · 2025-06-27T12:24:48Z

🚀 What does this PR do?

Hello! 👋 I'm truly impressed with Flux Kontext, but I noticed that inpainting functionality hasn’t been fully integrated yet. This PR adds support for inpainting using the 🤗 Diffusers library.

This contribution introduces:

🎯 Inpainting with text only

Example using FluxKontextInpaintPipeline with just a prompt:

import torch 
from diffusers import FluxKontextInpaintPipeline
from diffusers.utils import load_image

prompt = "Change the yellow dinosaur to green one"

img_url = "https://github.com/ZenAI-Vietnam/Flux-Kontext-pipelines/blob/main/assets/dinosaur_input.jpeg?raw=true"
mask_url = "https://github.com/ZenAI-Vietnam/Flux-Kontext-pipelines/blob/main/assets/dinosaur_mask.png?raw=true"

source = load_image(img_url)
mask = load_image(mask_url)

image = pipe(prompt=prompt, image=source, mask_image=mask,strength=1.0).images[0]
image.save("kontext_inpainting_normal.png")

🖼️ Original image and mask:

✅ Result using FluxKontextInpaintPipeline:
⚠️ When using the regular FluxKontext editing pipeline, the color change was not correctly applied to the target object:

🧩 Inpainting with image conditioning

In addition to text prompts, FluxKontextInpaintPipeline also supports conditioning on a reference image via the image_reference parameter:

import torch 
from diffusers import FluxKontextInpaintPipeline
from diffusers.utils import load_image

pipe = FluxKontextInpaintPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16)
pipe.to("cuda")
prompt = "Replace this ball"

img_url = "https://images.pexels.com/photos/39362/the-ball-stadion-football-the-pitch-39362.jpeg?auto=compress&cs=tinysrgb&dpr=1&w=500"
mask_url = "https://github.com/ZenAI-Vietnam/Flux-Kontext-pipelines/blob/main/assets/ball_mask.png?raw=true"
image_reference_url = "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcTah3x6OL_ECMBaZ5ZlJJhNsyC-OSMLWAI-xw&s"

source = load_image(img_url)
mask = load_image(mask_url)
image_reference = load_image(image_reference_url)

mask = pipe.mask_processor.blur(mask, blur_factor=12)
image = pipe(prompt=prompt, image=source, mask_image=mask,image_reference=image_reference,strength=1.0).images[0]
image.save("kontext_inpainting_ref.png")

📥 Input image, mask, and reference image:
🎉 Output using FluxKontextInpaintPipeline:

I hope this PR will be helpful for the community and contribute positively to the Diffusers ecosystem! 🌱

Core library:

Schedulers: @yiyixuxu
Pipelines and pipeline callbacks: @yiyixuxu and @asomoza
Training examples: @sayakpaul
Docs: @stevhliu and @sayakpaul
JAX and MPS: @pcuenca
Audio: @sanchit-gandhi
General functionalities: @sayakpaul @yiyixuxu @DN6

nitinmukesh · 2025-06-27T13:04:53Z

Awesome. 👍

apolinario · 2025-06-27T13:21:52Z

Fantastic! Would be cool to have a demo for it on Hugging Face Spaces while the PR gets reviewed :-)

HuggingFaceDocBuilderDev · 2025-06-27T19:33:39Z

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

thanks for the super cool PR! I left some comments
would you be able to provide more results for different strength values for both with and without the image_reference?

src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py

yiyixuxu · 2025-06-27T20:09:20Z

src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py

+        image: Optional[PipelineImageInput] = None,
+        image_reference: Optional[PipelineImageInput] = None,
+        mask_image: PipelineImageInput = None,
+        masked_image_latents: PipelineImageInput = None,


let's not support masked_image_latents for now so that we can simplify the logic a bit here - it's not used here

yiyixuxu · 2025-06-27T20:10:17Z

src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py

+        else:
+            masked_image = masked_image_latents
+
+        mask, masked_image_latents = self.prepare_mask_latents(


I don't see masked_image_latents being used in this pipeline, let's remove logics that we don't need here

src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py

vuongminh1907 · 2025-06-28T07:34:28Z

Hi @yiyixuxu, I just updated the code based on your comments — thank you so much!

About updating the results with more strength values, I did a quick test and here are the outputs:

Inpainting with text only

Inpainting with image conditioning

It seems that KontextInpaint works quite well with strength=1.0.

vuongminh1907 · 2025-06-30T10:30:44Z

Fantastic! Would be cool to have a demo for it on Hugging Face Spaces while the PR gets reviewed :-)

While waiting for the PR to be reviewed, feel free to test it using:
pip install git+https://github.com/vuongminh1907/diffusers

nitinmukesh · 2025-06-30T16:03:45Z

Thanks again @vuongminh1907

Tested using
pip install git+https://github.com/huggingface/diffusers.git@refs/pull/11820/head

Working good with nunchaku optimization. Will do more tests.

Kindly review and merge.

Nunchaku output

yiyixuxu

can we add a test case? will merge soon

vuongminh1907 · 2025-07-01T01:25:00Z

Just to confirm @yiyixuxu , should I add another test case, or will you take care of it?

nne998 · 2025-07-01T07:58:21Z

Hi! Nice work! Can you add a custom node to ComfyUI to support this ? @vuongminh1907

vuongminh1907 · 2025-07-01T08:02:27Z

Hi! Nice work! Can you add a custom node to ComfyUI to support this ? @vuongminh1907
okey @nne998, I will take a look soon

vuongminh1907 · 2025-07-01T10:01:16Z

@nne998, I think ComfyUI can also work, these are 2 examples:

Prompt: Make the girl take a phone instead of bread
Prompt: Make the girl take eating apple instead of bread

strawberrymelonpanda · 2025-07-01T11:49:35Z

I think ComfyUI can also work, these are 2 examples:

@vuongminh1907 How about the image_reference parameter? Will that work with existing nodes?
Great work, and thanks!

Edit: Apologies for being somewhat off-topic with a ComfyUI question. I was linked here from a ComfyUI issue and didn't realize I'd switched repos at first.

vuongminh1907 · 2025-07-01T12:25:02Z

@strawberrymelonpanda, I’ll take a look at the image_reference later. In the meantime, I’ve just released the Kontext Inpainting Node here: ComfyUI-Kontext-Inpainting

yiyixuxu · 2025-07-01T22:08:03Z

would you be able to add a test case? can reference the current tests for flux pipelines https://github.com/huggingface/diffusers/tree/main/tests/pipelines/flux

vuongminh1907 · 2025-07-02T07:38:32Z

thanks @yiyixuxu , I added test case and it passed all checks.

yiyixuxu · 2025-07-02T08:42:00Z

@bot /style

github-actions · 2025-07-02T08:42:33Z

Style bot fixed some files and pushed the changes.

yiyixuxu

thank you so much!

* update: FluxKontextInpaintPipeline support * fix: Refactor code, remove mask_image_latents and ruff check * feat: Add test case and fix with pytest * Apply style fixes * copies --------- Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

HaydenCupcake · 2025-07-07T13:54:01Z

Hi @yiyixuxu, I just updated the code based on your comments — thank you so much!

About updating the results with more strength values, I did a quick test and here are the outputs:

Inpainting with text only

Inpainting with image conditioning

It seems that KontextInpaint works quite well with strength=1.0.

Hi, is there a ComfyUI workflow that you can share that already includes "Inpainting with image conditioning"?

nne998 · 2025-07-08T02:24:42Z

Hi there, I’ve tested the 'Kontext Inpainting Conditioning' node in ComfyUI, and it works great! Could you add support for reference image conditioning to this node? @vuongminh1907

kkumar61 · 2025-07-08T21:39:46Z

Hi @yiyixuxu, I just updated the code based on your comments — thank you so much!
About updating the results with more strength values, I did a quick test and here are the outputs:

Inpainting with text only

Inpainting with image conditioning

It seems that KontextInpaint works quite well with strength=1.0.

Hi, is there a ComfyUI workflow that you can share that already includes "Inpainting with image conditioning"?

@HaydenCupcake I've been trying your pipeline and it seems to working only on your given example images when reference image is also given. It also doesn't work if we modify the image size ratios of your source image and reference image. Any idea what might be the issue?

HaydenCupcake · 2025-07-08T21:52:41Z

@HaydenCupcake I've been trying your pipeline and it seems to working only on your given example images when reference image is also given. It also doesn't work if we modify the image size ratios of your source image and reference image. Any idea what might be the issue?

+@yiyixuxu

Samuel-Recovery-123 · 2025-07-14T09:36:31Z

@yiyixuxu I have encountered the same problem with @HaydenCupcake

nitinmukesh · 2025-07-14T13:38:03Z

@Samuel-Recovery-123 @kkumar61

Please could you share the input images and prompt for me to check quickly.

yiyixuxu · 2025-07-14T17:35:09Z

hi @kkumar61
do you know that it does not work functionally or just the result is not as expected? can you share a code example?

it seems to working only on your given example images when reference image is also given

kkumar61 · 2025-07-14T17:57:51Z

hi @kkumar61 do you know that it does not work functionally or just the result is not as expected? can you share a code example?

it seems to working only on your given example images when reference image is also given

@yiyixuxu after enough tests, it seems to be an issue with the understanding of Kontext itself. May be it's not related to the pipeline. Thanks for your attention into this!

kkumar61 · 2025-07-16T20:34:50Z

@yiyixuxu @nitinmukesh I've an interesting finding, you might want to take a look at it and let me know what are your thoughts on this!

Source Image-

Target product to be placed-

Output--

now with another shot of the same target product--

output--

asomoza · 2025-07-17T01:12:07Z

@kkumar61 Kontext wasn't trained for using masks, the model was trained to use the whole image so your result makes sense, if you want a product placement you should not use inpainting but the regular pipeline or else you're going to get similar results until you win the lottery and the model places the sofa in the exact same spot that your mask allows it.

To me, this pipeline will only work for some very specific scenarios like the dinosaur one, if you have multiple objects in your image and you just want to change a specific one of them, then you can mask it to so the model works on that one, but again, you can just draw a rectangle over it an tell the model to change the "object" inside the rectangle.

vuongminh1907 · 2025-07-17T04:22:05Z

@kkumar61 @asomoza , Kontext Inpainting works well only with text conditioning. When using a reference image, it doesn't perform as well because the model wasn't trained for that task. Naturally, it works only with specific cases. If you want to try inpainting with references, let's train a LoRA for your own task( Like ACE did) using this KontextInpaintPipeline. I think that would work better.

update: FluxKontextInpaintPipeline support

a8ca715

yiyixuxu reviewed Jun 27, 2025

View reviewed changes

fix: Refactor code, remove mask_image_latents and ruff check

4845478

yiyixuxu reviewed Jun 30, 2025

View reviewed changes

Merge branch 'main' into main

1f7939d

yiyixuxu added the close-to-merge label Jun 30, 2025

This was referenced Jul 1, 2025

Need flux kontext mask inpaint support comfyanonymous/ComfyUI#8754

Open

【Feature Request】 Need a Flux kontext mask inpaint support node kijai/ComfyUI-KJNodes#324

Open

nitinmukesh mentioned this pull request Jul 1, 2025

Run FluxKontextPipeline with multi-images input, But An error occurred #11824

Closed

feat: Add test case and fix with pytest

73e4ebf

github-actions bot and others added 3 commits July 2, 2025 08:42

Apply style fixes

d91196e

copies

fd35734

Merge branch 'main' into main

7c8840a

yiyixuxu approved these changes Jul 2, 2025

View reviewed changes

yiyixuxu merged commit d6fa329 into huggingface:main Jul 2, 2025
17 of 18 checks passed

strawberrymelonpanda mentioned this pull request Jul 2, 2025

[feature request] image_reference parameter support ZenAI-Vietnam/ComfyUI-Kontext-Inpainting#2

Open

yiyixuxu removed the close-to-merge label Jul 3, 2025

update: FluxKontextInpaintPipeline support #11820

update: FluxKontextInpaintPipeline support #11820

Uh oh!

Conversation

vuongminh1907 commented Jun 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🚀 What does this PR do?

🎯 Inpainting with text only

🧩 Inpainting with image conditioning

Uh oh!

nitinmukesh commented Jun 27, 2025

Uh oh!

apolinario commented Jun 27, 2025

Uh oh!

HuggingFaceDocBuilderDev commented Jun 27, 2025

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yiyixuxu Jun 27, 2025

Choose a reason for hiding this comment

Uh oh!

yiyixuxu Jun 27, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

vuongminh1907 commented Jun 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Inpainting with text only

Inpainting with image conditioning

Uh oh!

vuongminh1907 commented Jun 30, 2025

Uh oh!

nitinmukesh commented Jun 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

vuongminh1907 commented Jul 1, 2025

Uh oh!

nne998 commented Jul 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vuongminh1907 commented Jul 1, 2025

Uh oh!

vuongminh1907 commented Jul 1, 2025

Uh oh!

strawberrymelonpanda commented Jul 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vuongminh1907 commented Jul 1, 2025

Uh oh!

yiyixuxu commented Jul 1, 2025

Uh oh!

vuongminh1907 commented Jul 2, 2025

Uh oh!

yiyixuxu commented Jul 2, 2025

Uh oh!

github-actions bot commented Jul 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yiyixuxu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

HaydenCupcake commented Jul 7, 2025

Inpainting with text only

Inpainting with image conditioning

Uh oh!

nne998 commented Jul 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kkumar61 commented Jul 8, 2025

Inpainting with text only

Inpainting with image conditioning

Uh oh!

HaydenCupcake commented Jul 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

vuongminh1907 commented Jun 27, 2025 •

edited

Loading

vuongminh1907 commented Jun 28, 2025 •

edited

Loading

nitinmukesh commented Jun 30, 2025 •

edited

Loading

nne998 commented Jul 1, 2025 •

edited

Loading

strawberrymelonpanda commented Jul 1, 2025 •

edited

Loading

github-actions bot commented Jul 2, 2025 •

edited

Loading

nne998 commented Jul 8, 2025 •

edited

Loading

HaydenCupcake commented Jul 8, 2025 •

edited

Loading