From a17a3c7b8a91c6b94965fe39932e833cc833599d Mon Sep 17 00:00:00 2001 From: atmaranto <35581227+atmaranto@users.noreply.github.com> Date: Mon, 20 Jul 2026 05:40:24 -0500 Subject: [PATCH] Fix bug with subgraphs, implement video/audio saving --- README.md | 8 +- comfyui_to_python.py | 68 ++++++- comfyui_to_python_utils.py | 355 ++++++++++++++++++++++++++++++++++++- 3 files changed, 419 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 99f8bdf..bc01fe6 100644 --- a/README.md +++ b/README.md @@ -14,11 +14,11 @@ Now, scripts exported with SaS can be imported as modules! Once you have a scrip >>> results = exported_script.main(text5="A prompt that would be sent to the command-line arguments", queue_size=1) ``` -The first `SaveImage` node reached will instead *return* the output to the calling function. +The first save node reached (`SaveImage`, `SaveImageAdvanced`, `SaveAudio`/`SaveAudioMP3`/`SaveAudioOpus`/`SaveAudioAdvanced`, `SaveVideo`, or `SaveWEBM`) will instead *return* the output to the calling function. ### Feature: Custom Output Path -When running the exported script normally, you can now specify an `--output` option that will override the default path of `SaveImage` nodes. +When running the exported script normally, you can now specify an `--output` option that will override the default path of save nodes. If only a single image is exported by the node, then the path will be used verbatim. Otherwise, the path will be used as a prefix, and `_#####.png` will be appended to ensure uniqueness. Note that files *will be clobbered* if only one image is exported. If the path is a directory, the `SaveImage` node's `filename_prefix` will be used. @@ -29,7 +29,7 @@ like: python3 script.py --text5 "A painting of outer space" --output - --queue-size 1 > image.png ``` -Note that this will *only* work with `SaveImage` nodes. If you use alternatives, such as the `AdvancedSaveImage` node, the program will not recognize them. +This works with `SaveImage`, `SaveImageAdvanced`, `SaveAudio`, `SaveAudioMP3`, `SaveAudioOpus`, `SaveAudioAdvanced`, `SaveVideo`, and `SaveWEBM` nodes. Other save nodes (e.g. from custom node packs) will not be recognized and will save to their default location. ### Usage (Web) @@ -107,6 +107,8 @@ options: The location to save the output image. Either a file path, a directory, or - for stdout (default: the ComfyUI output directory) --disable-metadata Disables writing workflow metadata to the outputs + --random-seed-type {random,fixed,increment,decement} + How to handle the random seed (default: random) ``` #### Passing Arguments to ComfyUI diff --git a/comfyui_to_python.py b/comfyui_to_python.py index 0006ca8..7960b99 100644 --- a/comfyui_to_python.py +++ b/comfyui_to_python.py @@ -12,7 +12,11 @@ import black from comfyui_to_python_utils import import_custom_nodes, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths,\ - get_value_at_index, parse_arg, save_image_wrapper + get_value_at_index, parse_arg, gen_noise, save_image_wrapper, resolve_save_metadata,\ + save_encoded_data, encode_audio_data, audio_metadata_dict,\ + save_image_advanced_wrapper, save_audio_wrapper, save_audio_mp3_wrapper,\ + save_audio_opus_wrapper, save_audio_advanced_wrapper, save_video_wrapper,\ + save_webm_wrapper PACKAGED_FUNCTIONS = [ get_value_at_index, @@ -21,9 +25,39 @@ PACKAGED_FUNCTIONS = [ add_extra_model_paths, import_custom_nodes, save_image_wrapper, + resolve_save_metadata, + save_encoded_data, + encode_audio_data, + audio_metadata_dict, + gen_noise, + save_image_advanced_wrapper, + save_audio_wrapper, + save_audio_mp3_wrapper, + save_audio_opus_wrapper, + save_audio_advanced_wrapper, + save_video_wrapper, + save_webm_wrapper, parse_arg ] +# Output ("save") nodes that the generated script intercepts, mapped to the wrapper +# that reimplements them: when imported as a module the script returns their inputs +# instead of writing files, and when run from the command line each wrapper encodes +# its output in memory so --output can redirect it. +SAVE_NODE_WRAPPERS = { + 'SaveImage': 'save_image_wrapper', + 'SaveImageAdvanced': 'save_image_advanced_wrapper', + 'SaveAudio': 'save_audio_wrapper', + 'SaveAudioMP3': 'save_audio_mp3_wrapper', + 'SaveAudioOpus': 'save_audio_opus_wrapper', + 'SaveAudioAdvanced': 'save_audio_advanced_wrapper', + 'SaveVideo': 'save_video_wrapper', + 'SaveWEBM': 'save_webm_wrapper', +} +SAVE_NODE_TYPES = set(SAVE_NODE_WRAPPERS) +# All but the classic SaveImage use the V3 ComfyNode API, and take the prompt for metadata +V3_SAVE_NODE_TYPES = SAVE_NODE_TYPES - {'SaveImage'} + add_comfyui_directory_to_sys_path() from nodes import NODE_CLASS_MAPPINGS import nodes @@ -209,6 +243,10 @@ class CodeGenerator: # Generate class definition and inputs from the data inputs, class_type = data['inputs'], data['class_type'] + if class_type in V3_SAVE_NODE_TYPES: + # The V3 save node wrapper embeds the prompt so metadata still works + include_prompt_data = True + input_types = self.node_class_mappings[class_type].INPUT_TYPES() missing = [] for i, input in enumerate(input_types.get("required", {}).keys()): @@ -251,10 +289,10 @@ class CodeGenerator: include_prompt_data = True # Create executed variable and generate code - executed_variables[idx] = f'{self.clean_variable_name(class_type)}_{idx}' + executed_variables[idx] = f'{self.clean_variable_name(class_type)}_{idx if ":" not in idx else "subnode_" + idx.replace(":", "_")}' inputs = self.update_inputs(inputs, executed_variables) - if class_type == 'SaveImage': + if class_type in SAVE_NODE_TYPES: save_code = self.create_function_call_code(initialized_objects[class_type], class_def.FUNCTION, executed_variables[idx], is_special_function, inputs).strip() return_code = ['if __name__ != "__main__":', '\treturn dict(' + ', '.join(self.format_arg(key, value) for key, value in inputs.items()) + ')', 'else:', '\t' + save_code] @@ -286,7 +324,17 @@ class CodeGenerator: Returns: str: The generated Python code. """ - args = ', '.join(self.format_arg(key, value) for key, value in kwargs.items()) + + multi_args = {key: value for key, value in kwargs.items() if "." in key} + args = ', '.join(self.format_arg(key, value) for key, value in kwargs.items() if "." not in key) + if multi_args: + unique_keys = set(key.split(".")[0] for key, _ in multi_args.items()) + key_dicts = {key: {} for key in unique_keys} + for key, value in multi_args.items(): + base_key, sub_key = key.split(".") + key_dicts[base_key][sub_key] = value + for key, sub_dict in key_dicts.items(): + args += f', {key}={{{", ".join(f"{repr(sub_key)}: {self.format_arg(sub_key, sub_value).split("=")[1]}" for sub_key, sub_value in sub_dict.items())}}}' # Generate the Python code code = f'{variable_name} = {obj_name}.{func}({args})\n' @@ -304,11 +352,13 @@ class CodeGenerator: str: Formatted argument as a string. """ # Randomize the seed if it's a set value - if isinstance(value, int) and (key == 'noise_seed' or key == 'seed'): - return f'{key}=random.randint(1, 2**64)' + if (isinstance(key, int) and (key == 'noise_seed' or key == 'seed')): + return f'{key}=gen_noise(args.random_seed_type, {value}, {repr(value)})' elif isinstance(value, str): return f'{key}={repr(value)}' elif isinstance(value, dict) and 'variable_name' in value: + if key == 'noise_seed' or key == 'seed': + return f'{key}=gen_noise(args.random_seed_type, {value["variable_name"]}, {repr(value["variable_name"])})' return f'{key}={value["variable_name"]}' return f'{key}={value}' @@ -340,6 +390,7 @@ class CodeGenerator: argparse_code.append('parser.add_argument("--comfyui-directory", "-c", default=None, help="Where to look for ComfyUI (default: current directory)")\n') argparse_code.append(f'parser.add_argument("--output", "-o", default=None, help="The location to save the output image. Either a file path, a directory, or - for stdout (default: the ComfyUI output directory)")\n') argparse_code.append(f'parser.add_argument("--disable-metadata", action="store_true", help="Disables writing workflow metadata to the outputs")\n') + argparse_code.append(f'parser.add_argument("--random-seed-type", default="random", choices=["random", "fixed", "increment", "decement"], help="How to handle the random seed (default: random)")\n') argparse_code.append(''' comfy_args = [sys.argv[0]] if __name__ == "__main__" and "--" in sys.argv: @@ -427,7 +478,10 @@ def main(*func_args, **func_kwargs): if class_type.strip() == 'SaveImage': before = 'save_image_wrapper(' + 'ctx, ' after = ')' - + elif class_type.strip() in V3_SAVE_NODE_TYPES: + before = SAVE_NODE_WRAPPERS[class_type.strip()] + '(' + 'ctx, ' + after = ', prompt_data=PROMPT_DATA)' + if self.can_be_imported(class_type): class_code = f'{variable_name} = {before}{class_type.strip()}{after}()' else: diff --git a/comfyui_to_python_utils.py b/comfyui_to_python_utils.py index 6700652..1186c2f 100644 --- a/comfyui_to_python_utils.py +++ b/comfyui_to_python_utils.py @@ -147,7 +147,7 @@ def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: def parse_arg(s: Any, default: Any = None) -> Any: """ Parses a JSON string, returning it unchanged if the parsing fails. """ - if __name__ == "__main__" or not isinstance(s, str): + if __name__ != "__main__" or not isinstance(s, str): return s try: @@ -155,6 +155,27 @@ def parse_arg(s: Any, default: Any = None) -> Any: except json.JSONDecodeError: return s +def gen_noise(seed_type: str, seed: int, seed_key: str) -> int: + import random + + if "_seeds" not in globals(): + global _seeds + _seeds = {} + + if seed_type == "random": + return random.randint(1, 2**64) + elif seed_type == "fixed": + return seed + elif seed_type == "increment": + _seeds[seed_key] = _seeds.get(seed_key, seed) + 1 + return _seeds[seed_key] + elif seed_type == "decement": + _seeds[seed_key] = _seeds.get(seed_key, seed) - 1 + return _seeds[seed_key] + else: + raise ValueError(f"Unknown seed_type: {seed_type}") + + def save_image_wrapper(context, cls): if args.output is None: return cls @@ -235,5 +256,335 @@ def save_image_wrapper(context, cls): }) return {"ui": {"images": results}} - + return WrappedSaveImage + +def resolve_save_metadata(prompt_data): + """Returns the (prompt, extra_pnginfo) a save node would have received from the server. + + The generated script only carries the API prompt, so extra_pnginfo (which holds the + UI workflow) is never available. + """ + if args is not None and args.disable_metadata: + return None, None + return prompt_data, None + +def save_encoded_data(context, datas, filename_prefix, extension, suffix=""): + """Routes already-encoded outputs to their destination. + + `datas` holds one encoded bytes object per batch item. With --output the bytes go + straight there (stdout, an exact file path, or a directory); without it they are + written to the ComfyUI output directory the way the node itself would have. + Returns the result dicts describing what was written. + """ + import folder_paths + + output = None if args is None else args.output + + if output == "-": + if len(datas) > 1: + raise ValueError("Cannot save multiple outputs to stdout") + for data in datas: + # Hack to briefly restore stdout + if context is not None: + context.__exit__(None, None, None) + try: + sys.stdout.buffer.write(data) + sys.stdout.buffer.flush() + finally: + if context is not None: + context.__enter__() + return [] + + # An exact file path is the only case where the node's filename_prefix is ignored + if output is not None and not os.path.isdir(output): + target = os.path.abspath(output) + directory = os.path.dirname(target) or os.getcwd() + os.makedirs(directory, exist_ok=True) + stem, target_ext = os.path.splitext(target) + results = [] + for batch_number, data in enumerate(datas): + if len(datas) == 1: + path = target + else: + path = f"{stem}_{batch_number:05}{target_ext or ('.' + extension)}" + with open(path, "wb") as f: + f.write(data) + print("Saved output to", path) + results.append({ + "filename": os.path.basename(path), + "subfolder": os.path.dirname(path), + "type": "output" + }) + return results + + root = folder_paths.get_output_directory() if output is None else os.path.abspath(output) + full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(filename_prefix, root) + os.makedirs(full_output_folder, exist_ok=True) + + results = [] + for batch_number, data in enumerate(datas): + name = filename.replace("%batch_num%", str(batch_number)) + file = f"{name}_{counter:05}{suffix}.{extension}" + with open(os.path.join(full_output_folder, file), "wb") as f: + f.write(data) + print("Saved output to", os.path.join(full_output_folder, file)) + results.append({"filename": file, "subfolder": subfolder, "type": "output"}) + counter += 1 + + return results + +def encode_audio_data(audio, format="flac", quality="128k", metadata=None): + """Encodes each batch item of an AUDIO input to `format`, in memory. + + Mirrors comfy_api's AudioSaveHelper.save_audio, but yields the encoded bytes + instead of writing them to disk. + """ + from io import BytesIO + import av + + opus_rates = [8000, 12000, 16000, 24000, 48000] + + datas = [] + for waveform in audio["waveform"].cpu(): + sample_rate = audio["sample_rate"] + + # Opus only accepts a fixed set of sample rates + if format == "opus": + if sample_rate > 48000: + sample_rate = 48000 + elif sample_rate not in opus_rates: + for rate in sorted(opus_rates): + if rate > sample_rate: + sample_rate = rate + break + if sample_rate not in opus_rates: + sample_rate = 48000 + + if sample_rate != audio["sample_rate"]: + import torchaudio + waveform = torchaudio.functional.resample(waveform, audio["sample_rate"], sample_rate) + + buffer = BytesIO() + container = av.open(buffer, mode="w", format=format) + + for key, value in (metadata or {}).items(): + container.metadata[key] = value + + layout = "mono" if waveform.shape[0] == 1 else "stereo" + if format == "opus": + stream = container.add_stream("libopus", rate=sample_rate, layout=layout) + if quality.endswith("k"): + stream.bit_rate = int(quality[:-1]) * 1000 + elif format == "mp3": + stream = container.add_stream("libmp3lame", rate=sample_rate, layout=layout) + if quality == "V0": + stream.codec_context.qscale = 1 + elif quality.endswith("k"): + stream.bit_rate = int(quality[:-1]) * 1000 + else: # flac + stream = container.add_stream("flac", rate=sample_rate, layout=layout) + + frame = av.AudioFrame.from_ndarray( + waveform.movedim(0, 1).reshape(1, -1).float().numpy(), + format="flt", + layout=layout, + ) + frame.sample_rate = sample_rate + frame.pts = 0 + container.mux(stream.encode(frame)) + container.mux(stream.encode(None)) + container.close() + + datas.append(buffer.getvalue()) + + return datas + +def audio_metadata_dict(prompt_data): + """Container metadata for the audio encoders (values must be strings).""" + prompt, extra_pnginfo = resolve_save_metadata(prompt_data) + metadata = {} + if prompt is not None: + metadata["prompt"] = json.dumps(prompt) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata[x] = json.dumps(extra_pnginfo[x]) + return metadata + +def save_image_advanced_wrapper(context, cls, prompt_data=None): + """SaveImageAdvanced: encodes PNG/EXR bytes in memory, then routes them.""" + + class WrappedSaveImageAdvanced(cls): + @classmethod + def execute(cls, images, filename_prefix, format, **kwargs): + from comfy_api.latest import io as comfy_io + from comfy_extras.nodes_images import _encode_image, inject_png_metadata, inject_exr_metadata + + file_format = format["format"] + bit_depth = format["bit_depth"] + colorspace = format.get("input_color_space", "sRGB") + prompt, extra_pnginfo = resolve_save_metadata(prompt_data) + + datas = [] + for image in images: + encoded = _encode_image(image, file_format, bit_depth, colorspace) + if prompt is not None or extra_pnginfo is not None: + if file_format == "png": + encoded = inject_png_metadata(encoded, prompt, extra_pnginfo) + elif file_format == "exr": + encoded = inject_exr_metadata(encoded, prompt, extra_pnginfo, colorspace) + datas.append(encoded) + + results = save_encoded_data(context, datas, filename_prefix, file_format) + return comfy_io.NodeOutput(images, ui={"images": results}) + + return WrappedSaveImageAdvanced + +def save_audio_wrapper(context, cls, prompt_data=None): + """SaveAudio (FLAC): encodes in memory, then routes the bytes.""" + + class WrappedSaveAudio(cls): + @classmethod + def execute(cls, audio, filename_prefix="ComfyUI", format="flac", **kwargs): + from comfy_api.latest import io as comfy_io + + if audio is None: + raise ValueError("SaveAudio: input audio is None (source video may have no audio track).") + datas = encode_audio_data(audio, format=format, metadata=audio_metadata_dict(prompt_data)) + results = save_encoded_data(context, datas, filename_prefix, format) + return comfy_io.NodeOutput(audio, ui={"audio": results}) + + return WrappedSaveAudio + +def save_audio_mp3_wrapper(context, cls, prompt_data=None): + """SaveAudioMP3: encodes in memory, then routes the bytes.""" + + class WrappedSaveAudioMP3(cls): + @classmethod + def execute(cls, audio, filename_prefix="ComfyUI", format="mp3", quality="128k", **kwargs): + from comfy_api.latest import io as comfy_io + + if audio is None: + raise ValueError("SaveAudioMP3: input audio is None (source video may have no audio track).") + datas = encode_audio_data(audio, format=format, quality=quality, metadata=audio_metadata_dict(prompt_data)) + results = save_encoded_data(context, datas, filename_prefix, format) + return comfy_io.NodeOutput(audio, ui={"audio": results}) + + return WrappedSaveAudioMP3 + +def save_audio_opus_wrapper(context, cls, prompt_data=None): + """SaveAudioOpus: encodes in memory, then routes the bytes.""" + + class WrappedSaveAudioOpus(cls): + @classmethod + def execute(cls, audio, filename_prefix="ComfyUI", format="opus", quality="V3", **kwargs): + from comfy_api.latest import io as comfy_io + + if audio is None: + raise ValueError("SaveAudioOpus: input audio is None (source video may have no audio track).") + datas = encode_audio_data(audio, format=format, quality=quality, metadata=audio_metadata_dict(prompt_data)) + results = save_encoded_data(context, datas, filename_prefix, format) + return comfy_io.NodeOutput(audio, ui={"audio": results}) + + return WrappedSaveAudioOpus + +def save_audio_advanced_wrapper(context, cls, prompt_data=None): + """SaveAudioAdvanced: format and quality arrive together as a dynamic combo dict.""" + + class WrappedSaveAudioAdvanced(cls): + @classmethod + def execute(cls, audio, filename_prefix, format, **kwargs): + from comfy_api.latest import io as comfy_io + + file_format = format.get("format", "flac") + quality = format.get("quality", None) + extra = {} if quality is None else {"quality": quality} + datas = encode_audio_data(audio, format=file_format, metadata=audio_metadata_dict(prompt_data), **extra) + results = save_encoded_data(context, datas, filename_prefix, file_format) + return comfy_io.NodeOutput(audio, ui={"audio": results}) + + return WrappedSaveAudioAdvanced + +def save_video_wrapper(context, cls, prompt_data=None): + """SaveVideo: muxes into a BytesIO buffer rather than a path.""" + + class WrappedSaveVideo(cls): + @classmethod + def execute(cls, video, filename_prefix, format, codec, **kwargs): + from io import BytesIO + from comfy_api.latest import io as comfy_io, Types + + prompt, extra_pnginfo = resolve_save_metadata(prompt_data) + metadata = {} + if extra_pnginfo is not None: + metadata.update(extra_pnginfo) + if prompt is not None: + metadata["prompt"] = prompt + saved_metadata = metadata if len(metadata) > 0 else None + + buffer = BytesIO() + video.save_to( + buffer, + format=Types.VideoContainer(format), + codec=codec, + metadata=saved_metadata, + ) + + extension = Types.VideoContainer.get_extension(format) + results = save_encoded_data(context, [buffer.getvalue()], filename_prefix, extension, suffix="_") + return comfy_io.NodeOutput(video, ui={"images": results, "animated": (True,)}) + + return WrappedSaveVideo + +def save_webm_wrapper(context, cls, prompt_data=None): + """SaveWEBM: encodes the frames into a BytesIO buffer rather than a path.""" + + class WrappedSaveWEBM(cls): + @classmethod + def execute(cls, images, codec, fps, filename_prefix, crf, **kwargs): + from io import BytesIO + from fractions import Fraction + import av + import torch + from comfy_api.latest import io as comfy_io + + prompt, extra_pnginfo = resolve_save_metadata(prompt_data) + + buffer = BytesIO() + container = av.open(buffer, mode="w", format="webm") + + if prompt is not None: + container.metadata["prompt"] = json.dumps(prompt) + if extra_pnginfo is not None: + for x in extra_pnginfo: + container.metadata[x] = json.dumps(extra_pnginfo[x]) + + # Save transparency when the images carry an alpha channel (RGBA) and the codec supports it. + # vp9 -> yuva420p; other codecs have no usable alpha path, so the alpha is ignored. + save_alpha = images.shape[-1] == 4 and codec == "vp9" + + codec_map = {"vp9": "libvpx-vp9", "av1": "libsvtav1"} + stream = container.add_stream(codec_map[codec], rate=Fraction(round(fps * 1000), 1000)) + stream.width = images.shape[-2] + stream.height = images.shape[-3] + stream.pix_fmt = "yuva420p" if save_alpha else ("yuv420p10le" if codec == "av1" else "yuv420p") + stream.bit_rate = 0 + stream.options = {"crf": str(crf)} + if codec == "av1": + stream.options["preset"] = "6" + + for frame in images: + if save_alpha: + frame = av.VideoFrame.from_ndarray(torch.clamp(frame[..., :4] * 255, min=0, max=255).to(device=torch.device("cpu"), dtype=torch.uint8).numpy(), format="rgba") + else: + frame = av.VideoFrame.from_ndarray(torch.clamp(frame[..., :3] * 255, min=0, max=255).to(device=torch.device("cpu"), dtype=torch.uint8).numpy(), format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + container.mux(stream.encode()) + container.close() + + results = save_encoded_data(context, [buffer.getvalue()], filename_prefix, "webm", suffix="_") + return comfy_io.NodeOutput(images, ui={"images": results, "animated": (True,)}) + + return WrappedSaveWEBM +