monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Specify precision for causal attention mask

Browse Source
This commit is contained in:
filipstrand 2025-01-14 21:32:02 +01:00
parent 3b0f5a4516
commit 16a0225b41
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py
View File

@ -2,6 +2,8 @@ import mlx.core as mx
from mlx import nn from mlx import nn
from mlx.core.fast import scaled_dot_product_attention from mlx.core.fast import scaled_dot_product_attention
from mflux import Config
class CLIPSdpaAttention(nn.Module): class CLIPSdpaAttention(nn.Module):
head_dimension = 64 head_dimension = 64
@ -16,6 +18,8 @@ class CLIPSdpaAttention(nn.Module):
self.out_proj = nn.Linear(input_dims=768, output_dims=768) self.out_proj = nn.Linear(input_dims=768, output_dims=768)
def __call__(self, hidden_states: mx.array, causal_attention_mask: mx.array) -> mx.array: def __call__(self, hidden_states: mx.array, causal_attention_mask: mx.array) -> mx.array:
causal_attention_mask = causal_attention_mask.astype(Config.precision)
query = self.q_proj(hidden_states) query = self.q_proj(hidden_states)
key = self.k_proj(hidden_states) key = self.k_proj(hidden_states)
value = self.v_proj(hidden_states) value = self.v_proj(hidden_states)