Add prompt to prompt

cccntu · Oct 16, 2022 · b43bc04 · b43bc04
1 parent a593a66
commit b43bc04
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py
@@ -250,8 +250,8 @@ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.
     def forward(self, x, context=None, mask=None):
         q = self.to_q(x)
         context = default(context, x)
-        k = self.to_k(context)
-        v = self.to_v(context)
+        k = self.to_k(context if isinstance(context, torch.Tensor) else context[0])
+        v = self.to_v(context if isinstance(context, torch.Tensor) else context[1])
 
         b, _, _ = q.shape
         q, k, v = map(
@@ -331,8 +331,8 @@ def forward(self, hidden_states, context=None, mask=None):
 
         query = self.to_q(hidden_states)
         context = context if context is not None else hidden_states
-        key = self.to_k(context)
-        value = self.to_v(context)
+        key = self.to_k(context if isinstance(context, torch.Tensor) else context[0])
+        value = self.to_v(context if isinstance(context, torch.Tensor) else context[1])
 
         dim = query.shape[-1]