Fix dtype mismatch for ZeroGPU SDPA patching

vui/model.py

@@ -100,7 +100,7 @@ class MHA(nn.Module):
 
         dropout_p = self.dropout if self.training else 0.0
 
-        qkv = self.Wqkv(x)
+        qkv = self.Wqkv(x).to(x.dtype)
         if self.n_heads == self.n_kv_heads:
             qkv = rearrange(
                 qkv, "B T (three h d) -> B three h T d", three=3, h=self.n_heads
@@ -125,6 +125,8 @@ class MHA(nn.Module):
         if self.kv_cache is not None:
             k, v = self.kv_cache.update(input_pos, k, v)
 
+        q, k, v = q.to(x.dtype), k.to(x.dtype), v.to(x.dtype)
+
         if self.n_reps > 1:
             k = repeat_kv(k, self.n_reps)
             v = repeat_kv(v, self.n_reps)