split out save audio to separate nodes per filetype

comfy_extras/nodes_audio.py

@@ -89,6 +89,119 @@ class VAEDecodeAudio:
         audio /= std
         return ({"waveform": audio, "sample_rate": 44100}, )
 
+
+def save_audio(self, audio, filename_prefix="ComfyUI", format="flac", prompt=None, extra_pnginfo=None, quality="128k"):
+
+    filename_prefix += self.prefix_append
+    full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+    results: list[FileLocator] = []
+
+    # Prepare metadata dictionary
+    metadata = {}
+    if not args.disable_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])
+
+    # Opus supported sample rates
+    OPUS_RATES = [8000, 12000, 16000, 24000, 48000]
+
+    for (batch_number, waveform) in enumerate(audio["waveform"].cpu()):
+        filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
+        file = f"{filename_with_batch_num}_{counter:05}_.{format}"
+        output_path = os.path.join(full_output_folder, file)
+
+        # Use original sample rate initially
+        sample_rate = audio["sample_rate"]
+
+        # Handle Opus sample rate requirements
+        if format == "opus":
+            if sample_rate > 48000:
+                sample_rate = 48000
+            elif sample_rate not in OPUS_RATES:
+                # Find the next highest supported rate
+                for rate in sorted(OPUS_RATES):
+                    if rate > sample_rate:
+                        sample_rate = rate
+                        break
+                if sample_rate not in OPUS_RATES:  # Fallback if still not supported
+                    sample_rate = 48000
+
+            # Resample if necessary
+            if sample_rate != audio["sample_rate"]:
+                waveform = torchaudio.functional.resample(waveform, audio["sample_rate"], sample_rate)
+
+        # Create in-memory WAV buffer
+        wav_buffer = io.BytesIO()
+        torchaudio.save(wav_buffer, waveform, sample_rate, format="WAV")
+        wav_buffer.seek(0)  # Rewind for reading
+
+        # Use PyAV to convert and add metadata
+        input_container = av.open(wav_buffer)
+
+        # Create output with specified format
+        output_buffer = io.BytesIO()
+        output_container = av.open(output_buffer, mode='w', format=format)
+
+        # Set metadata on the container
+        for key, value in metadata.items():
+            output_container.metadata[key] = value
+
+        # Set up the output stream with appropriate properties
+        input_container.streams.audio[0]
+        if format == "opus":
+            out_stream = output_container.add_stream("libopus", rate=sample_rate)
+            if quality == "64k":
+                out_stream.bit_rate = 64000
+            elif quality == "96k":
+                out_stream.bit_rate = 96000
+            elif quality == "128k":
+                out_stream.bit_rate = 128000
+            elif quality == "192k":
+                out_stream.bit_rate = 192000
+            elif quality == "320k":
+                out_stream.bit_rate = 320000
+        elif format == "mp3":
+            out_stream = output_container.add_stream("libmp3lame", rate=sample_rate)
+            if quality == "V0":
+                #TODO i would really love to support V3 and V5 but there doesn't seem to be a way to set the qscale level, the property below is a bool
+                out_stream.codec_context.qscale = 1
+            elif quality == "128k":
+                out_stream.bit_rate = 128000
+            elif quality == "320k":
+                out_stream.bit_rate = 320000
+        else: #format == "flac":
+            out_stream = output_container.add_stream("flac", rate=sample_rate)
+
+
+        # Copy frames from input to output
+        for frame in input_container.decode(audio=0):
+            frame.pts = None  # Let PyAV handle timestamps
+            output_container.mux(out_stream.encode(frame))
+
+        # Flush encoder
+        output_container.mux(out_stream.encode(None))
+
+        # Close containers
+        output_container.close()
+        input_container.close()
+
+        # Write the output to file
+        output_buffer.seek(0)
+        with open(output_path, 'wb') as f:
+            f.write(output_buffer.getbuffer())
+
+        results.append({
+            "filename": file,
+            "subfolder": subfolder,
+            "type": self.type
+        })
+        counter += 1
+
+    return { "ui": { "audio": results } }
+
 class SaveAudio:
     def __init__(self):
         self.output_dir = folder_paths.get_output_directory()
@@ -99,115 +212,69 @@ class SaveAudio:
     def INPUT_TYPES(s):
         return {"required": { "audio": ("AUDIO", ),
                             "filename_prefix": ("STRING", {"default": "audio/ComfyUI"}),
-                            "format": (["flac", "mp3", "opus"], {"default": "flac"})},
+                            },
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                 }
 
     RETURN_TYPES = ()
-    FUNCTION = "save_audio"
+    FUNCTION = "save_flac"
 
     OUTPUT_NODE = True
 
     CATEGORY = "audio"
-    def save_audio(self, audio, filename_prefix="ComfyUI", format="flac", prompt=None, extra_pnginfo=None):
 
-        filename_prefix += self.prefix_append
+    def save_flac(self, audio, filename_prefix="ComfyUI", format="flac", prompt=None, extra_pnginfo=None):
-        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+        return save_audio(self, audio, filename_prefix, format, prompt, extra_pnginfo)
-        results: list[FileLocator] = []
 
-        # Prepare metadata dictionary
+class SaveAudioMP3:
-        metadata = {}
+    def __init__(self):
-        if not args.disable_metadata:
+        self.output_dir = folder_paths.get_output_directory()
-            if prompt is not None:
+        self.type = "output"
-                metadata["prompt"] = json.dumps(prompt)
+        self.prefix_append = ""
-            if extra_pnginfo is not None:
-                for x in extra_pnginfo:
-                    metadata[x] = json.dumps(extra_pnginfo[x])
 
-        # Opus supported sample rates
+    @classmethod
-        OPUS_RATES = [8000, 12000, 16000, 24000, 48000]
+    def INPUT_TYPES(s):
+        return {"required": { "audio": ("AUDIO", ),
+                            "filename_prefix": ("STRING", {"default": "audio/ComfyUI"}),
+                            "quality": (["V0", "128k", "320k"], {"default": "V0"}),
+                            },
+                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+                }
 
-        for (batch_number, waveform) in enumerate(audio["waveform"].cpu()):
+    RETURN_TYPES = ()
-            filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
+    FUNCTION = "save_mp3"
-            file = f"{filename_with_batch_num}_{counter:05}_.{format}"
-            output_path = os.path.join(full_output_folder, file)
 
-            # Use original sample rate initially
+    OUTPUT_NODE = True
-            sample_rate = audio["sample_rate"]
 
-            # Handle Opus sample rate requirements
+    CATEGORY = "audio"
-            if format == "opus":
-                if sample_rate > 48000:
-                    sample_rate = 48000
-                elif sample_rate not in OPUS_RATES:
-                    # Find the next highest supported rate
-                    for rate in sorted(OPUS_RATES):
-                        if rate > sample_rate:
-                            sample_rate = rate
-                            break
-                    if sample_rate not in OPUS_RATES:  # Fallback if still not supported
-                        sample_rate = 48000
 
-                # Resample if necessary
+    def save_mp3(self, audio, filename_prefix="ComfyUI", format="mp3", prompt=None, extra_pnginfo=None, quality="128k"):
-                if sample_rate != audio["sample_rate"]:
+        return save_audio(self, audio, filename_prefix, format, prompt, extra_pnginfo, quality)
-                    waveform = torchaudio.functional.resample(waveform, audio["sample_rate"], sample_rate)
 
-            # Create in-memory WAV buffer
+class SaveAudioOpus:
-            wav_buffer = io.BytesIO()
+    def __init__(self):
-            torchaudio.save(wav_buffer, waveform, sample_rate, format="WAV")
+        self.output_dir = folder_paths.get_output_directory()
-            wav_buffer.seek(0)  # Rewind for reading
+        self.type = "output"
+        self.prefix_append = ""
 
-            # Use PyAV to convert and add metadata
+    @classmethod
-            input_container = av.open(wav_buffer)
+    def INPUT_TYPES(s):
+        return {"required": { "audio": ("AUDIO", ),
+                            "filename_prefix": ("STRING", {"default": "audio/ComfyUI"}),
+                            "quality": (["64k", "96k", "128k", "192k", "320k"], {"default": "128k"}),
+                            },
+                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+                }
 
-            # Create output with specified format
+    RETURN_TYPES = ()
-            output_buffer = io.BytesIO()
+    FUNCTION = "save_opus"
-            output_container = av.open(output_buffer, mode='w', format=format)
 
-            # Set metadata on the container
+    OUTPUT_NODE = True
-            for key, value in metadata.items():
-                output_container.metadata[key] = value
 
-            # Set up the output stream with appropriate properties
+    CATEGORY = "audio"
-            input_container.streams.audio[0]
-            if format == "opus":
-                out_stream = output_container.add_stream("libopus", rate=sample_rate)
-            elif format == "mp3":
-                out_stream = output_container.add_stream("mp3", rate=sample_rate)
-            elif format == "flac":
-                out_stream = output_container.add_stream("flac", rate=sample_rate)
-            else:  # wav
-                out_stream = output_container.add_stream("pcm_s16le", rate=sample_rate)
 
-            # Set channel layout
+    def save_opus(self, audio, filename_prefix="ComfyUI", format="opus", prompt=None, extra_pnginfo=None, quality="V3"):
-            # this causes mp3 to fail
+        return save_audio(self, audio, filename_prefix, format, prompt, extra_pnginfo, quality)
-            # out_stream.layout = in_stream.layout
-
-            # Copy frames from input to output
-            for frame in input_container.decode(audio=0):
-                frame.pts = None  # Let PyAV handle timestamps
-                output_container.mux(out_stream.encode(frame))
-
-            # Flush encoder
-            output_container.mux(out_stream.encode(None))
-
-            # Close containers
-            output_container.close()
-            input_container.close()
-
-            # Write the output to file
-            output_buffer.seek(0)
-            with open(output_path, 'wb') as f:
-                f.write(output_buffer.getbuffer())
-
-            results.append({
-                "filename": file,
-                "subfolder": subfolder,
-                "type": self.type
-            })
-            counter += 1
-
-        return { "ui": { "audio": results } }
 
 class PreviewAudio(SaveAudio):
     def __init__(self):
@@ -259,7 +326,20 @@ NODE_CLASS_MAPPINGS = {
     "VAEEncodeAudio": VAEEncodeAudio,
     "VAEDecodeAudio": VAEDecodeAudio,
     "SaveAudio": SaveAudio,
+    "SaveAudioMP3": SaveAudioMP3,
+    "SaveAudioOpus": SaveAudioOpus,
     "LoadAudio": LoadAudio,
     "PreviewAudio": PreviewAudio,
     "ConditioningStableAudio": ConditioningStableAudio,
 }
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "EmptyLatentAudio": "Empty Latent Audio",
+    "VAEEncodeAudio": "VAE Encode Audio",
+    "VAEDecodeAudio": "VAE Decode Audio",
+    "PreviewAudio": "Preview Audio",
+    "LoadAudio": "Load Audio",
+    "SaveAudio": "Save Audio (FLAC)",
+    "SaveAudioMP3": "Save Audio (MP3)",
+    "SaveAudioOpus": "Save Audio (Opus)",
+}