Merge branch 'inpaint' into main

yishaoai · Sep 4, 2024 · 4ce54be · 4ce54be
2 parents 2ee3215 + 3de208c
commit 4ce54be
Show file tree

Hide file tree

Showing 8 changed files with 2,180 additions and 10 deletions.
diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
@@ -258,6 +258,8 @@
             "CogVideoXVideoToVideoPipeline",
             "CycleDiffusionPipeline",
             "FluxControlNetPipeline",
+            "FluxControlNetInpaintPipeline",
+            "FluxInpaintPipeline",
             "FluxPipeline",
             "HunyuanDiTControlNetPipeline",
             "HunyuanDiTPAGPipeline",
@@ -347,6 +349,7 @@
             "StableDiffusionUpscalePipeline",
             "StableDiffusionXLAdapterPipeline",
             "StableDiffusionXLControlNetImg2ImgPipeline",
+            "StableDiffusionXLBrushNetPipeline",
             "StableDiffusionXLControlNetInpaintPipeline",
             "StableDiffusionXLControlNetPAGImg2ImgPipeline",
             "StableDiffusionXLControlNetPAGPipeline",
@@ -557,6 +560,7 @@
             DiTTransformer2DModel,
             FluxControlNetModel,
             FluxMultiControlNetModel,
+            FluxMultiControlNetModel,
             FluxTransformer2DModel,
             HunyuanDiT2DControlNetModel,
             HunyuanDiT2DModel,
@@ -703,6 +707,8 @@
             CogVideoXVideoToVideoPipeline,
             CycleDiffusionPipeline,
             FluxControlNetPipeline,
+            FluxControlNetInpaintPipeline,
+            FluxInpaintPipeline,
             FluxPipeline,
             HunyuanDiTControlNetPipeline,
             HunyuanDiTPAGPipeline,
@@ -791,6 +797,7 @@
             StableDiffusionUpscalePipeline,
             StableDiffusionXLAdapterPipeline,
             StableDiffusionXLControlNetImg2ImgPipeline,
+            StableDiffusionXLBrushNetPipeline,
             StableDiffusionXLControlNetInpaintPipeline,
             StableDiffusionXLControlNetPAGImg2ImgPipeline,
             StableDiffusionXLControlNetPAGPipeline,

diff --git a/src/diffusers/models/transformers/transformer_flux.py b/src/diffusers/models/transformers/transformer_flux.py
@@ -454,16 +454,16 @@ def forward(
         encoder_hidden_states = self.context_embedder(encoder_hidden_states)
 
         if txt_ids.ndim == 3:
-            logger.warning(
-                "Passing `txt_ids` 3d torch.Tensor is deprecated."
-                "Please remove the batch dimension and pass it as a 2d torch Tensor"
-            )
+            #logger.warning(
+            #    "Passing `txt_ids` 3d torch.Tensor is deprecated."
+            #    "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            #)
             txt_ids = txt_ids[0]
         if img_ids.ndim == 3:
-            logger.warning(
-                "Passing `img_ids` 3d torch.Tensor is deprecated."
-                "Please remove the batch dimension and pass it as a 2d torch Tensor"
-            )
+            #logger.warning(
+            #    "Passing `img_ids` 3d torch.Tensor is deprecated."
+            #    "Please remove the batch dimension and pass it as a 2d torch Tensor"
+            #)
             img_ids = img_ids[0]
         ids = torch.cat((txt_ids, img_ids), dim=0)
         image_rotary_emb = self.pos_embed(ids)

diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py
@@ -124,7 +124,7 @@
         "AnimateDiffSparseControlNetPipeline",
         "AnimateDiffVideoToVideoPipeline",
     ]
-    _import_structure["flux"] = ["FluxPipeline", "FluxControlNetPipeline"]
+    _import_structure["flux"] = ["FluxPipeline", "FluxControlNetPipeline", "FluxControlNetInpaintPipeline", "FluxInpaintPipeline"]
     _import_structure["audioldm"] = ["AudioLDMPipeline"]
     _import_structure["audioldm2"] = [
         "AudioLDM2Pipeline",
@@ -455,6 +455,9 @@
         from .aura_flow import AuraFlowPipeline
         from .blip_diffusion import BlipDiffusionPipeline
         from .cogvideo import CogVideoXPipeline, CogVideoXVideoToVideoPipeline
+
+        from .brushnet import StableDiffusionBrushNetPipeline, StableDiffusionXLBrushNetPipeline
+
         from .controlnet import (
             BlipDiffusionControlNetPipeline,
             StableDiffusionControlNetImg2ImgPipeline,
@@ -494,7 +497,7 @@
             VersatileDiffusionTextToImagePipeline,
             VQDiffusionPipeline,
         )
-        from .flux import FluxControlNetPipeline, FluxPipeline
+        from .flux import FluxControlNetPipeline, FluxPipeline, FluxControlNetInpaintPipeline, FluxInpaintPipeline
         from .hunyuandit import HunyuanDiTPipeline
         from .i2vgen_xl import I2VGenXLPipeline
         from .kandinsky import (

diff --git a/src/diffusers/pipelines/flux/__init__.py b/src/diffusers/pipelines/flux/__init__.py
@@ -24,6 +24,8 @@
 else:
     _import_structure["pipeline_flux"] = ["FluxPipeline"]
     _import_structure["pipeline_flux_controlnet"] = ["FluxControlNetPipeline"]
+    _import_structure["pipeline_flux_controlnet_inpaint"] = ["FluxControlNetInpaintPipeline"]
+    _import_structure["pipeline_flux_inpaint"] = ["FluxInpaintPipeline"]
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
         if not (is_transformers_available() and is_torch_available()):
@@ -33,6 +35,8 @@
     else:
         from .pipeline_flux import FluxPipeline
         from .pipeline_flux_controlnet import FluxControlNetPipeline
+        from .pipeline_flux_controlnet_inpaint import FluxControlNetInpaintPipeline
+        from .pipeline_flux_inpaint import FluxInpaintPipeline
 else:
     import sys