diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py index 136ad6159..7cd449c92 100644 --- a/comfy_extras/nodes_audio.py +++ b/comfy_extras/nodes_audio.py @@ -1,5 +1,7 @@ from __future__ import annotations +import av +import tempfile import torchaudio import torch import comfy.model_management @@ -153,7 +155,8 @@ class SaveAudio: @classmethod def INPUT_TYPES(s): return {"required": { "audio": ("AUDIO", ), - "filename_prefix": ("STRING", {"default": "audio/ComfyUI"})}, + "filename_prefix": ("STRING", {"default": "audio/ComfyUI"}), + "format": (["flac", "mp3", "opus"], {"default": "flac"})}, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } @@ -163,12 +166,16 @@ class SaveAudio: OUTPUT_NODE = True CATEGORY = "audio" - - def save_audio(self, audio, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + def save_audio(self, audio, filename_prefix="ComfyUI", format="flac", prompt=None, extra_pnginfo=None): + import av + import io + import numpy as np + 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: @@ -177,25 +184,88 @@ class SaveAudio: 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}_.flac" - - buff = io.BytesIO() - torchaudio.save(buff, waveform, audio["sample_rate"], format="FLAC") - - buff = insert_or_replace_vorbis_comment(buff, metadata) - - with open(os.path.join(full_output_folder, file), 'wb') as f: - f.write(buff.getbuffer()) - + 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 + in_stream = 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("libmp3lame", 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 + 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):