first pass at opus and mp3 as well as migrating flac to pyav

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):