add qwen-image-edit-controlnet

Author: dimitribarbot

src/diffusers/__init__.py

@@ -512,6 +512,7 @@
             "PixArtSigmaPipeline",
             "QwenImageControlNetInpaintPipeline",
             "QwenImageControlNetPipeline",
+            "QwenImageEditControlNetPipeline",
             "QwenImageEditInpaintPipeline",
             "QwenImageEditPipeline",
             "QwenImageImg2ImgPipeline",
@@ -1166,6 +1167,7 @@
             PixArtSigmaPipeline,
             QwenImageControlNetInpaintPipeline,
             QwenImageControlNetPipeline,
+            QwenImageEditControlNetPipeline,
             QwenImageEditInpaintPipeline,
             QwenImageEditPipeline,
             QwenImageImg2ImgPipeline,

src/diffusers/models/transformers/transformer_qwenimage.py

@@ -639,7 +639,9 @@ def forward(
             if controlnet_block_samples is not None:
                 interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
                 interval_control = int(np.ceil(interval_control))
-                hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+                sample = controlnet_block_samples[index_block // interval_control]
+                sample_size = min(sample.size(1), hidden_states.size(1))
+                hidden_states[:, :sample_size] = hidden_states[:, :sample_size] + sample[:, :sample_size]
 
         # Use only the image part (hidden_states) from the dual-stream blocks
         hidden_states = self.norm_out(hidden_states, temb)

src/diffusers/pipelines/__init__.py

@@ -394,6 +394,7 @@
         "QwenImageInpaintPipeline",
         "QwenImageEditPipeline",
         "QwenImageEditInpaintPipeline",
+        "QwenImageEditControlNetPipeline",
         "QwenImageControlNetInpaintPipeline",
         "QwenImageControlNetPipeline",
     ]
@@ -717,6 +718,7 @@
         from .qwenimage import (
             QwenImageControlNetInpaintPipeline,
             QwenImageControlNetPipeline,
+            QwenImageEditControlNetPipeline,
             QwenImageEditInpaintPipeline,
             QwenImageEditPipeline,
             QwenImageImg2ImgPipeline,

src/diffusers/pipelines/qwenimage/__init__.py

@@ -27,6 +27,7 @@
     _import_structure["pipeline_qwenimage_controlnet"] = ["QwenImageControlNetPipeline"]
     _import_structure["pipeline_qwenimage_controlnet_inpaint"] = ["QwenImageControlNetInpaintPipeline"]
     _import_structure["pipeline_qwenimage_edit"] = ["QwenImageEditPipeline"]
+    _import_structure["pipeline_qwenimage_edit_controlnet"] = ["QwenImageEditControlNetPipeline"]
     _import_structure["pipeline_qwenimage_edit_inpaint"] = ["QwenImageEditInpaintPipeline"]
     _import_structure["pipeline_qwenimage_img2img"] = ["QwenImageImg2ImgPipeline"]
     _import_structure["pipeline_qwenimage_inpaint"] = ["QwenImageInpaintPipeline"]
@@ -42,6 +43,7 @@
         from .pipeline_qwenimage_controlnet import QwenImageControlNetPipeline
         from .pipeline_qwenimage_controlnet_inpaint import QwenImageControlNetInpaintPipeline
         from .pipeline_qwenimage_edit import QwenImageEditPipeline
+        from .pipeline_qwenimage_edit_controlnet import QwenImageEditControlNetPipeline
         from .pipeline_qwenimage_edit_inpaint import QwenImageEditInpaintPipeline
         from .pipeline_qwenimage_img2img import QwenImageImg2ImgPipeline
         from .pipeline_qwenimage_inpaint import QwenImageInpaintPipeline