File size: 4,560 Bytes
7d45cf1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | #!/usr/bin/env python3
"""
PuLID-Flux ComfyUI v0.25 Compatibility Patch
=============================================
PuLID's forward_orig in pulidflux.py doesn't pass attn_mask and
transformer_options to its internal block forward calls. ComfyUI v0.25
needs transformer_options for fp8 quantization patches. Without this fix,
generation crashes (OOM / exit 137) within seconds.
Usage (on the HF Space, via SSH):
# Upload via base64 pipe (SCP doesn't work on HF Spaces)
cat patch_pulid_v025.py | base64 | ssh SPACE@ssh.hf.space \\
"base64 -d > /tmp/patch_pulid.py && python3 /tmp/patch_pulid.py"
# Then restart ComfyUI to reload the patched module:
ssh SPACE@ssh.hf.space "pkill -f 'python3 main.py'; cd /ComfyUI && python3 main.py --listen 0.0.0.0 --port 8188 --output-directory /tmp/output &"
IMPORTANT: This patch is LIVE-FILESYSTEM ONLY. It reverts on factory rebuild.
Re-run after every container rebuild, or bake into Dockerfile/start.sh.
Verified: ComfyUI v0.25.0, PyTorch 2.5.1+cu121, a10g-large (46GB RAM).
"""
import re
import sys
PULID_FILE = "/home/ubuntu/ComfyUI/custom_nodes/ComfyUI-PuLID-Flux/pulidflux.py"
def patch():
with open(PULID_FILE, "r") as f:
content = f.read()
original = content
changes = []
# 1. Extract attn_mask and transformer_options from kwargs at the top of forward_orig
# Look for the kwargs extraction pattern after **kwargs in the signature
old_extract = "attn_mask = kwargs.get('attn_mask')"
if old_extract not in content:
# Add extraction if not present
# Find the kwargs parameter in forward_orig signature
content = re.sub(
r"(def forward_orig\(self.*?\*\*kwargs,?\s*\):)",
r"\1\n attn_mask = kwargs.get('attn_mask')\n transformer_options = kwargs.get('transformer_options', {})",
content,
count=1,
)
changes.append("Added attn_mask/transformer_options extraction from kwargs")
else:
# Ensure transformer_options is also extracted
if "transformer_options = kwargs.get('transformer_options'" not in content:
content = content.replace(
old_extract,
old_extract + "\n transformer_options = kwargs.get('transformer_options', {})",
)
changes.append("Added transformer_options extraction")
# 2. Patch double-block calls to pass attn_mask and transformer_options
# Pattern: img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
old_double = "img, txt = block(img=img, txt=txt, vec=vec, pe=pe)"
new_double = "img, txt = block(img=img, txt=txt, vec=vec, pe=pe, attn_mask=attn_mask, transformer_options=transformer_options)"
if old_double in content and new_double not in content:
content = content.replace(old_double, new_double)
changes.append("Patched double-block call with attn_mask + transformer_options")
# Also handle alternative double-block patterns (varies by PuLID version)
old_double_alt = "img, txt = block(img=img, txt=txt, vec=vec, pe=pe,"
if old_double_alt in content and "attn_mask=attn_mask" not in content.split(old_double_alt)[1].split(")")[0]:
# Already partially has extra kwargs - just add ours
pass # The simple replace above should handle the common case
# 3. Patch single-block calls to pass attn_mask and transformer_options
# Pattern: img = block(img, vec=vec, pe=pe)
old_single = "img = block(img, vec=vec, pe=pe)"
new_single = "img = block(img, vec=vec, pe=pe, attn_mask=attn_mask, transformer_options=transformer_options)"
if old_single in content and new_single not in content:
content = content.replace(old_single, new_single)
changes.append("Patched single-block call with attn_mask + transformer_options")
if content == original:
print("✅ No changes needed — patch already applied or pattern not found.")
print(" Verify manually:")
print(f" grep -n 'attn_mask' {PULID_FILE}")
sys.exit(0)
# Write patched file
with open(PULID_FILE, "w") as f:
f.write(content)
print(f"✅ Patched {PULID_FILE}")
for c in changes:
print(f" - {c}")
print("\n⚠️ Restart ComfyUI to reload the patched module:")
print(" pkill -f 'python3 main.py'; cd /ComfyUI && python3 main.py --listen 0.0.0.0 --port 8188 &")
print("\n⚠️ This patch is LIVE-FILESYSTEM ONLY — reverts on factory rebuild.")
if __name__ == "__main__":
patch()
|