diff --git a/README.md b/README.md index 32afd50..1b8c09c 100644 --- a/README.md +++ b/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: diff --git a/comfyui_to_python.py b/comfyui_to_python.py index df08184..6eec402 100644 --- a/comfyui_to_python.py +++ b/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.") diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py index 656a3e8..7beed03 100644 --- a/tests/test_upscale_model_loader_export.py +++ b/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()