diff --git a/src/comfy_script/transpile/__init__.py b/src/comfy_script/transpile/__init__.py index cae4b0e..01bfb19 100644 --- a/src/comfy_script/transpile/__init__.py +++ b/src/comfy_script/transpile/__init__.py @@ -58,8 +58,9 @@ def _assign_id(self, name: str) -> str: return name def _get_input_types(self, node_type: str) -> dict: + # UI-only virtual nodes # registerNodeType: Reroute, PrimitiveNode, Note - if node_type == 'Reroute': + if node_type in passes.REROUTE_NODES: return { 'required': { '': ('*',) @@ -78,6 +79,8 @@ def _get_input_types(self, node_type: str) -> dict: } } else: + if node_type not in self.nodes_info: + raise KeyError(f'Node not found: {node_type}. If this node is an UI-only virtual node (e.g. custom Reroute, PrimitiveNode, Note nodes), you can workaround this problem by exporting the API format workflow and transpiling it. You can report this issue in https://github.com/Chaoses-Ib/ComfyScript/issues .') return self.nodes_info[node_type]['input'] def _get_widget_value_names(self, node_type: str) -> list[str]: diff --git a/src/comfy_script/transpile/passes/__init__.py b/src/comfy_script/transpile/passes/__init__.py index 849737c..9017995 100644 --- a/src/comfy_script/transpile/passes/__init__.py +++ b/src/comfy_script/transpile/passes/__init__.py @@ -10,10 +10,11 @@ class AssignContext: vars: list c: str +REROUTE_NODES = ('Reroute', 'Reroute (rgthree)') def reroute_elimination(ctx: AssignContext): - if ctx.v.type != 'Reroute': + if ctx.v.type not in REROUTE_NODES: return - assert re.fullmatch(r'(?:# )?(?:_ = Reroute\(\S+\)|(\S+) = Reroute\(\1\))\s*', ctx.c), ctx.c + assert re.fullmatch(r'(?:# )?(?:_ = [A-Za-z_0-9]+\(\S+\)|(\S+) = [A-Za-z_0-9]+\(\1\))\s*', ctx.c), ctx.c ctx.c = '' def bypass_move_elimination(ctx: AssignContext): diff --git a/tests/transpile/rgthree-comfy.json b/tests/transpile/rgthree-comfy.json new file mode 100644 index 0000000..275fc1a --- /dev/null +++ b/tests/transpile/rgthree-comfy.json @@ -0,0 +1,501 @@ +{ + "last_node_id": 12, + "last_link_id": 18, + "nodes": [ + { + "id": 5, + "type": "EmptyLatentImage", + "pos": [ + 473, + 609 + ], + "size": { + "0": 315, + "1": 106 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 2 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage" + }, + "widgets_values": [ + 512, + 512, + 1 + ] + }, + { + "id": 4, + "type": "CheckpointLoaderSimple", + "pos": [ + 26, + 474 + ], + "size": { + "0": 315, + "1": 98 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 1 + ], + "slot_index": 0 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 5, + 18 + ], + "slot_index": 1 + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 8 + ], + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "v1-5-pruned-emaonly.ckpt" + ] + }, + { + "id": 8, + "type": "VAEDecode", + "pos": [ + 1265, + 182 + ], + "size": { + "0": 210, + "1": 46 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": 8 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 9 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 9, + "type": "SaveImage", + "pos": [ + 1502, + 174 + ], + "size": { + "0": 210, + "1": 58 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "properties": {}, + "widgets_values": [ + "ComfyUI" + ] + }, + { + "id": 6, + "type": "CLIPTextEncode", + "pos": [ + 415, + 186 + ], + "size": { + "0": 422.84503173828125, + "1": 164.31304931640625 + }, + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 11 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 15 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "p" + ] + }, + { + "id": 7, + "type": "CLIPTextEncode", + "pos": [ + 413, + 389 + ], + "size": { + "0": 425.27801513671875, + "1": 180.6060791015625 + }, + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 5 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "n" + ] + }, + { + "id": 3, + "type": "KSampler", + "pos": [ + 920, + 181 + ], + "size": { + "0": 315, + "1": 262 + }, + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 16 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 13 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 2 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 7 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 0, + "randomize", + 20, + 8, + "euler", + "normal", + 1 + ] + }, + { + "id": 12, + "type": "Reroute (rgthree)", + "pos": { + "0": 853, + "1": 54, + "2": 0, + "3": 0, + "4": 0, + "5": 0, + "6": 0, + "7": 0, + "8": 0, + "9": 0 + }, + "size": [ + 40, + 30 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 15, + "dir": 3, + "has_old_label": true, + "label": " ", + "slot_index": 0 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 13, + 16 + ], + "dir": 4, + "has_old_label": true, + "label": " ", + "slot_index": 0 + } + ], + "properties": { + "resizable": false, + "size": [ + 40, + 30 + ] + } + }, + { + "id": 10, + "type": "Reroute (rgthree)", + "pos": { + "0": 360, + "1": 52, + "2": 0, + "3": 0, + "4": 0, + "5": 0, + "6": 0, + "7": 0, + "8": 0, + "9": 0 + }, + "size": [ + 40, + 30 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 18, + "dir": 3, + "has_old_label": true, + "label": " " + } + ], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 11 + ], + "dir": 4, + "has_old_label": true, + "label": " ", + "slot_index": 0 + } + ], + "properties": { + "resizable": false, + "size": [ + 40, + 30 + ] + } + } + ], + "links": [ + [ + 1, + 4, + 0, + 3, + 0, + "MODEL" + ], + [ + 2, + 5, + 0, + 3, + 3, + "LATENT" + ], + [ + 5, + 4, + 1, + 7, + 0, + "CLIP" + ], + [ + 7, + 3, + 0, + 8, + 0, + "LATENT" + ], + [ + 8, + 4, + 2, + 8, + 1, + "VAE" + ], + [ + 9, + 8, + 0, + 9, + 0, + "IMAGE" + ], + [ + 11, + 10, + 0, + 6, + 0, + "CLIP" + ], + [ + 13, + 12, + 0, + 3, + 2, + "CONDITIONING" + ], + [ + 15, + 6, + 0, + 12, + 0, + "*" + ], + [ + 16, + 12, + 0, + 3, + 1, + "CONDITIONING" + ], + [ + 18, + 4, + 1, + 10, + 0, + "*" + ] + ], + "groups": [], + "config": {}, + "extra": { + "0246.VERSION": [ + 0, + 0, + 4 + ], + "ds": { + "scale": 1, + "offset": [ + 54.39996337890625, + 239.2000274658203 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/tests/transpile/test_transpiler.py b/tests/transpile/test_transpiler.py index df57cee..2ef8143 100644 --- a/tests/transpile/test_transpiler.py +++ b/tests/transpile/test_transpiler.py @@ -34,6 +34,15 @@ image5, _, _, _ = DetailerForEachPipe(image3, segs2, 1024, True, 1024, 403808226377311, 10, 3, 'lcm', 'ddim_uniform', 0.1, 50, True, True, basic_pipe, '', 0, 1, None, None, True, 50) PreviewImage(image5) PreviewImage(image) +"""), + ('rgthree-comfy.json', +r"""model, clip, vae = CheckpointLoaderSimple('v1-5-pruned-emaonly.ckpt') +# _ = CLIPTextEncode('n', clip) +conditioning = CLIPTextEncode('p', clip) +latent = EmptyLatentImage(512, 512, 1) +latent = KSampler(model, 0, 20, 8, 'euler', 'normal', conditioning, conditioning, latent, 1) +image = VAEDecode(latent, vae) +SaveImage(image, 'ComfyUI') """) ]) def test_workflow(workflow, script):