update: FluxKontextInpaintPipeline support

[vuongminh1907]

🚀 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

1. 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")

2. 🖼️ Original image and mask:

3. ✅ Result using FluxKontextInpaintPipeline:
   

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

🧩 Inpainting with image conditioning

1. 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")

2. 📥 Input image, mask, and reference image:
   
   
   

3. 🎉 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]

Awesome. 👍

[apolinario]

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

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

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?

[yiyixuxu]

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]

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

[vuongminh1907]

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]

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]

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]

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

[nne998]

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

[vuongminh1907]

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

[vuongminh1907]

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

1. Prompt: Make the girl take a phone instead of bread
   

2. Prompt: Make the girl take eating apple instead of bread
   

[strawberrymelonpanda]

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]

@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]

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]

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

[yiyixuxu]

@bot /style

[github-actions]

Style bot fixed some files and pushed the changes.

[yiyixuxu]

thank you so much!