mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-15 05:57:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactored Preview/Save of audios

Browse Source
This commit is contained in:
bigcat88
2025-07-23 10:13:47 +03:00
parent 333d942f30
commit bed60d6ed9
2 changed files with 161 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
comfy_extras/v3/nodes_audio.py
View File

@@ -1,18 +1,14 @@
from __future__ import annotations
import hashlib
import json
import os
from io import BytesIO
import av
import torch
import torchaudio
import comfy.model_management
import folder_paths
import node_helpers
from comfy.cli_args import args
from comfy_api.v3 import io, ui
@@ -142,8 +138,12 @@ class SaveAudioMP3(io.ComfyNodeV3):
)
@classmethod
def execute(self, audio, filename_prefix="ComfyUI", format="mp3", quality="V0") -> io.NodeOutput:
return _save_audio(self, audio, filename_prefix, format, quality)
def execute(cls, audio, filename_prefix="ComfyUI", format="mp3", quality="V0") -> io.NodeOutput:
return io.NodeOutput(
ui=ui.AudioSaveHelper.get_save_audio_ui(
audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality
)
)
class SaveAudioOpus(io.ComfyNodeV3):
@@ -163,8 +163,12 @@ class SaveAudioOpus(io.ComfyNodeV3):
)
@classmethod
def execute(self, audio, filename_prefix="ComfyUI", format="opus", quality="128k") -> io.NodeOutput:
return _save_audio(self, audio, filename_prefix, format, quality)
def execute(cls, audio, filename_prefix="ComfyUI", format="opus", quality="128k") -> io.NodeOutput:
return io.NodeOutput(
ui=ui.AudioSaveHelper.get_save_audio_ui(
audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality
)
)
class SaveAudio(io.ComfyNodeV3):
@@ -184,7 +188,9 @@ class SaveAudio(io.ComfyNodeV3):
@classmethod
def execute(cls, audio, filename_prefix="ComfyUI", format="flac") -> io.NodeOutput:
return _save_audio(cls, audio, filename_prefix, format)
return io.NodeOutput(
ui=ui.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=format)
)
class VAEDecodeAudio(io.ComfyNodeV3):
@@ -232,108 +238,6 @@ class VAEEncodeAudio(io.ComfyNodeV3):
return io.NodeOutput({"samples": vae.encode(waveform.movedim(1, -1))})
def _save_audio(cls, audio, filename_prefix="ComfyUI", format="flac", quality="128k") -> io.NodeOutput:
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
filename_prefix, folder_paths.get_output_directory()
)
# Prepare metadata dictionary
metadata = {}
if not args.disable_metadata:
if cls.hidden.prompt is not None:
metadata["prompt"] = json.dumps(cls.hidden.prompt)
if cls.hidden.extra_pnginfo is not None:
for x in cls.hidden.extra_pnginfo:
metadata[x] = json.dumps(cls.hidden.extra_pnginfo[x])
# Opus supported sample rates
OPUS_RATES = [8000, 12000, 16000, 24000, 48000]
results = []
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 output with specified format
output_buffer = 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
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)
frame = av.AudioFrame.from_ndarray(
waveform.movedim(0, 1).reshape(1, -1).float().numpy(),
format="flt",
layout="mono" if waveform.shape[0] == 1 else "stereo",
)
frame.sample_rate = sample_rate
frame.pts = 0
output_container.mux(out_stream.encode(frame))
# Flush encoder
output_container.mux(out_stream.encode(None))
# Close containers
output_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(ui.SavedResult(file, subfolder, io.FolderType.output))
counter += 1
return io.NodeOutput(ui={"audio": results})
NODES_LIST: list[type[io.ComfyNodeV3]] = [
ConditioningStableAudio,
EmptyLatentAudio,