Merge pull request #119 from longredzhong/main

Author: Chaoses-Ib

docs/Transpiler.md

@@ -45,6 +45,8 @@ Usage: python -m comfy_script.transpile [OPTIONS] WORKFLOW
 Options:
   --api TEXT  [default: http://127.0.0.1:8188/]
   --runtime   Wrap the script with runtime imports and workflow context.
+  --args [pos|pos2orkwd|kwd]  Format node inputs as positional or keyword
+                              arguments.  [default: Pos2OrKwd]
   --help      Show this message and exit.
 ```
 

pyproject.toml

@@ -125,6 +125,9 @@ build-backend = "hatchling.build"
 # TODO: Exclude docs and examples in sdist?
 
 [tool.hatch.envs.test]
+features = [
+  "default",
+]
 dependencies = [
   "pytest"
 ]

src/comfy_script/astutil.py

@@ -16,6 +16,11 @@ class FloatEnum(str, Enum):
 import re
 import keyword
 
+if sys.version_info < (3, 10):
+    _dataclass_kw_only = {}
+else:
+    _dataclass_kw_only = {'kw_only': True}
+
 def is_xid_start(s: str) -> bool:
     return s.isidentifier()
 
@@ -167,6 +172,32 @@ def to_float_enum(id: str, values: Iterable[float], indent: str) -> (str, FloatE
     '''
     return to_enum(id, { str(v): v for v in values }, indent, FloatEnum)
 
+class ArgsFormat(StrEnum):
+    '''
+    Format function arguments as positional or keyword arguments.
+
+    Members are named in CamelCase to simplify CLI usage.
+    '''
+
+    Pos = 'pos'
+    '''Format as positional arguments.'''
+
+    Pos2OrKwd = 'pos2orkwd'
+    '''Format as positional arguments if there are 2 or fewer args;
+    otherwise, format as keyword arguments.
+    '''
+
+    Kwd = 'kwd'
+    '''Format as keyword arguments.'''
+
+    def _format_as_kwd(self, args: dict) -> bool:
+        if self == self.Pos:
+            return False
+        elif self == self.Pos2OrKwd:
+            return len(args) > 2
+        else:
+            return True
+
 def find_spec_from_fullname(fullname: str) -> importlib.ModuleSpec | None:
     module, _, _ = fullname.rpartition('.')
     while module != '':

src/comfy_script/transpile/__init__.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+import dataclasses
 import json
 from pathlib import Path
 from types import SimpleNamespace
@@ -9,10 +10,26 @@
 
 from .. import client
 from .. import astutil
+from ..astutil import ArgsFormat
 from . import passes
 
+@dataclasses.dataclass(**astutil._dataclass_kw_only)
+class FormatOptions:
+    runtime: bool = False
+    '''Whether to wrap the script with runtime imports and workflow context.'''
+
+    args: ArgsFormat = ArgsFormat.Pos2OrKwd
+    '''Format arguments as positional arguments or keyword arguments.'''
+
+_format = FormatOptions()
+'''Global default format options.'''
+
 class WorkflowToScriptTranspiler:
-    def __init__(self, workflow: str | dict, api_endpoint: str = None):
+    def __init__(
+        self,
+        workflow: str | dict,
+        api_endpoint: str | None = None,
+    ):
         '''
         - `workflow`: Can be either in the web UI format or the API format.
         '''
@@ -47,9 +64,15 @@ def __init__(self, workflow: str | dict, api_endpoint: str = None):
         self.links = links
 
     @staticmethod
-    def from_image(image: Image.Image, comfyui_api: str = None) -> WorkflowToScriptTranspiler:
+    def from_image(
+        image: Image.Image,
+        comfyui_api: str | None = None,
+    ) -> WorkflowToScriptTranspiler:
         '''
         Support PNG images generated by ComfyScript/ComfyUI.
+        
+        - `image`: A PIL image containing embedded workflow or prompt metadata.
+        - `comfyui_api`: Optional ComfyUI API endpoint to initialize the client with.
         '''
         # TODO: webp
         workflow = None
@@ -63,9 +86,15 @@ def from_image(image: Image.Image, comfyui_api: str = None) -> WorkflowToScriptT
         return WorkflowToScriptTranspiler(workflow, comfyui_api)
 
     @staticmethod
-    def from_file(path: str | Path, comfyui_api: str = None) -> WorkflowToScriptTranspiler:
+    def from_file(
+        path: str | Path,
+        comfyui_api: str | None = None,
+    ) -> WorkflowToScriptTranspiler:
         '''
         Support PNG images generated by ComfyScript/ComfyUI and workflow JSON files either in the web UI format or the API format.
+        
+        - `path`: Path to a PNG image generated by ComfyScript/ComfyUI or to a workflow JSON file in either the web UI format or the API format.
+        - `comfyui_api`: Optional ComfyUI API endpoint to use when creating the client.
         '''
         path = Path(path)
         if path.suffix == '.json':
@@ -171,8 +200,32 @@ def _keyword_args_to_positional(self, node_type: str, kwargs: dict) -> list:
                     # Optional inputs
                     args.append({'exp': 'None', 'value': None})
         return args
-
-    def _node_to_assign_st(self, node):
+    
+    def _format_args_as_keyword(self, node_type: str, args_dict: dict) -> list[str]:
+        """Format arguments as keyword arguments like 'param_name=value'."""
+        result = []
+        input_types = self._get_input_types(node_type)
+        
+        for group in 'required', 'optional':
+            group_dict: dict = input_types.get(group)
+            if group_dict is None:
+                continue
+            for name in group_dict:
+                value = args_dict.get(name)
+                if value is not None:
+                    # Format as keyword argument
+                    result.append(f"{name}={value['exp']}")
+                elif group == 'required':
+                    # Only include None for required parameters
+                    # Optional parameters with None value are skipped
+                    result.append(f"{name}=None")
+        return result
+
+    def _node_to_assign_st(
+        self,
+        node,
+        format: FormatOptions,
+    ):
         G = self.G
         links = self.links
 
@@ -290,7 +343,14 @@ def _node_to_assign_st(self, node):
         if len(vars) != 0:
             c += f"{astutil.to_assign_target_list(vars)} = "
         if mode != 4:
-            c += f"{class_id}({', '.join(arg['exp'] for arg in args)})"
+            # JS node args must be positional for passes to work
+            if format.args._format_as_kwd(args_dict) and not v.type in passes.JS_NODES:
+                # Generate keyword arguments
+                formatted_args = self._format_args_as_keyword(v.type, args_dict)
+                c += f"{class_id}({', '.join(formatted_args)})"
+            else:
+                # Generate positional arguments (default behavior)
+                c += f"{class_id}({', '.join(arg['exp'] for arg in args)})"
         else:
             # Bypass
             if len(vars) > 0 and len(vars_args_of_same_type) == len(vars):
@@ -364,11 +424,21 @@ def visit(node):
         for v in end_nodes:
             yield from visit(v)
 
-    def to_script(self, end_nodes: list[int | str] | None = None, *, runtime: bool = False) -> str:
+    def to_script(
+        self,
+        end_nodes: list[int | str] | None = None,
+        *,
+        runtime: bool | None = None,
+        format: FormatOptions = _format,
+    ) -> str:
         '''
         - `end_nodes`: The id can be of a different type than the type used by the workflow.
         - `runtime`: Whether to wrap the script with runtime imports and workflow context.
+        - `format`: See [`FormatOptions`].
         '''
+        if runtime is not None:
+            format = dataclasses.replace(format, runtime=runtime)
+
         # From leaves to roots or roots to leaves?
         # ComfyUI now executes workflows from leaves to roots, but there is a PR to change this to from roots to leaves with topological sort: https://github.com/comfyanonymous/ComfyUI/pull/931
         # To minimize future maintenance cost and suit the mental model better, we choose **from roots to leaves** too.
@@ -390,9 +460,9 @@ def to_script(self, end_nodes: list[int | str] | None = None, *, runtime: bool =
         c = ''
         for node in self._topological_generations_ordered_dfs(end_nodes):
             # TODO: Add line breaks if a node has multiple inputs
-            c += self._node_to_assign_st(self.G.nodes[node])
+            c += self._node_to_assign_st(self.G.nodes[node], format=format)
 
-        if runtime:
+        if format.runtime:
             import textwrap
 
             c = textwrap.indent(c, '    ')
@@ -408,4 +478,6 @@ def to_script(self, end_nodes: list[int | str] | None = None, *, runtime: bool =
 
 __all__ = [
     'WorkflowToScriptTranspiler',
+    'FormatOptions',
+    'ArgsFormat',
 ]

src/comfy_script/transpile/__main__.py

@@ -7,9 +7,24 @@
 @click.argument('workflow', type=click.File('r', encoding='utf-8'))
 @click.option('--api', type=click.STRING, default='http://127.0.0.1:8188/', show_default=True)
 @click.option('--runtime', is_flag=True, default=False, show_default=True, help='Wrap the script with runtime imports and workflow context.')
-def cli(workflow: TextIO, api: str, runtime: bool):
-    workflow = workflow.read()
-    script = WorkflowToScriptTranspiler(workflow, api).to_script(runtime=runtime)
+@click.option('--args',
+    type=click.Choice(ArgsFormat, case_sensitive=False),
+    default=ArgsFormat.Pos2OrKwd,
+    show_default=True,
+    help='Format node inputs as positional or keyword arguments.',
+)
+def cli(
+    workflow: TextIO,
+    api: str,
+    runtime: bool,
+    args: ArgsFormat,
+):
+    workflow_str = workflow.read()
+    format = FormatOptions(
+        runtime=runtime,
+        args=args,
+    )
+    script = WorkflowToScriptTranspiler(workflow_str, api).to_script(format=format)
     print(script)
 
 cli()

src/comfy_script/transpile/passes/__init__.py

@@ -30,6 +30,12 @@ def primitive_node_elimination(ctx: AssignContext):
     assert new_c != ctx.c, ctx.c
     ctx.c = new_c
 
+JS_NODES = [
+    *REROUTE_NODES,
+    'PrimitiveNode',
+    'Note'
+]
+
 SWITCH_NODES = {
     'HypernetworkLoader': [{'strength': 0}],
     'CLIPSetLastLayer': [{'stop_at_clip_layer': -1}],

tests/test.ps1

@@ -0,0 +1,5 @@
+
+# python -m pip install hatch
+mv __init__.py __init__.py.bak
+hatch env run -e test pytest
+mv __init__.py.bak __init__.py

tests/transpile/test_transpiler.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-import comfy_script.transpile as transpile
+from comfy_script.transpile import *
 
 @pytest.mark.parametrize('workflow, script', [
     ('default.json',
@@ -68,5 +68,73 @@
 """)
 ])
 def test_workflow(workflow, script):
+    format = FormatOptions(args=ArgsFormat.Pos)
     with open(Path(__file__).parent / workflow) as f:
-        assert transpile.WorkflowToScriptTranspiler(f.read()).to_script() == script
+        assert WorkflowToScriptTranspiler(f.read()).to_script(format=format) == script
+
+@pytest.mark.parametrize('workflow, script', [
+    ('default.json',
+r"""model, clip, vae = CheckpointLoaderSimple(ckpt_name='v1-5-pruned-emaonly.ckpt')
+conditioning = CLIPTextEncode(text='beautiful scenery nature glass bottle landscape, , purple galaxy bottle,', clip=clip)
+conditioning2 = CLIPTextEncode(text='text, watermark', clip=clip)
+latent = EmptyLatentImage(width=512, height=512, batch_size=1)
+latent = KSampler(model=model, seed=156680208700286, steps=20, cfg=8, sampler_name='euler', scheduler='normal', positive=conditioning, negative=conditioning2, latent_image=latent, denoise=1)
+image = VAEDecode(samples=latent, vae=vae)
+SaveImage(images=image, filename_prefix='ComfyUI')
+"""),
+    ('bypass.json',
+r"""image, _ = LoadImage(image='ComfyUI_temp_rcuxh_00001_.png')
+image2 = ImageScaleToSide(image=image, side_length=1024, side='Longest', upscale_method='nearest-exact', crop='disabled')
+PreviewImage(images=image2)
+image3, _ = CRUpscaleImage(image=image2, upscale_model='8x_NMKD-Superscale_150000_G.pth', mode='rescale', rescale_factor=2, resize_width=1024, resampling_method='lanczos', supersample='true', rounding_modulus=8)
+segs = ImpactMakeTileSEGS(image=image3, bbox_size=600, crop_factor=1.5, min_overlap=200, max_overlap=100, sub_batch_size_for_dilation=0, filter_segs_dilation='Reuse fast', mask_irregularity=None, irregular_mask_mode=None)
+# _ = SEGSPreview(segs=segs, alpha_mode=True, falloff=0.1, image=image3)
+image4 = image3
+PreviewImage(images=image4)
+segs2 = segs
+model, clip, vae = CheckpointLoaderSimple(ckpt_name=r'XL\turbovisionxlSuperFastXLBasedOnNew_alphaV0101Bakedvae.safetensors')
+lora_stack, _ = CRLoRAStack(switch_1='On', lora_name_1=r'xl\LCMTurboMix_LCM_Sampler.safetensors', model_weight_1=1, clip_weight_1=1, switch_2='On', lora_name_2=r'xl\xl_more_art-full_v1.safetensors', model_weight_2=1, clip_weight_2=1, switch_3='On', lora_name_3=r'xl\add-detail-xl.safetensors', model_weight_3=1, clip_weight_3=1, lora_stack=None)
+model, clip, _ = CRApplyLoRAStack(model=model, clip=clip, lora_stack=lora_stack)
+conditioning = CLIPTextEncode(text='Shot Size - extreme wide shot,( Marrakech market at night time:1.5), Moroccan young beautiful woman, smiling, exotic, (loose hijab:0.1)', clip=clip)
+conditioning2 = CLIPTextEncode(text='(worst quality, low quality, normal quality:2), blurry, depth of field, nsfw', clip=clip)
+basic_pipe = ToBasicPipe(model=model, clip=clip, vae=vae, positive=conditioning, negative=conditioning2)
+image5, _, _, _ = DetailerForEachPipe(image=image3, segs=segs2, guide_size=1024, guide_size_for=True, max_size=1024, seed=403808226377311, steps=10, cfg=3, sampler_name='lcm', scheduler='ddim_uniform', denoise=0.1, feather=50, noise_mask=True, force_inpaint=True, basic_pipe=basic_pipe, wildcard='', cycle=0, inpaint_model=1, noise_mask_feather=None, scheduler_func_opt=None, detailer_hook=True, refiner_ratio=50)
+PreviewImage(images=image5)
+PreviewImage(images=image)
+"""),
+    ('rgthree-comfy.json',
+r"""model, clip, vae = CheckpointLoaderSimple(ckpt_name='v1-5-pruned-emaonly.ckpt')
+# _ = CLIPTextEncode(text='n', clip=clip)
+conditioning = CLIPTextEncode(text='p', clip=clip)
+latent = EmptyLatentImage(width=512, height=512, batch_size=1)
+latent = KSampler(model=model, seed=0, steps=20, cfg=8, sampler_name='euler', scheduler='normal', positive=conditioning, negative=conditioning, latent_image=latent, denoise=1)
+image = VAEDecode(samples=latent, vae=vae)
+SaveImage(images=image, filename_prefix='ComfyUI')
+"""),
+    ('SplitSigmasDenoise.api.json',
+r"""noise = DisableNoise()
+width, height, _, _, _, empty_latent, _ = CRAspectRatio(width=512, height=768, aspect_ratio='custom', swap_dimensions='Off', upscale_factor=1, prescale_factor=1, batch_size=1)
+model = UNETLoader(unet_name='flux1-dev.safetensors', weight_dtype='fp8_e4m3fn')
+model = LoraLoaderModelOnly(model=model, lora_name='a.safetensors', strength_model=0.7000000000000001)
+model = LoraLoaderModelOnly(model=model, lora_name='b.safetensors', strength_model=0.7000000000000001)
+model = ModelSamplingFlux(model=model, max_shift=1.1500000000000001, base_shift=0.5, width=width, height=height)
+clip = DualCLIPLoader(clip_name1='t5.safetensors', clip_name2='clip_l.safetensors', type='flux')
+conditioning = CLIPTextEncode(text='prompt text', clip=clip)
+conditioning = FluxGuidance(conditioning=conditioning, guidance=3.5)
+guider = BasicGuider(model=model, conditioning=conditioning)
+sampler = KSamplerSelect(sampler_name='deis')
+sigmas = BasicScheduler(model=model, scheduler='beta', steps=30, denoise=1)
+sigmas, low_sigmas = SplitSigmasDenoise(sigmas=sigmas, denoise=0.4)
+noise2 = RandomNoise(noise_seed=149684926930931)
+empty_latent, _ = SamplerCustomAdvanced(noise=noise2, guider=guider, sampler=sampler, sigmas=sigmas, latent_image=empty_latent)
+empty_latent = InjectLatentNoise(latent=empty_latent, noise_seed=49328841076664, noise_strength=0.3, normalize='true')
+empty_latent, _ = SamplerCustomAdvanced(noise=noise, guider=guider, sampler=sampler, sigmas=low_sigmas, latent_image=empty_latent)
+vae = VAELoader(vae_name='ae.safetensors')
+image = VAEDecode(samples=empty_latent, vae=vae)
+SaveImage(images=image, filename_prefix='ComfyUI')
+""")
+])
+def test_workflow_with_keyword_args(workflow, script):
+    format = FormatOptions(args=ArgsFormat.Kwd)
+    with open(Path(__file__).parent / workflow) as f:
+        assert WorkflowToScriptTranspiler(f.read()).to_script(format=format) == script