add fixes.py: attn_mask + sdpa dtype patches
This commit is contained in:
@@ -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.')
|
||||
Reference in New Issue
Block a user