fix(transpile): fallback to API format if web UI format breaks the tr…

…anspiler (#42)
Chaoses-Ib · Jun 14, 2024 · 7e0f37e · 7e0f37e
1 parent 867aee6
commit 7e0f37e
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/src/comfy_script/nodes/__init__.py b/src/comfy_script/nodes/__init__.py
@@ -15,6 +15,7 @@ def setup():
     import inspect
     import traceback
     import json
+    from warnings import warn
 
     import PIL.PngImagePlugin
 
@@ -83,11 +84,6 @@ def chunks(self):
                 if print_script is True or save_script is True:
                     if 'workflow' in self._texts or 'prompt' in self._texts:
                         try:
-                            if not prefer_api_format:
-                                workflow, zip = self._texts['workflow' if 'workflow' in self._texts else 'prompt']
-                            else:
-                                workflow, zip = self._texts['prompt']
-
                             end_nodes = None
                             # TODO: UNIQUE_ID
                             frame = inspect.currentframe()
@@ -98,7 +94,20 @@ def chunks(self):
                             else:
                                 print('ComfyScript: Failed to resolve the id of current node.')
 
-                            comfy_script = transpile.WorkflowToScriptTranspiler(workflow).to_script(end_nodes)
+                            if not prefer_api_format and 'workflow' in self._texts:
+                                workflow, zip = self._texts['workflow']
+                                try:
+                                    comfy_script = transpile.WorkflowToScriptTranspiler(workflow).to_script(end_nodes)
+                                except Exception:
+                                    traceback.print_exc()
+
+                                    warn('Failed to transpile workflow in web UI format, fallbacking to API format... If this fixes the problem, you can suppress this warning by setting `transpile.hook.prefer_api_format` to `true` in a `settings.toml`, see `settings.example.toml` for details. Please report this issue in https://github.com/Chaoses-Ib/ComfyScript/issues .')
+                                    workflow, zip = self._texts['prompt']
+                                    comfy_script = transpile.WorkflowToScriptTranspiler(workflow).to_script(end_nodes)
+                            else:
+                                workflow, zip = self._texts['prompt']
+                                comfy_script = transpile.WorkflowToScriptTranspiler(workflow).to_script(end_nodes)
+
                             if print_script is True:
                                 # TODO: Syntax highlight?
                                 print('ComfyScript:', comfy_script, sep='\n')

diff --git a/src/comfy_script/transpile/__init__.py b/src/comfy_script/transpile/__init__.py
@@ -80,7 +80,7 @@ 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. If you are using the transpiler hook, you can config it with `transpile.hook.prefer_api_format` in a `settings.toml`, see `settings.example.toml` for details. You can report this issue in https://github.com/Chaoses-Ib/ComfyScript/issues .')
+                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. Please 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]: