51 lines
1.5 KiB
Python
51 lines
1.5 KiB
Python
import os
|
|
# ComfyUI imports
|
|
import folder_paths
|
|
# Local imports
|
|
from .util import increment_filename_no_overwrite
|
|
|
|
class AudioSinkNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"wav_bytes": ("WAV_BYTES",),
|
|
"filename": ("STRING", {"default": "AudioSink"}),
|
|
"save_format": (['wav'],),
|
|
"save_output": ("BOOLEAN", {"default": True}),
|
|
"overwrite_existing": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = tuple()
|
|
#RETURN_NAMES = ("image_output_name",)
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "audio"
|
|
|
|
FUNCTION = "save_audio"
|
|
|
|
def save_audio(self, wav_bytes, filename, save_format, save_output, overwrite_existing):
|
|
if not save_output:
|
|
return (1, )
|
|
|
|
filename = filename.strip()
|
|
assert filename, "Filename cannot be empty"
|
|
|
|
base_output_dir = folder_paths.get_output_directory()
|
|
assert os.path.exists(base_output_dir), f"Output directory {base_output_dir} does not exist"
|
|
output_dir = os.path.join(base_output_dir, 'audio')
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
output_path = os.path.join(output_dir, filename + '.' + save_format)
|
|
|
|
if os.path.exists(output_path) and not overwrite_existing:
|
|
output_path = increment_filename_no_overwrite(output_path)
|
|
|
|
with open(output_path, 'wb') as f:
|
|
f.write(wav_bytes)
|
|
return (0, )
|