67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
import logging
|
|
|
|
logger = logging.getLogger("SmartLoRA.AnimaPatch")
|
|
|
|
_PATCHED = False
|
|
|
|
|
|
def apply():
|
|
global _PATCHED
|
|
if _PATCHED:
|
|
return
|
|
_PATCHED = True
|
|
|
|
try:
|
|
import torch.nn.functional as F
|
|
from comfy.ldm.anima.model import Attention, apply_rotary_pos_emb
|
|
except ImportError:
|
|
return
|
|
|
|
_original_forward = Attention.forward
|
|
|
|
def _dtype_safe_forward(
|
|
self,
|
|
x,
|
|
mask=None,
|
|
context=None,
|
|
position_embeddings=None,
|
|
position_embeddings_context=None,
|
|
):
|
|
context = x if context is None else context
|
|
input_shape = x.shape[:-1]
|
|
q_shape = (*input_shape, self.n_heads, self.head_dim)
|
|
context_shape = context.shape[:-1]
|
|
kv_shape = (*context_shape, self.n_heads, self.head_dim)
|
|
|
|
query_states = self.q_norm(
|
|
self.q_proj(x).view(q_shape)
|
|
).transpose(1, 2)
|
|
key_states = self.k_norm(
|
|
self.k_proj(context).view(kv_shape)
|
|
).transpose(1, 2)
|
|
value_states = self.v_proj(context).view(kv_shape).transpose(1, 2)
|
|
|
|
if position_embeddings is not None:
|
|
assert position_embeddings_context is not None
|
|
cos, sin = position_embeddings
|
|
query_states = apply_rotary_pos_emb(query_states, cos, sin)
|
|
cos, sin = position_embeddings_context
|
|
key_states = apply_rotary_pos_emb(key_states, cos, sin)
|
|
|
|
target_dtype = query_states.dtype
|
|
key_states = key_states.to(target_dtype)
|
|
value_states = value_states.to(target_dtype)
|
|
|
|
attn_output = F.scaled_dot_product_attention(
|
|
query_states, key_states, value_states, attn_mask=mask,
|
|
)
|
|
attn_output = (
|
|
attn_output.transpose(1, 2)
|
|
.reshape(*input_shape, -1)
|
|
.contiguous()
|
|
)
|
|
attn_output = self.o_proj(attn_output)
|
|
return attn_output
|
|
|
|
Attention.forward = _dtype_safe_forward
|