38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
"""Make YuE's HF token samplers honor Comfy's cancel flag. Run with the node directory."""
|
|
import ast
|
|
import pathlib
|
|
import shutil
|
|
import sys
|
|
|
|
CHECK = ' from comfy.model_management import throw_exception_if_processing_interrupted\n throw_exception_if_processing_interrupted()\n'
|
|
|
|
def patch(source):
|
|
tree = ast.parse(source)
|
|
cls = next(n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == 'BlockTokenRangeProcessor')
|
|
method = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == '__call__')
|
|
if 'throw_exception_if_processing_interrupted' in ast.get_source_segment(source, method):
|
|
return source
|
|
lines = source.splitlines(keepends=True)
|
|
lines.insert(method.body[0].lineno - 1, CHECK)
|
|
result = ''.join(lines)
|
|
ast.parse(result)
|
|
return result
|
|
|
|
if __name__ == '__main__':
|
|
root = pathlib.Path(sys.argv[1]).resolve(strict=True)
|
|
plans = []
|
|
for name in ('infer.py', 'common.py'):
|
|
path = root / 'inference' / name
|
|
source = path.read_text(encoding='utf-8')
|
|
plans.append((path, source, patch(source)))
|
|
for path, source, result in plans:
|
|
if source == result:
|
|
print(f'Already patched: {path.name}')
|
|
continue
|
|
backup = path.with_suffix('.py.before-aigen-cancel')
|
|
if backup.exists():
|
|
raise RuntimeError(f'Backup already exists: {backup}')
|
|
shutil.copy2(path, backup)
|
|
path.write_text(result, encoding='utf-8')
|
|
print(f'Patched {path.name}; original backed up.')
|