MajoorWaldi
ac06dabd89
- Refactor 26 nodes to io.ComfyNode with define_schema and execute\n- Implement ComfyExtension subclass with async comfy_entrypoint\n- Preserve majoor_imageops synthetic namespace loader to avoid nodes module collision\n- Update and validate unit, registry and frontend test suites\n- Rebuild frontend dist files
274 lines
12 KiB
Python
274 lines
12 KiB
Python
import ast
|
|
import re
|
|
from pathlib import Path
|
|
|
|
BASE_DIR = Path(__file__).resolve().parent.parent
|
|
|
|
def format_input(name, spec, is_optional):
|
|
raw_type = spec[0] if isinstance(spec, (tuple, list)) and spec else "STRING"
|
|
opts = spec[1] if isinstance(spec, (tuple, list)) and len(spec) > 1 and isinstance(spec[1], dict) else {}
|
|
|
|
type_str = "String"
|
|
if isinstance(raw_type, (list, tuple)) and not isinstance(raw_type, str):
|
|
options_str = repr(list(raw_type))
|
|
kwargs_str = ", ".join(f"{k}={repr(v)}" for k, v in opts.items())
|
|
if kwargs_str:
|
|
kwargs_str = ", " + kwargs_str
|
|
return f'io.String.Input("{name}", options={options_str}{kwargs_str})'
|
|
|
|
type_name = str(raw_type).upper()
|
|
if type_name == "IMAGE,VIDEO":
|
|
type_str = "MultiType"
|
|
elif type_name == "COLOR" or ("COLOR" in name.lower() and type_name == "STRING"):
|
|
type_str = "Color"
|
|
elif type_name in {"BOOLEAN", "BOOL"}:
|
|
type_str = "Boolean"
|
|
elif type_name == "INT":
|
|
type_str = "Int"
|
|
elif type_name == "FLOAT":
|
|
type_str = "Float"
|
|
elif type_name == "MASK":
|
|
type_str = "Mask"
|
|
elif "IMAGE" in type_name or "VIDEO" in type_name:
|
|
type_str = "Image"
|
|
|
|
kwargs_str = ""
|
|
for k, v in opts.items():
|
|
if k == "extra_dict" and type_str in ("Image", "Video", "Mask", "MultiType"):
|
|
kwargs_str += f", extra_dict={repr(v)}"
|
|
elif k == "display":
|
|
kwargs_str += f", display_mode=io.NumberDisplay({repr(v)})"
|
|
elif k == "forceInput":
|
|
kwargs_str += f", force_input={repr(v)}"
|
|
else:
|
|
kwargs_str += f", {k}={repr(v)}"
|
|
|
|
if is_optional:
|
|
kwargs_str += ", optional=True"
|
|
|
|
if type_str == "MultiType":
|
|
return f'io.MultiType.Input("{name}", types=[io.Image, io.Video]{kwargs_str})'
|
|
return f'io.{type_str}.Input("{name}"{kwargs_str})'
|
|
|
|
|
|
class NodeMigrator(ast.NodeTransformer):
|
|
def __init__(self, node_id, display_name):
|
|
self.node_id = node_id
|
|
self.display_name = display_name
|
|
self.modified = False
|
|
|
|
# We will parse these from the legacy class attributes
|
|
self.function_name = "apply"
|
|
self.category = "image/imageops"
|
|
self.return_types = []
|
|
self.return_names = []
|
|
self.inputs_dict = {"required": {}, "optional": {}, "hidden": {}}
|
|
|
|
def extract_dict_from_ast(self, dict_node):
|
|
"""Very basic AST to dict converter for INPUT_TYPES"""
|
|
result = {}
|
|
if not isinstance(dict_node, ast.Dict):
|
|
return result
|
|
for k, v in zip(dict_node.keys, dict_node.values):
|
|
if isinstance(k, ast.Constant) and isinstance(k.value, str):
|
|
if isinstance(v, ast.Tuple) or isinstance(v, ast.List):
|
|
# Tuple (TYPE, {options})
|
|
elements = v.elts
|
|
spec = []
|
|
if len(elements) > 0:
|
|
if isinstance(elements[0], ast.Constant):
|
|
spec.append(elements[0].value)
|
|
elif isinstance(elements[0], ast.Name):
|
|
# like _BLUR_TYPES, just use string repr for now, we'll fix it manually later if needed
|
|
# actually we can just output it as raw type
|
|
spec.append(elements[0].id)
|
|
else:
|
|
spec.append("STRING")
|
|
|
|
if len(elements) > 1 and isinstance(elements[1], ast.Dict):
|
|
opts = {}
|
|
for dk, dv in zip(elements[1].keys, elements[1].values):
|
|
if isinstance(dk, ast.Constant) and isinstance(dk.value, str):
|
|
if isinstance(dv, ast.Constant):
|
|
opts[dk.value] = dv.value
|
|
elif isinstance(dv, ast.UnaryOp) and isinstance(dv.op, ast.USub) and isinstance(dv.operand, ast.Constant):
|
|
opts[dk.value] = -dv.operand.value
|
|
spec.append(opts)
|
|
result[k.value] = spec
|
|
elif isinstance(v, ast.Dict):
|
|
result[k.value] = self.extract_dict_from_ast(v)
|
|
elif isinstance(v, ast.Constant):
|
|
result[k.value] = v.value
|
|
return result
|
|
|
|
def generate_schema_code(self):
|
|
inputs_code = []
|
|
required = self.inputs_dict.get("required", {})
|
|
optional = self.inputs_dict.get("optional", {})
|
|
hidden = self.inputs_dict.get("hidden", {})
|
|
|
|
for name, spec in required.items():
|
|
if not isinstance(spec, (list, tuple)): spec = [spec]
|
|
inputs_code.append(" " + format_input(name, spec, False) + ",")
|
|
for name, spec in optional.items():
|
|
if not isinstance(spec, (list, tuple)): spec = [spec]
|
|
inputs_code.append(" " + format_input(name, spec, True) + ",")
|
|
|
|
outputs_code = []
|
|
for index, output_type in enumerate(self.return_types):
|
|
name = self.return_names[index] if index < len(self.return_names) else f"output_{index + 1}"
|
|
type_name = str(output_type).upper()
|
|
type_str = "String"
|
|
if type_name == "MASK": type_str = "Mask"
|
|
elif "IMAGE" in type_name or "VIDEO" in type_name: type_str = "Image"
|
|
elif type_name == "INT": type_str = "Int"
|
|
elif type_name == "FLOAT": type_str = "Float"
|
|
elif type_name in {"BOOLEAN", "BOOL"}: type_str = "Boolean"
|
|
outputs_code.append(f' io.{type_str}.Output("{name}", display_name="{name}"),')
|
|
|
|
hidden_code = []
|
|
for name in hidden:
|
|
hidden_code.append(f"io.Hidden.{name}")
|
|
hidden_str = f",\n hidden=[{', '.join(hidden_code)}]" if hidden_code else ""
|
|
|
|
return f"""@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="{self.node_id}",
|
|
display_name="{self.display_name}",
|
|
category="{self.category}",
|
|
inputs=[
|
|
{chr(10).join(inputs_code)}
|
|
],
|
|
outputs=[
|
|
{chr(10).join(outputs_code)}
|
|
]{hidden_str}
|
|
)
|
|
"""
|
|
|
|
def visit_ClassDef(self, node):
|
|
if not node.name.startswith("ImageOps"):
|
|
return node
|
|
|
|
# Already migrated?
|
|
has_define_schema = any(isinstance(n, ast.FunctionDef) and n.name == "define_schema" for n in node.body)
|
|
if has_define_schema:
|
|
return node
|
|
|
|
self.modified = True
|
|
|
|
# Modify bases to inherit from io.ComfyNode
|
|
if not any(isinstance(b, ast.Attribute) and b.value.id == "io" and b.attr == "ComfyNode" for b in node.bases):
|
|
node.bases = [ast.parse("io.ComfyNode").body[0].value]
|
|
|
|
# Extract metadata
|
|
for child in node.body:
|
|
if isinstance(child, ast.Assign):
|
|
targets = [t.id for t in child.targets if isinstance(t, ast.Name)]
|
|
if "CATEGORY" in targets and isinstance(child.value, ast.Constant):
|
|
self.category = child.value.value
|
|
elif "FUNCTION" in targets and isinstance(child.value, ast.Constant):
|
|
self.function_name = child.value.value
|
|
elif "RETURN_TYPES" in targets and isinstance(child.value, ast.Tuple):
|
|
self.return_types = [e.value for e in child.value.elts if isinstance(e, ast.Constant)]
|
|
elif "RETURN_NAMES" in targets and isinstance(child.value, ast.Tuple):
|
|
self.return_names = [e.value for e in child.value.elts if isinstance(e, ast.Constant)]
|
|
|
|
if isinstance(child, ast.FunctionDef) and child.name == "INPUT_TYPES":
|
|
# Find return dictionary
|
|
for stmt in child.body:
|
|
if isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.Dict):
|
|
self.inputs_dict = self.extract_dict_from_ast(stmt.value)
|
|
|
|
# Now remove legacy attributes and inject schema
|
|
new_body = []
|
|
for child in node.body:
|
|
if isinstance(child, ast.Assign):
|
|
targets = [t.id for t in child.targets if isinstance(t, ast.Name)]
|
|
if any(t in {"CATEGORY", "RETURN_TYPES", "RETURN_NAMES", "FUNCTION"} for t in targets):
|
|
continue
|
|
if isinstance(child, ast.FunctionDef):
|
|
if child.name == "INPUT_TYPES":
|
|
# inject define_schema
|
|
schema_code = self.generate_schema_code()
|
|
schema_ast = ast.parse(schema_code).body[0]
|
|
new_body.append(schema_ast)
|
|
continue
|
|
if child.name == self.function_name:
|
|
# Rename function to execute, make it classmethod, replace self with cls
|
|
child.name = "execute"
|
|
if not any(isinstance(d, ast.Name) and d.id == "classmethod" for d in child.decorator_list):
|
|
child.decorator_list.insert(0, ast.Name(id="classmethod", ctx=ast.Load()))
|
|
if child.args.args and child.args.args[0].arg == "self":
|
|
child.args.args[0].arg = "cls"
|
|
|
|
class SelfReplacer(ast.NodeTransformer):
|
|
def visit_Name(self, n):
|
|
if n.id == "self":
|
|
n.id = "cls"
|
|
return n
|
|
child = SelfReplacer().visit(child)
|
|
new_body.append(child)
|
|
|
|
node.body = new_body
|
|
return node
|
|
|
|
|
|
def process_file(filepath):
|
|
content = filepath.read_text("utf-8")
|
|
|
|
if "io.Schema" in content:
|
|
return
|
|
|
|
class_match = re.search(r"class (ImageOps[A-Za-z]+)(?:\(\w+\))?:", content)
|
|
if not class_match:
|
|
return
|
|
class_name = class_match.group(1)
|
|
|
|
display_name = re.sub(r"(?<!^)(?=[A-Z])", " ", class_name).strip()
|
|
display_name = display_name.replace("Ajust", "Color Correct").strip()
|
|
display_name = "〽️ " + display_name
|
|
|
|
try:
|
|
tree = ast.parse(content)
|
|
except Exception as e:
|
|
print(f"Error parsing {filepath.name}: {e}")
|
|
return
|
|
|
|
migrator = NodeMigrator(class_name, display_name)
|
|
new_tree = migrator.visit(tree)
|
|
|
|
if not migrator.modified:
|
|
return
|
|
|
|
# Unparse code
|
|
new_content = ast.unparse(new_tree)
|
|
|
|
# Add imports
|
|
if "from comfy_api.latest import io" not in new_content and "import io" not in new_content:
|
|
if "from __future__ import annotations" in new_content:
|
|
new_content = new_content.replace("from __future__ import annotations", "from __future__ import annotations\nfrom comfy_api.latest import io")
|
|
else:
|
|
new_content = "from comfy_api.latest import io\n" + new_content
|
|
|
|
# Fix variables that got quoted by AST extraction
|
|
new_content = re.sub(r'options="(_[A-Z_]+)"', r'options=\1', new_content)
|
|
new_content = re.sub(r'types=\[\'io.Image\', \'io.Video\'\]', r'types=[io.Image, io.Video]', new_content)
|
|
# Add **kwargs to execute if missing, but only if no ** exists
|
|
new_content = re.sub(r'def execute\(cls, ([^)]+)\):', lambda m: f'def execute(cls, {m.group(1)}):' if '**' in m.group(1) else f'def execute(cls, {m.group(1)}, **kwargs):', new_content)
|
|
|
|
filepath.write_text(new_content, "utf-8")
|
|
print(f"Migrated {filepath.name} to V3")
|
|
|
|
|
|
def run():
|
|
nodes_dir = BASE_DIR / "nodes"
|
|
for py_file in nodes_dir.glob("*.py"):
|
|
if py_file.name.startswith("_") or py_file.name == "comp.py":
|
|
continue
|
|
process_file(py_file)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run()
|