(
self,
query_dim: int,
added_kv_proj_dim: int,
processor: "MochiAttnProcessor2_0",
heads: int = 8,
dim_head: int = 64,
dropout: float = 0.0,
bias: bool = False,
added_proj_bias: bool = True,
out_dim: int | None = None,
out_context_dim: int | None = None,
out_bias: bool = True,
context_pre_only: bool = False,
eps: float = 1e-5,
)
| 931 | |
| 932 | class MochiAttention(nn.Module): |
| 933 | def __init__( |
| 934 | self, |
| 935 | query_dim: int, |
| 936 | added_kv_proj_dim: int, |
| 937 | processor: "MochiAttnProcessor2_0", |
| 938 | heads: int = 8, |
| 939 | dim_head: int = 64, |
| 940 | dropout: float = 0.0, |
| 941 | bias: bool = False, |
| 942 | added_proj_bias: bool = True, |
| 943 | out_dim: int | None = None, |
| 944 | out_context_dim: int | None = None, |
| 945 | out_bias: bool = True, |
| 946 | context_pre_only: bool = False, |
| 947 | eps: float = 1e-5, |
| 948 | ): |
| 949 | super().__init__() |
| 950 | from .normalization import MochiRMSNorm |
| 951 | |
| 952 | self.inner_dim = out_dim if out_dim is not None else dim_head * heads |
| 953 | self.out_dim = out_dim if out_dim is not None else query_dim |
| 954 | self.out_context_dim = out_context_dim if out_context_dim else query_dim |
| 955 | self.context_pre_only = context_pre_only |
| 956 | |
| 957 | self.heads = out_dim // dim_head if out_dim is not None else heads |
| 958 | |
| 959 | self.norm_q = MochiRMSNorm(dim_head, eps, True) |
| 960 | self.norm_k = MochiRMSNorm(dim_head, eps, True) |
| 961 | self.norm_added_q = MochiRMSNorm(dim_head, eps, True) |
| 962 | self.norm_added_k = MochiRMSNorm(dim_head, eps, True) |
| 963 | |
| 964 | self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias) |
| 965 | self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias) |
| 966 | self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias) |
| 967 | |
| 968 | self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) |
| 969 | self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) |
| 970 | if self.context_pre_only is not None: |
| 971 | self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias) |
| 972 | |
| 973 | self.to_out = nn.ModuleList([]) |
| 974 | self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) |
| 975 | self.to_out.append(nn.Dropout(dropout)) |
| 976 | |
| 977 | if not self.context_pre_only: |
| 978 | self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias) |
| 979 | |
| 980 | self.processor = processor |
| 981 | |
| 982 | def forward( |
| 983 | self, |
nothing calls this directly
no test coverage detected