From f5c7f97f69c03a9f549347dd92110e5a57955f26 Mon Sep 17 00:00:00 2001
From: gitadmin <gitadmin@d2k.nl>
Date: Mon, 22 Jun 2026 18:42:51 +0000
Subject: [PATCH] add fixes.py: attn_mask + sdpa dtype patches

---
 fixes.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 fixes.py

diff --git a/fixes.py b/fixes.py
new file mode 100644
index 0000000..1a39932
--- /dev/null
+++ b/fixes.py
@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+"""Patches applied to x-flux-comfyui for ComfyUI compatibility."""
+import re
+
+# Fix 1: xflux DoubleStreamBlock.forward doesn't accept attn_mask/transformer_options
+# ComfyUI calls block(..., attn_mask=attn_mask, transformer_options=...) which fails
+path1 = '/app/custom_nodes/x-flux-comfyui/xflux/src/flux/modules/layers.py'
+with open(path1) as f:
+    c = f.read()
+old1 = '    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor) -> tuple[Tensor, Tensor]:\n      return self.processor(self, img, txt, vec, pe)'
+new1 = '    def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, transformer_options={}, **kwargs):\n      return self.processor(self, img, txt, vec, pe)'
+assert old1 in c, f'Fix 1 pattern not found in {path1}'
+with open(path1, 'w') as f:
+    f.write(c.replace(old1, new1))
+print('Fix 1 applied: DoubleStreamBlock.forward accepts attn_mask')
+
+# Fix 2: IPProcessor.forward sdpa dtype mismatch (ip_query=float32, ip_key/value=bfloat16)
+# Caused by projection layers loading in bfloat16 while img tensor stays float32
+path2 = '/app/custom_nodes/x-flux-comfyui/layers.py'
+with open(path2) as f:
+    c = f.read()
+old2 = '''        ip_attention = F.scaled_dot_product_attention(
+            ip_query, 
+            ip_key, 
+            ip_value, 
+            dropout_p=0.0, 
+            is_causal=False
+        )'''
+new2 = '''        ip_attention = F.scaled_dot_product_attention(
+            ip_query, 
+            ip_key.to(ip_query.dtype), 
+            ip_value.to(ip_query.dtype), 
+            dropout_p=0.0, 
+            is_causal=False
+        )'''
+assert old2 in c, f'Fix 2 pattern not found in {path2}'
+with open(path2, 'w') as f:
+    f.write(c.replace(old2, new2))
+print('Fix 2 applied: IPProcessor sdpa dtype cast')
+
+print('All fixes applied successfully.')
\ No newline at end of file