pipeline updated

Author: akiseakusa

src/diffusers/models/controlnets/controlnet_qwenimage_blockwise.py

@@ -25,6 +25,7 @@
 from ..cache_utils import CacheMixin
 from .controlnet import zero_module
 from ..modeling_outputs import Transformer2DModelOutput
+from ..normalization import AdaLayerNormContinuous, RMSNorm
 from ..modeling_utils import ModelMixin
 from ..transformers.qwenimage_dit import (
     QwenEmbedRope,
@@ -109,12 +110,16 @@ def __init__(
 
         # controlnet_blocks
         self.controlnet_blocks = nn.ModuleList([])
-        for _ in range(len(self.transformer_blocks)):
+        for _ in range(num_layers):
             self.controlnet_blocks.append(zero_module(BlockWiseControlBlock(self.inner_dim)))
         self.controlnet_x_embedder = zero_module(
             torch.nn.Linear(in_channels + extra_condition_channels, self.inner_dim)
         )
 
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        
         self.gradient_checkpointing = False
 
     @property
@@ -265,10 +270,6 @@ def forward(
 
         hidden_states_seq_len = hidden_states.shape[1]
         hidden_states = self.img_in(hidden_states)
-
-        # add
-        hidden_states = hidden_states + self.controlnet_x_embedder(controlnet_cond)
-
         temb = self.time_text_embed(timestep, hidden_states)
 
         image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
@@ -277,7 +278,6 @@ def forward(
         encoder_hidden_states = self.txt_norm(encoder_hidden_states)
         encoder_hidden_states = self.txt_in(encoder_hidden_states)
 
-        block_samples = ()
         for index_block, block in enumerate(self.transformer_blocks):
             if torch.is_grad_enabled() and self.gradient_checkpointing:
                 encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
@@ -298,26 +298,33 @@ def forward(
                     image_rotary_emb=image_rotary_emb,
                     joint_attention_kwargs=joint_attention_kwargs,
                 )
-            
+           
             # controlnet block
             controlnet_block_samples = ()
-            for block_sample, controlnet_block in zip(block_samples, self.controlnet_blocks):
-                block_sample = controlnet_block(block_sample)
-                controlnet_block_samples = controlnet_block_samples + (block_sample,)
-
-        # scaling
-        controlnet_block_samples = [sample * conditioning_scale for sample in controlnet_block_samples]
-        controlnet_block_samples = None if len(controlnet_block_samples) == 0 else controlnet_block_samples
-
+            #running for net
+            if controlnet_cond is not None : 
+                hidden_states_slice = hidden_states[:,:hidden_states_seq_len].clone()
+                for conditioning in controlnet_cond:
+                    controlnet_block = self.controlnet_blocks[index_block]
+                    sample = controlnet_block(hidden_states_slice, conditioning)
+                    controlnet_block_samples.append(sample)
+                    # scaling
+                controlnet_block_samples = [sample * conditioning_scale for sample in controlnet_block_samples]
+                controlnet_block_samples = None if len(controlnet_block_samples) == 0 else controlnet_block_samples
+                    
+                hidden_states[:, :hidden_states_seq_len] = hidden_states_slice + controlnet_block_samples
+        hidden_states = self.norm_out(hidden_states, controlnet_cond)
+        hidden_states = self.proj_out(hidden_states)
+        hidden_states = hidden_states[:, :hidden_states_seq_len]
         if USE_PEFT_BACKEND:
             # remove `lora_scale` from each PEFT layer
             unscale_lora_layers(self, lora_scale)
 
         if not return_dict:
-            return controlnet_block_samples
+            return hidden_states
 
         return QwenImageBlockControlNetOutput(
-            controlnet_block_samples=controlnet_block_samples,
+            controlnet_block_samples=hidden_states,
         )
 
 class QwenImageBlockwiseMultiControlNetModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin):
@@ -350,60 +357,36 @@ def forward(
         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
         return_dict: bool = True,
     ) -> Union[QwenImageBlockControlNetOutput, Tuple]:
-        # ControlNet-Union with multiple conditions
+       
+      
+        # if len(self.nets) ==1 ControlNet-Union with multiple conditions
         # only load one ControlNet for saving memories
-        if len(self.nets) == 1:
-            controlnet = self.nets[0]
-
-            for i, (image, scale) in enumerate(zip(controlnet_cond, conditioning_scale)):
-                block_samples = controlnet(
-                    hidden_states=hidden_states,
-                    controlnet_cond=image,
-                    conditioning_scale=scale,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_hidden_states_mask=encoder_hidden_states_mask,
-                    timestep=timestep,
-                    img_shapes=img_shapes,
-                    txt_seq_lens=txt_seq_lens,
-                    joint_attention_kwargs=joint_attention_kwargs,
-                    return_dict=return_dict,
-                )
-
-                # merge samples
-                if i == 0:
-                    control_block_samples = block_samples
-                else:
-                    if block_samples is not None and control_block_samples is not None:
-                        control_block_samples = [
-                            control_block_sample + block_sample
-                            for control_block_sample, block_sample in zip(control_block_samples, block_samples)
-                        ]
-         # Regular Multi-ControlNets
+        
+        # else Regular Multi-ControlNets
         # load all ControlNets into memories
-        else:
-            for i, (image, scale, controlnet) in enumerate(
-                zip(controlnet_cond, conditioning_scale, self.nets)
-            ):
-                block_samples = controlnet(
-                    hidden_states=hidden_states,
-                    controlnet_cond=image,
-                    conditioning_scale=scale,
-                    timestep=timestep,
-                    encoder_hidden_states=encoder_hidden_states,
-                    joint_attention_kwargs=joint_attention_kwargs,
-                    return_dict=return_dict,
-                )
+    
+            
+        nets_to_use = [self.nets[0]] * len(controlnet_cond) if len(self.nets) == 1 else self.nets
+        controlnet_calls = list(zip(controlnet_cond, conditioning_scale, nets_to_use))
+
+
+        # Process and merge outputs
+        for  image, scale, controlnet in controlnet_calls:
+            control_block_samples = controlnet(
+                hidden_states=hidden_states,
+                controlnet_cond=image,
+                conditioning_scale=scale,
+                timestep=timestep,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_hidden_states_mask=encoder_hidden_states_mask if len(self.nets) == 1 else None,
+                img_shapes=img_shapes if len(self.nets) == 1 else None,
+                txt_seq_lens=txt_seq_lens if len(self.nets) == 1 else None,
+                joint_attention_kwargs=joint_attention_kwargs,
+                return_dict=return_dict,
+            )
 
-                # merge samples
-                if i == 0:
-                    control_block_samples = block_samples
-                else:
-                    if block_samples is not None and control_block_samples is not None:
-                        control_block_samples = [
-                            control_block_sample + block_sample
-                            for control_block_sample, block_sample in zip(control_block_samples, block_samples)
-                        ]
-                
+       
 
         return control_block_samples
 
+        

src/diffusers/pipelines/qwenimage/pipeline_qwenimage_blockcontrolnet.py

@@ -721,8 +721,7 @@ def __call__(
                     height=control_image_.shape[3],
                     width=control_image_.shape[4],
                 ).to(dtype=prompt_embeds.dtype, device=device)
-
-                self.controlnet.img_in(control_image)
+                control_image_ = self.controlnet.controlnet_x_embedder(control_image_)
                 control_images.append(control_image_)
 
             control_image = control_images
@@ -787,6 +786,7 @@ def __call__(
 
         # 6. Denoising loop
         self.scheduler.set_begin_index(0)
+        
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 if self.interrupt:
@@ -795,7 +795,15 @@ def __call__(
                 self._current_timestep = t
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                 timestep = t.expand(latents.shape[0]).to(latents.dtype)
-
+                controlnet_keep = []
+                
+                progress = (num_inference_steps - 1 - i) / max(num_inference_steps - 1, 1)
+                keeps = [
+                    1.0 - float(progress > s + 1e-4 or progress < e - 1e-4)
+                    for s, e in zip(control_guidance_start, control_guidance_end)
+                ]
+                controlnet_keep.append(keeps[0] if isinstance(self.controlnet, QwenImageBlockwiseMultiControlNetModel) else keeps)
+    
                 if isinstance(controlnet_keep[i], list):
                     cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]
                 else: