xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-04 12:47:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix ruff

Browse Source
This commit is contained in:
thot-experiment 2025-05-09 00:38:41 -07:00
parent 6554a96d20
commit 269d141ce7
1 changed files with 17 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
comfy_extras/nodes_audio.py
View File

@ -1,7 +1,6 @@
from __future__ import annotations from __future__ import annotations
import av import av
import tempfile
import torchaudio import torchaudio
import torch import torch
import comfy.model_management import comfy.model_management
@ -167,10 +166,8 @@ class SaveAudio:
CATEGORY = "audio" CATEGORY = "audio"
def save_audio(self, audio, filename_prefix="ComfyUI", format="flac", 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 io
import numpy as np
filename_prefix += self.prefix_append filename_prefix += self.prefix_append
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
results: list[FileLocator] = [] results: list[FileLocator] = []
@ -191,10 +188,10 @@ class SaveAudio:
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
file = f"{filename_with_batch_num}_{counter:05}_.{format}" file = f"{filename_with_batch_num}_{counter:05}_.{format}"
output_path = os.path.join(full_output_folder, file) output_path = os.path.join(full_output_folder, file)
# Use original sample rate initially # Use original sample rate initially
sample_rate = audio["sample_rate"] sample_rate = audio["sample_rate"]
# Handle Opus sample rate requirements # Handle Opus sample rate requirements
if format == "opus": if format == "opus":
if sample_rate > 48000: if sample_rate > 48000:
@ -207,29 +204,29 @@ class SaveAudio:
break break
if sample_rate not in OPUS_RATES: # Fallback if still not supported if sample_rate not in OPUS_RATES: # Fallback if still not supported
sample_rate = 48000 sample_rate = 48000
# Resample if necessary # Resample if necessary
if sample_rate != audio["sample_rate"]: if sample_rate != audio["sample_rate"]:
waveform = torchaudio.functional.resample(waveform, audio["sample_rate"], sample_rate) waveform = torchaudio.functional.resample(waveform, audio["sample_rate"], sample_rate)
# Create in-memory WAV buffer # Create in-memory WAV buffer
wav_buffer = io.BytesIO() wav_buffer = io.BytesIO()
torchaudio.save(wav_buffer, waveform, sample_rate, format="WAV") torchaudio.save(wav_buffer, waveform, sample_rate, format="WAV")
wav_buffer.seek(0) # Rewind for reading wav_buffer.seek(0) # Rewind for reading
# Use PyAV to convert and add metadata # Use PyAV to convert and add metadata
input_container = av.open(wav_buffer) input_container = av.open(wav_buffer)
# Create output with specified format # Create output with specified format
output_buffer = io.BytesIO() output_buffer = io.BytesIO()
output_container = av.open(output_buffer, mode='w', format=format) output_container = av.open(output_buffer, mode='w', format=format)
# Set metadata on the container # Set metadata on the container
for key, value in metadata.items(): for key, value in metadata.items():
output_container.metadata[key] = value output_container.metadata[key] = value
# Set up the output stream with appropriate properties # Set up the output stream with appropriate properties
in_stream = input_container.streams.audio[0] input_container.streams.audio[0]
if format == "opus": if format == "opus":
out_stream = output_container.add_stream("libopus", rate=sample_rate) out_stream = output_container.add_stream("libopus", rate=sample_rate)
elif format == "mp3": elif format == "mp3":
@ -238,34 +235,34 @@ class SaveAudio:
out_stream = output_container.add_stream("flac", rate=sample_rate) out_stream = output_container.add_stream("flac", rate=sample_rate)
else: # wav else: # wav
out_stream = output_container.add_stream("pcm_s16le", rate=sample_rate) out_stream = output_container.add_stream("pcm_s16le", rate=sample_rate)
# Set channel layout # Set channel layout
# out_stream.layout = in_stream.layout # out_stream.layout = in_stream.layout
# Copy frames from input to output # Copy frames from input to output
for frame in input_container.decode(audio=0): for frame in input_container.decode(audio=0):
frame.pts = None # Let PyAV handle timestamps frame.pts = None # Let PyAV handle timestamps
output_container.mux(out_stream.encode(frame)) output_container.mux(out_stream.encode(frame))
# Flush encoder # Flush encoder
output_container.mux(out_stream.encode(None)) output_container.mux(out_stream.encode(None))
# Close containers # Close containers
output_container.close() output_container.close()
input_container.close() input_container.close()
# Write the output to file # Write the output to file
output_buffer.seek(0) output_buffer.seek(0)
with open(output_path, 'wb') as f: with open(output_path, 'wb') as f:
f.write(output_buffer.getbuffer()) f.write(output_buffer.getbuffer())
results.append({ results.append({
"filename": file, "filename": file,
"subfolder": subfolder, "subfolder": subfolder,
"type": self.type "type": self.type
}) })
counter += 1 counter += 1
return { "ui": { "audio": results } } return { "ui": { "audio": results } }
class PreviewAudio(SaveAudio): class PreviewAudio(SaveAudio):