Fix bug with subgraphs, implement video/audio saving

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

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:

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
+