Support Flux Kontext in modular (T2I and I2I)

examples/text_to_image/requirements.txt

@@ -5,4 +5,4 @@ datasets>=2.19.1
 ftfy
 tensorboard
 Jinja2
-peft>=0.17.0
+peft==0.7.0

examples/text_to_image/requirements_sdxl.txt

@@ -5,4 +5,4 @@ ftfy
 tensorboard
 Jinja2
 datasets
-peft>=0.17.0
+peft==0.7.0

src/diffusers/__init__.py

@@ -390,8 +390,6 @@
             "QwenImageAutoBlocks",
             "QwenImageEditAutoBlocks",
             "QwenImageEditModularPipeline",
-            "QwenImageEditPlusAutoBlocks",
-            "QwenImageEditPlusModularPipeline",
             "QwenImageModularPipeline",
             "StableDiffusionXLAutoBlocks",
             "StableDiffusionXLModularPipeline",
@@ -1054,8 +1052,6 @@
             QwenImageAutoBlocks,
             QwenImageEditAutoBlocks,
             QwenImageEditModularPipeline,
-            QwenImageEditPlusAutoBlocks,
-            QwenImageEditPlusModularPipeline,
             QwenImageModularPipeline,
             StableDiffusionXLAutoBlocks,
             StableDiffusionXLModularPipeline,

src/diffusers/models/transformers/transformer_ltx.py

@@ -353,9 +353,7 @@ def forward(
         norm_hidden_states = self.norm1(hidden_states)
 
         num_ada_params = self.scale_shift_table.shape[0]
-        ada_values = self.scale_shift_table[None, None].to(temb.device) + temb.reshape(
-            batch_size, temb.size(1), num_ada_params, -1
-        )
+        ada_values = self.scale_shift_table[None, None] + temb.reshape(batch_size, temb.size(1), num_ada_params, -1)
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
         norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
 

src/diffusers/models/transformers/transformer_wan.py

@@ -682,12 +682,12 @@ def forward(
         # 5. Output norm, projection & unpatchify
         if temb.ndim == 3:
             # batch_size, seq_len, inner_dim (wan 2.2 ti2v)
-            shift, scale = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2)
+            shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).chunk(2, dim=2)
             shift = shift.squeeze(2)
             scale = scale.squeeze(2)
         else:
             # batch_size, inner_dim
-            shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1)
+            shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
 
         # Move the shift and scale tensors to the same device as hidden_states.
         # When using multi-GPU inference via accelerate these will be on the

src/diffusers/models/transformers/transformer_wan_vace.py

@@ -103,7 +103,7 @@ def forward(
             control_hidden_states = control_hidden_states + hidden_states
 
         shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
-            self.scale_shift_table.to(temb.device) + temb.float()
+            self.scale_shift_table + temb.float()
         ).chunk(6, dim=1)
 
         # 1. Self-attention
@@ -361,7 +361,7 @@ def forward(
                     hidden_states = hidden_states + control_hint * scale
 
         # 6. Output norm, projection & unpatchify
-        shift, scale = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1)
+        shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
 
         # Move the shift and scale tensors to the same device as hidden_states.
         # When using multi-GPU inference via accelerate these will be on the

src/diffusers/modular_pipelines/__init__.py

@@ -46,14 +46,13 @@
     ]
     _import_structure["stable_diffusion_xl"] = ["StableDiffusionXLAutoBlocks", "StableDiffusionXLModularPipeline"]
     _import_structure["wan"] = ["WanAutoBlocks", "WanModularPipeline"]
+    _import_structure["flux"] = ["FluxAutoBlocks", "FluxModularPipeline", "FluxKontextModularPipeline"]
     _import_structure["flux"] = ["FluxAutoBlocks", "FluxModularPipeline"]
     _import_structure["qwenimage"] = [
         "QwenImageAutoBlocks",
         "QwenImageModularPipeline",
         "QwenImageEditModularPipeline",
         "QwenImageEditAutoBlocks",
-        "QwenImageEditPlusModularPipeline",
-        "QwenImageEditPlusAutoBlocks",
     ]
     _import_structure["components_manager"] = ["ComponentsManager"]
 
@@ -65,7 +64,7 @@
         from ..utils.dummy_pt_objects import *  # noqa F403
     else:
         from .components_manager import ComponentsManager
-        from .flux import FluxAutoBlocks, FluxModularPipeline
+        from .flux import FluxAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline
         from .modular_pipeline import (
             AutoPipelineBlocks,
             BlockState,
@@ -80,8 +79,6 @@
             QwenImageAutoBlocks,
             QwenImageEditAutoBlocks,
             QwenImageEditModularPipeline,
-            QwenImageEditPlusAutoBlocks,
-            QwenImageEditPlusModularPipeline,
             QwenImageModularPipeline,
         )
         from .stable_diffusion_xl import StableDiffusionXLAutoBlocks, StableDiffusionXLModularPipeline

src/diffusers/modular_pipelines/flux/__init__.py

@@ -24,15 +24,19 @@
     _import_structure["encoders"] = ["FluxTextEncoderStep"]
     _import_structure["modular_blocks"] = [
         "ALL_BLOCKS",
+        "ALL_BLOCKS_KONTEXT",
         "AUTO_BLOCKS",
+        "AUTO_BLOCKS_KONTEXT",
         "TEXT2IMAGE_BLOCKS",
         "FluxAutoBeforeDenoiseStep",
         "FluxAutoBlocks",
-        "FluxAutoBlocks",
         "FluxAutoDecodeStep",
         "FluxAutoDenoiseStep",
+        "FluxKontextAutoBeforeDenoiseStep",
+        "FluxKontextAutoBlocks",
+        "FluxKontextAutoDenoiseStep",
     ]
-    _import_structure["modular_pipeline"] = ["FluxModularPipeline"]
+    _import_structure["modular_pipeline"] = ["FluxKontextModularPipeline", "FluxModularPipeline"]
 
 if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
     try:
@@ -44,14 +48,18 @@
         from .encoders import FluxTextEncoderStep
         from .modular_blocks import (
             ALL_BLOCKS,
+            ALL_BLOCKS_KONTEXT,
             AUTO_BLOCKS,
+            AUTO_BLOCKS_KONTEXT,
             TEXT2IMAGE_BLOCKS,
             FluxAutoBeforeDenoiseStep,
             FluxAutoBlocks,
             FluxAutoDecodeStep,
             FluxAutoDenoiseStep,
+            FluxKontextAutoBeforeDenoiseStep,
+            FluxKontextAutoDenoiseStep,
         )
-        from .modular_pipeline import FluxModularPipeline
+        from .modular_pipeline import FluxKontextModularPipeline, FluxModularPipeline
 else:
     import sys