prepare_encoder_hidden_states -> norm_encoder_hidden_states re: @patrickvonplaten

Author: williamberman

src/diffusers/models/attention_processor.py

@@ -311,10 +311,7 @@ def prepare_attention_mask(self, attention_mask, target_length, batch_size=None)
             attention_mask = attention_mask.repeat_interleave(head_size, dim=0)
         return attention_mask
 
-    def prepare_encoder_hidden_states(self, hidden_states, encoder_hidden_states=None):
-        if encoder_hidden_states is None:
-            return hidden_states
-
+    def norm_encoder_hidden_states(self, encoder_hidden_states):
         if self.norm_cross is None:
             return encoder_hidden_states
 
@@ -349,7 +346,10 @@ def __call__(
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
         query = attn.to_q(hidden_states)
 
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -415,7 +415,10 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
         query = attn.to_q(hidden_states) + scale * self.to_q_lora(hidden_states)
         query = attn.head_to_batch_dim(query)
 
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            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)
         value = attn.to_v(encoder_hidden_states) + scale * self.to_v_lora(encoder_hidden_states)
@@ -442,7 +445,11 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
         batch_size, sequence_length, _ = hidden_states.shape
 
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
 
@@ -490,7 +497,10 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
 
         query = attn.to_q(hidden_states)
 
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -531,7 +541,10 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
 
         query = attn.to_q(hidden_states)
 
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -580,7 +593,10 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
         query = attn.to_q(hidden_states) + scale * self.to_q_lora(hidden_states)
         query = attn.head_to_batch_dim(query).contiguous()
 
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            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)
         value = attn.to_v(encoder_hidden_states) + scale * self.to_v_lora(encoder_hidden_states)
@@ -615,7 +631,10 @@ def __call__(self, attn: Attention, hidden_states, encoder_hidden_states=None, a
         dim = query.shape[-1]
         query = attn.head_to_batch_dim(query)
 
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
@@ -662,7 +681,11 @@ def __call__(self, attn: "Attention", hidden_states, encoder_hidden_states=None,
         batch_size, sequence_length, _ = hidden_states.shape
 
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            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

@@ -241,7 +241,10 @@ def __call__(
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
         query = attn.to_q(hidden_states)
 
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_sag.py

@@ -63,7 +63,10 @@ def __call__(
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
         query = attn.to_q(hidden_states)
 
-        encoder_hidden_states = attn.prepare_encoder_hidden_states(hidden_states, encoder_hidden_states)
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+        else:
+            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
 
         key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)