Add FLUX.2 Klein Inpaint Pipeline (#13050)

* Add Flux2KleinInpaintPipeline

* Fixed mask channel mismatch and a bit of cleaning

* Added tests and minor refactors

* Added support for reference images for inpainting

* Style fixes

* Fixed the example docstring

* Corrected mask latent preparation for correct dimensional alignment

* replace masked_image_latents context with clean_source_latents, fix mask spatial alignment and remove unused VAE encoding

* Fix T-coordinate collision for conditioning

* Changed the default strength from 0.6 to 0.8

* Added reference image test and updated the frozenset

* Validated ref image, latent passing support and fixed ref image preprocessing

* Refined preprocessing with 1MP resolution cap and timestep tracking

* Updated typing, improved validation and changed the example docstring

* Style fixes

* Fixed batch inference discrepancy and addressed review comments

* Fixed a typo

Co-authored-by: Álvaro Somoza <asomoza@users.noreply.github.com>

* Apply suggestion from @asomoza

Co-authored-by: Álvaro Somoza <asomoza@users.noreply.github.com>

* Reused encoded latents and fix channel check consistency

* fixed pre-encoded latent preprocessing for source and ref images

* Apply style fixes

* Updated the docstring with the shape requirements

* Apply style fixes

* Fixed copies

Author: adi776borate

src/diffusers/__init__.py

@@ -533,6 +533,7 @@
             "EasyAnimateInpaintPipeline",
             "EasyAnimatePipeline",
             "ErnieImagePipeline",
+            "Flux2KleinInpaintPipeline",
             "Flux2KleinKVPipeline",
             "Flux2KleinPipeline",
             "Flux2Pipeline",
@@ -1317,6 +1318,7 @@
             EasyAnimateInpaintPipeline,
             EasyAnimatePipeline,
             ErnieImagePipeline,
+            Flux2KleinInpaintPipeline,
             Flux2KleinKVPipeline,
             Flux2KleinPipeline,
             Flux2Pipeline,

src/diffusers/pipelines/__init__.py

@@ -160,7 +160,12 @@
     ]
     _import_structure["bria"] = ["BriaPipeline"]
     _import_structure["bria_fibo"] = ["BriaFiboPipeline", "BriaFiboEditPipeline"]
-    _import_structure["flux2"] = ["Flux2Pipeline", "Flux2KleinPipeline", "Flux2KleinKVPipeline"]
+    _import_structure["flux2"] = [
+        "Flux2Pipeline",
+        "Flux2KleinPipeline",
+        "Flux2KleinInpaintPipeline",
+        "Flux2KleinKVPipeline",
+    ]
     _import_structure["flux"] = [
         "FluxControlPipeline",
         "FluxControlInpaintPipeline",
@@ -697,7 +702,7 @@
             FluxPriorReduxPipeline,
             ReduxImageEncoder,
         )
-        from .flux2 import Flux2KleinKVPipeline, Flux2KleinPipeline, Flux2Pipeline
+        from .flux2 import Flux2KleinInpaintPipeline, Flux2KleinKVPipeline, Flux2KleinPipeline, Flux2Pipeline
         from .glm_image import GlmImagePipeline
         from .helios import HeliosPipeline, HeliosPyramidPipeline
         from .hidream_image import HiDreamImagePipeline

src/diffusers/pipelines/flux2/__init__.py

@@ -24,6 +24,7 @@
 else:
     _import_structure["pipeline_flux2"] = ["Flux2Pipeline"]
     _import_structure["pipeline_flux2_klein"] = ["Flux2KleinPipeline"]
+    _import_structure["pipeline_flux2_klein_inpaint"] = ["Flux2KleinInpaintPipeline"]
     _import_structure["pipeline_flux2_klein_kv"] = ["Flux2KleinKVPipeline"]
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
@@ -34,6 +35,7 @@
     else:
         from .pipeline_flux2 import Flux2Pipeline
         from .pipeline_flux2_klein import Flux2KleinPipeline
+        from .pipeline_flux2_klein_inpaint import Flux2KleinInpaintPipeline
         from .pipeline_flux2_klein_kv import Flux2KleinKVPipeline
 else:
     import sys

src/diffusers/pipelines/flux2/image_processor.py

@@ -35,8 +35,12 @@ class Flux2ImageProcessor(VaeImageProcessor):
             VAE latent channels.
         do_normalize (`bool`, *optional*, defaults to `True`):
             Whether to normalize the image to [-1,1].
+        do_binarize (`bool`, *optional*, defaults to `False`):
+            Whether to binarize the image to 0/1.
         do_convert_rgb (`bool`, *optional*, defaults to be `True`):
             Whether to convert the images to RGB format.
+        do_convert_grayscale (`bool`, *optional*, defaults to be `False`):
+            Whether to convert the images to grayscale format.
     """
 
     @register_to_config
@@ -46,14 +50,18 @@ def __init__(
         vae_scale_factor: int = 16,
         vae_latent_channels: int = 32,
         do_normalize: bool = True,
+        do_binarize: bool = False,
         do_convert_rgb: bool = True,
+        do_convert_grayscale: bool = False,
     ):
         super().__init__(
             do_resize=do_resize,
             vae_scale_factor=vae_scale_factor,
             vae_latent_channels=vae_latent_channels,
             do_normalize=do_normalize,
+            do_binarize=do_binarize,
             do_convert_rgb=do_convert_rgb,
+            do_convert_grayscale=do_convert_grayscale,
         )
 
     @staticmethod