Fix workflow metadata and subgraph export names

This commit is contained in:
Peyton DeNiro
2026-03-28 16:58:47 -05:00
parent 6a9593ee70
commit 342e1b8749
2 changed files with 87 additions and 15 deletions
+28 -14
View File
@@ -194,6 +194,15 @@ class CodeGenerator:
self.node_class_mappings = node_class_mappings
self.base_node_class_mappings = base_node_class_mappings
@staticmethod
def sanitize_node_id(node_id: str) -> str:
"""Convert node IDs into variable-safe tokens without collapsing separators."""
sanitized = re.sub(r"[^a-z0-9_]", "_", str(node_id).lower().strip())
sanitized = re.sub(r"_+", "_", sanitized).strip("_")
if not sanitized:
sanitized = "node"
return sanitized
def generate_workflow(
self,
load_order: List,
@@ -284,7 +293,7 @@ class CodeGenerator:
# Create executed variable and generate code
executed_variables[idx] = (
f"{self.clean_variable_name(class_type)}_"
f"{self.clean_variable_name(str(idx))}"
f"{self.sanitize_node_id(str(idx))}"
)
inputs = self.update_inputs(inputs, executed_variables)
@@ -315,7 +324,7 @@ class CodeGenerator:
special_functions_code,
code,
workflow_data,
metadata_workflow_data or workflow_data,
metadata_workflow_data,
queue_size,
custom_nodes,
)
@@ -386,7 +395,7 @@ class CodeGenerator:
speical_functions_code: List[str],
code: List[str],
workflow_data: Dict,
metadata_workflow_data: Dict,
metadata_workflow_data: Dict | None,
queue_size: int,
custom_nodes=False,
) -> str:
@@ -397,13 +406,23 @@ class CodeGenerator:
speical_functions_code (List[str]): A list of special functions code strings.
code (List[str]): A list of code strings.
workflow_data (Dict): The API workflow data used for runtime prompt execution.
metadata_workflow_data (Dict): The workflow metadata to embed into saved outputs.
metadata_workflow_data (Dict | None): The workflow metadata to embed into saved outputs.
queue_size (int): Number of photos that will be generated by the script.
custom_nodes (bool): Whether to include custom nodes in the code.
Returns:
str: Generated final code as a string.
"""
if metadata_workflow_data is None:
extra_pnginfo_code = "extra_pnginfo = None"
else:
extra_pnginfo_code = (
"extra_pnginfo = "
'{"workflow": json.loads('
+ json.dumps(json.dumps(metadata_workflow_data))
+ ")}"
)
# Get the source code of the utils functions as a string
func_strings = []
for func in [
@@ -429,12 +448,7 @@ class CodeGenerator:
"\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n",
f"workflow = json.loads({json.dumps(json.dumps(workflow_data))})",
"prompt = json.loads(json.dumps(workflow))",
(
"extra_pnginfo = "
'{"workflow": json.loads('
+ json.dumps(json.dumps(metadata_workflow_data))
+ ")}"
),
extra_pnginfo_code,
]
)
# Check if custom nodes should be included
@@ -628,12 +642,12 @@ class ComfyUItoPython:
else:
data = json.loads(self.workflow)
frontend_workflow_data = None
metadata_workflow_data = None
if self.frontend_workflow:
if isinstance(self.frontend_workflow, str):
frontend_workflow_data = json.loads(self.frontend_workflow)
metadata_workflow_data = json.loads(self.frontend_workflow)
else:
frontend_workflow_data = self.frontend_workflow
metadata_workflow_data = self.frontend_workflow
# Step 2: Initialize extra/custom nodes when requested or when the workflow references
# a node class that is not currently loaded in the runtime.
@@ -657,7 +671,7 @@ class ComfyUItoPython:
generated_code = code_generator.generate_workflow(
load_order,
data,
frontend_workflow_data,
metadata_workflow_data,
queue_size=self.queue_size,
)
+59 -1
View File
@@ -57,6 +57,23 @@ class ImageUpscaleWithModel:
return (image,)
class MergeImages:
CATEGORY = "image"
FUNCTION = "merge"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"left": ("IMAGE",),
"right": ("IMAGE",),
}
}
def merge(self, left, right):
return (left, right)
class UpscaleModelLoaderExportTest(unittest.TestCase):
def test_upscale_workflow_uses_direct_upscale_model_loader_init(self):
workflow = {
@@ -102,7 +119,7 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
self.assertIn("loadimage = LoadImage()", generated)
self.assertIn("upscalemodelloader = UpscaleModelLoader()", generated)
self.assertIn(
"imageupscalewithmodel__3 = imageupscalewithmodel.upscale(",
"imageupscalewithmodel_3 = imageupscalewithmodel.upscale(",
generated,
)
self.assertNotIn('NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()', generated)
@@ -169,6 +186,47 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
self.assertIn("extra_pnginfo = None", generated)
self.assertNotIn('"workflow": json.loads(', generated)
def test_export_preserves_unique_variable_names_for_subgraph_node_ids(self):
workflow = {
"1:23": {
"class_type": "LoadImage",
"inputs": {
"image": "left.png",
},
},
"12:3": {
"class_type": "LoadImage",
"inputs": {
"image": "right.png",
},
},
"20": {
"class_type": "MergeImages",
"inputs": {
"left": ["1:23", 0],
"right": ["12:3", 0],
},
},
}
output = StringIO()
ComfyUItoPython(
workflow=json.dumps(workflow),
output_file=output,
node_class_mappings={
"LoadImage": LoadImage,
"MergeImages": MergeImages,
},
)
generated = output.getvalue()
self.assertIn("loadimage_1_23 = loadimage.load_image(", generated)
self.assertIn("loadimage_12_3 = loadimage.load_image(", generated)
self.assertIn("mergeimages_20 = mergeimages.merge(", generated)
self.assertIn("left=get_value_at_index(loadimage_1_23, 0)", generated)
self.assertIn("right=get_value_at_index(loadimage_12_3, 0)", generated)
def test_run_accepts_frontend_workflow_file_for_cli_metadata(self):
workflow = {
"1": {