Merge pull request #1 from longredzhong/copilot/add-keyword-argument-support

Add keyword argument support to transpiler

Author: longredzhong

src/comfy_script/transpile/__init__.py

@@ -12,13 +12,15 @@
 from . import passes
 
 class WorkflowToScriptTranspiler:
-    def __init__(self, workflow: str | dict, api_endpoint: str = None):
+    def __init__(self, workflow: str | dict, api_endpoint: str = None, use_keyword_args: bool = False):
         '''
         - `workflow`: Can be either in the web UI format or the API format.
+        - `use_keyword_args`: Whether to generate keyword arguments instead of positional arguments.
         '''
         if api_endpoint is not None:
             client.client = client.Client(api_endpoint)
         self.nodes_info = client.get_nodes_info()
+        self.use_keyword_args = use_keyword_args
 
         if isinstance(workflow, str):
             workflow = json.loads(workflow)
@@ -47,9 +49,13 @@ 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, use_keyword_args: bool = False) -> 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.
+        - `use_keyword_args`: Whether to generate keyword arguments instead of positional arguments.
         '''
         # TODO: webp
         workflow = None
@@ -60,21 +66,25 @@ def from_image(image: Image.Image, comfyui_api: str = None) -> WorkflowToScriptT
             workflow = image_info['prompt']
         else:
             raise ValueError('No workflow in the image')
-        return WorkflowToScriptTranspiler(workflow, comfyui_api)
+        return WorkflowToScriptTranspiler(workflow, comfyui_api, use_keyword_args=use_keyword_args)
 
     @staticmethod
-    def from_file(path: str | Path, comfyui_api: str = None) -> WorkflowToScriptTranspiler:
+    def from_file(path: str | Path, comfyui_api: str = None, use_keyword_args: bool = False) -> 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.
+        - `use_keyword_args`: Whether to generate keyword arguments instead of positional arguments.
         '''
         path = Path(path)
         if path.suffix == '.json':
             with path.open() as f:
                 workflow = json.load(f)
-            return WorkflowToScriptTranspiler(workflow, comfyui_api)
+            return WorkflowToScriptTranspiler(workflow, comfyui_api, use_keyword_args=use_keyword_args)
         else:
             image = Image.open(path)
-            return WorkflowToScriptTranspiler.from_image(image, comfyui_api)
+            return WorkflowToScriptTranspiler.from_image(image, comfyui_api, use_keyword_args=use_keyword_args)
 
     def _declare_id(self, id: str) -> str:
         if id not in self.ids:
@@ -171,6 +181,26 @@ def _keyword_args_to_positional(self, node_type: str, kwargs: dict) -> list:
                     # Optional inputs
                     args.append({'exp': 'None', 'value': None})
         return args
+    
+    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):
         G = self.G
@@ -290,7 +320,13 @@ 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)})"
+            if self.use_keyword_args:
+                # 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):

src/comfy_script/transpile/__main__.py

@@ -7,9 +7,10 @@
 @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):
+@click.option('--use-keyword-args', is_flag=True, default=False, show_default=True, help='Generate keyword arguments instead of positional arguments.')
+def cli(workflow: TextIO, api: str, runtime: bool, use_keyword_args: bool):
     workflow = workflow.read()
-    script = WorkflowToScriptTranspiler(workflow, api).to_script(runtime=runtime)
+    script = WorkflowToScriptTranspiler(workflow, api, use_keyword_args=use_keyword_args).to_script(runtime=runtime)
     print(script)
 
 cli()

tests/transpile/test_transpiler.py

@@ -69,4 +69,70 @@
 ])
 def test_workflow(workflow, script):
     with open(Path(__file__).parent / workflow) as f:
-        assert transpile.WorkflowToScriptTranspiler(f.read()).to_script() == script
+        assert transpile.WorkflowToScriptTranspiler(f.read()).to_script() == 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, seed=49328841076664, strength=0.3, normalize='true', average=None)
+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):
+    with open(Path(__file__).parent / workflow) as f:
+        assert transpile.WorkflowToScriptTranspiler(f.read(), use_keyword_args=True).to_script() == script