move norm_cross defined check to outside norm_encoder_hidden_states

Author: williamberman

src/diffusers/models/attention_processor.py

@@ -312,8 +312,7 @@ def prepare_attention_mask(self, attention_mask, target_length, batch_size=None)
         return attention_mask
 
     def norm_encoder_hidden_states(self, encoder_hidden_states):
-        if self.norm_cross is None:
-            return encoder_hidden_states
+        assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states"
 
         if isinstance(self.norm_cross, nn.LayerNorm):
             encoder_hidden_states = self.norm_cross(encoder_hidden_states)
@@ -348,7 +347,7 @@ def __call__(
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
-        else:
+        elif attn.norm_cross:
             encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
@@ -448,7 +447,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
-        else:
+        elif attn.norm_cross:
             encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
@@ -499,7 +498,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
-        else:
+        elif attn.norm_cross:
             encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
@@ -543,7 +542,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
-        else:
+        elif attn.norm_cross:
             encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
@@ -595,7 +594,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
-        else:
+        elif attn.norm_cross:
             encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states) + scale * self.to_k_lora(encoder_hidden_states)
@@ -633,7 +632,7 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
-        else:
+        elif attn.norm_cross:
             encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
@@ -684,7 +683,7 @@ def __call__(self, attn: "Attention", hidden_states, encoder_hidden_states=None,
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
-        else:
+        elif attn.norm_cross:
             encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_pix2pix_zero.py

@@ -243,7 +243,7 @@ def __call__(
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
-        else:
+        elif attn.norm_cross:
             encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_sag.py

@@ -65,7 +65,7 @@ def __call__(
 
         if encoder_hidden_states is None:
             encoder_hidden_states = hidden_states
-        else:
+        elif attn.norm_cross:
             encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)