Add CLI support for frontend workflow metadata

README.md

@@ -87,12 +87,14 @@ Options:
 ```bash
 uv run python comfyui_to_python.py \
   --input_file workflow_api.json \
+  --frontend_workflow_file workflow.json \
   --output_file workflow_api.py \
   --queue_size 10
 ```
 
 Flags:
 - `--input_file`: input workflow JSON, default `workflow_api.json`
+- `--frontend_workflow_file`: optional full frontend workflow JSON for reimportable PNG metadata
 - `--output_file`: output Python file, default `workflow_api.py`
 - `--queue_size`: default execution count in the generated script, default `10`
 
@@ -110,6 +112,10 @@ export COMFYUI_PATH=/path/to/ComfyUI
 
 The generated script is a workflow export. It does not automatically turn workflow inputs into command-line arguments.
 
+If you want generated images to drag-drop back into ComfyUI from the CLI flow, provide both:
+- the API workflow JSON via `--input_file`
+- the full frontend workflow JSON via `--frontend_workflow_file`
+
 ## Troubleshooting
 
 - `Save As Script` not visible:

comfyui_to_python.py

@@ -671,6 +671,7 @@ def run(
     input_file: str = DEFAULT_INPUT_FILE,
     output_file: str = DEFAULT_OUTPUT_FILE,
     queue_size: int = DEFAULT_QUEUE_SIZE,
+    frontend_workflow_file: str = "",
 ) -> None:
     """Generate Python code from a ComfyUI workflow_api.json file.
 
@@ -680,12 +681,19 @@ def run(
             Defaults to "workflow_api.py".
         queue_size (int): The number of times a workflow will be executed by the script.
             Defaults to 1.
+        frontend_workflow_file (str): Optional path to a full frontend workflow JSON file.
+            When provided, generated images can preserve reimportable ComfyUI workflow metadata.
 
     Returns:
         None
     """
     ComfyUItoPython(
         input_file=input_file,
+        frontend_workflow=(
+            FileHandler.read_json_file(frontend_workflow_file)
+            if frontend_workflow_file
+            else None
+        ),
         output_file=output_file,
         queue_size=queue_size,
         needs_init_custom_nodes=True,
@@ -718,6 +726,12 @@ def main() -> None:
         help="number of times the workflow will be executed by default",
         default=DEFAULT_QUEUE_SIZE,
     )
+    parser.add_argument(
+        "--frontend_workflow_file",
+        type=str,
+        help="optional path to a full frontend workflow JSON file for reimportable PNG metadata",
+        default="",
+    )
     pargs = parser.parse_args()
     run(**vars(pargs))
     print("Done.")

tests/test_upscale_model_loader_export.py

@@ -1,8 +1,11 @@
 import json
+import tempfile
 import unittest
 from io import StringIO
+from pathlib import Path
+from unittest.mock import patch
 
-from comfyui_to_python import ComfyUItoPython
+from comfyui_to_python import ComfyUItoPython, run
 
 
 class UpscaleModelLoader:
@@ -142,6 +145,65 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
         self.assertIn('"nodes": []', generated)
         self.assertNotIn('"source": "workflow_api"', generated)
 
+    def test_run_accepts_frontend_workflow_file_for_cli_metadata(self):
+        workflow = {
+            "1": {
+                "class_type": "LoadImage",
+                "inputs": {
+                    "image": "example.png",
+                },
+            }
+        }
+        frontend_workflow = {
+            "version": 0.4,
+            "last_node_id": 1,
+            "last_link_id": 0,
+            "nodes": [
+                {
+                    "id": 1,
+                    "type": "LoadImage",
+                    "pos": [0, 0],
+                    "size": [210, 60],
+                    "flags": {},
+                    "order": 0,
+                    "mode": 0,
+                    "inputs": [],
+                    "outputs": [],
+                    "properties": {},
+                    "widgets_values": ["example.png"],
+                }
+            ],
+            "links": [],
+            "groups": [],
+            "config": {},
+            "extra": {},
+        }
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmpdir_path = Path(tmpdir)
+            input_file = tmpdir_path / "workflow_api.json"
+            frontend_file = tmpdir_path / "workflow.json"
+            output_file = tmpdir_path / "workflow_api.py"
+
+            input_file.write_text(json.dumps(workflow), encoding="utf-8")
+            frontend_file.write_text(json.dumps(frontend_workflow), encoding="utf-8")
+
+            with patch(
+                "comfyui_to_python.get_node_class_mappings",
+                return_value={"LoadImage": LoadImage},
+            ), patch("comfyui_to_python.import_custom_nodes"):
+                run(
+                    input_file=str(input_file),
+                    frontend_workflow_file=str(frontend_file),
+                    output_file=str(output_file),
+                    queue_size=1,
+                )
+
+            generated = output_file.read_text(encoding="utf-8")
+
+        self.assertIn('"version": 0.4', generated)
+        self.assertIn('"type": "LoadImage"', generated)
+
 
 if __name__ == "__main__":
     unittest.main()