Update src/utils.py

src/utils.py

@@ -600,8 +600,10 @@ def register_faster_forward(model, mod = '50ls'):
                 timestep_cond: Optional[torch.Tensor] = None,
                 attention_mask: Optional[torch.Tensor] = None,
                 cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+                added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, # ADDED
                 down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
                 mid_block_additional_residual: Optional[torch.Tensor] = None,
+                down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None, # ADDED
                 return_dict: bool = True,
             ) -> Union[UNet2DConditionOutput, Tuple]:
                 r"""
@@ -739,18 +741,27 @@ def register_faster_forward(model, mod = '50ls'):
                     down_block_res_samples = (sample,)
                     for downsample_block in self.down_blocks:
                         if hasattr(downsample_block, "has_cross_attention") and downsample_block.has_cross_attention:
+                            #added for t2i adapters
+                            additional_residuals = {}
+                            if down_intrablock_additional_residuals is not None and len(down_intrablock_additional_residuals) > 0:
+                                additional_residuals["additional_residuals"] = down_intrablock_additional_residuals.pop(0)
                             sample, res_samples = downsample_block(
                                 hidden_states=sample,
                                 temb=emb,
                                 encoder_hidden_states=encoder_hidden_states,
                                 attention_mask=attention_mask,
                                 cross_attention_kwargs=cross_attention_kwargs,
+                                **additional_residuals
                             )
                         else:
                             sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
+                            if down_intrablock_additional_residuals is not None and len(down_intrablock_additional_residuals) > 0:
+                                sample += down_intrablock_additional_residuals.pop(0)
 
                         down_block_res_samples += res_samples
 
+                    
+
                     if down_block_additional_residuals is not None:
                         new_down_block_res_samples = ()
 
@@ -762,15 +773,31 @@ def register_faster_forward(model, mod = '50ls'):
 
                         down_block_res_samples = new_down_block_res_samples
 
+                    # Handle ControlNet additional residuals
+                    if down_block_additional_residuals is not None:
+                        new_down_block_res_samples = ()
+                        for down_block_res_sample, down_block_additional_residual in zip(
+                                down_block_res_samples, down_block_additional_residuals
+                        ):
+                            down_block_res_sample = down_block_res_sample + down_block_additional_residual
+                            new_down_block_res_samples += (down_block_res_sample,)
+                        down_block_res_samples = new_down_block_res_samples
+                        
                     # 4. mid
                     if self.mid_block is not None:
-                        sample = self.mid_block(
-                            sample,
-                            emb,
-                            encoder_hidden_states=encoder_hidden_states,
-                            attention_mask=attention_mask,
-                            cross_attention_kwargs=cross_attention_kwargs,
-                        )
+                        if hasattr(self.mid_block, "has_cross_attention") and self.mid_block.has_cross_attention:
+                            sample = self.mid_block(
+                                sample,
+                                emb,
+                                encoder_hidden_states=encoder_hidden_states,
+                                attention_mask=attention_mask,
+                                cross_attention_kwargs=cross_attention_kwargs,
+                            )
+                        else:
+                            sample = self.mid_block(sample, emb)
+                        #Handle T2I-Adapter-XL
+                        if down_intrablock_additional_residuals is not None and len(down_intrablock_additional_residuals) > 0 and sample.shape == down_intrablock_additional_residuals[0].shape:
+                            sample += down_intrablock_additional_residuals.pop(0)
 
                     if mid_block_additional_residual is not None:
                         sample = sample + mid_block_additional_residual