Fix workflow metadata and subgraph export names
This commit is contained in:
+28
-14
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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": {
|
||||
|
||||
Reference in New Issue
Block a user