From 16a0225b41959dd69f882039912d12cab67ccd92 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Tue, 14 Jan 2025 21:32:02 +0100 Subject: [PATCH] Specify precision for causal attention mask --- .../models/text_encoder/clip_encoder/clip_sdpa_attention.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py b/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py index 36937ba..e067d1d 100644 --- a/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py +++ b/src/mflux/models/text_encoder/clip_encoder/clip_sdpa_attention.py @@ -2,6 +2,8 @@ import mlx.core as mx from mlx import nn from mlx.core.fast import scaled_dot_product_attention +from mflux import Config + class CLIPSdpaAttention(nn.Module): head_dimension = 64 @@ -16,6 +18,8 @@ class CLIPSdpaAttention(nn.Module): 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: + causal_attention_mask = causal_attention_mask.astype(Config.precision) + query = self.q_proj(hidden_states) key = self.k_proj(hidden_states) value = self.v_proj(hidden_states)