36 lines
2.0 KiB
Python
36 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Patches applied to x-flux-comfyui for ComfyUI compatibility."""
|
|
import os, glob
|
|
|
|
# Fix 1: xflux DoubleStreamBlock.forward doesn't accept attn_mask/transformer_options
|
|
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)
|
|
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(\n ip_query, \n ip_key, \n ip_value, \n dropout_p=0.0, \n is_causal=False\n )'
|
|
new2 = ' ip_attention = F.scaled_dot_product_attention(\n ip_query.float(), \n ip_key.float(), \n ip_value.float(), \n dropout_p=0.0, \n is_causal=False\n )'
|
|
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 cast to float32')
|
|
|
|
# Delete all stale .pyc files so Python recompiles from patched .py
|
|
for pyc in glob.glob('/app/custom_nodes/x-flux-comfyui/**/*.pyc', recursive=True):
|
|
os.remove(pyc)
|
|
print(f'Removed stale pyc: {pyc}')
|
|
|
|
# Recompile from patched sources
|
|
import compileall
|
|
compileall.compile_dir('/app/custom_nodes/x-flux-comfyui', quiet=True)
|
|
print('Recompiled all x-flux-comfyui sources')
|
|
print('All fixes applied successfully.') |