72 lines
2.6 KiB
Python
72 lines
2.6 KiB
Python
import os
|
|
import re
|
|
|
|
import folder_paths
|
|
|
|
from .util import get_wav_bytes_from_file, audio_extensions, video_extensions
|
|
|
|
|
|
accepted_extensions = audio_extensions + video_extensions
|
|
|
|
class LoadBatchNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base_folder": (['input', 'root'],),
|
|
"path": ("STRING", {"default": ""}),
|
|
"regex_match": ("STRING", {"default": ""}),
|
|
"regex_exclude": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = "audio"
|
|
|
|
RETURN_TYPES = ("WAV_BYTES_BATCH", "INT", "LIST[STRING]")
|
|
RETURN_NAMES = ("wav_bytes_batch", "count", "filenames")
|
|
|
|
FUNCTION = "load_audio"
|
|
|
|
def load_audio(self, base_folder, path, regex_match, regex_exclude):
|
|
base_dir = None
|
|
if base_folder == 'input':
|
|
base_dir = folder_paths.get_input_directory()
|
|
elif base_folder == 'root':
|
|
base_dir = ""
|
|
path_dir = os.path.join(base_dir, path)
|
|
assert os.path.exists(path_dir), f"Path {path_dir} does not exist"
|
|
assert os.path.isdir(path_dir), f"Path {path_dir} is not a directory"
|
|
files = os.listdir(path_dir)
|
|
files = filter(lambda f: os.path.isfile(os.path.join(path_dir, f)), files)
|
|
files = filter(lambda f: os.path.splitext(f)[1].lstrip('.') in accepted_extensions, files)
|
|
if regex_match:
|
|
files = filter(lambda f: re.search(regex_match, f), files)
|
|
if regex_exclude:
|
|
files = filter(lambda f: not re.search(regex_exclude, f), files)
|
|
files = list(files)
|
|
audio_batch = []
|
|
for f in files:
|
|
filepath = os.path.join(path_dir, f)
|
|
item = {
|
|
"filename": f,
|
|
"file_path": filepath,
|
|
}
|
|
# Need to initialize get_audio this way because of the way references are handled.
|
|
# Something to do with closures and late binding.
|
|
# Else get_audio will always refer to the last item in the loop.
|
|
item["get_audio"] = lambda x=filepath: get_wav_bytes_from_file(x)
|
|
|
|
audio_batch.append(item)
|
|
return (audio_batch, len(audio_batch), list(files))
|
|
|
|
# @classmethod
|
|
# def IS_CHANGED(cls, video, **kwargs):
|
|
# image_path = folder_paths.get_annotated_filepath(video)
|
|
# return calculate_file_hash(image_path)
|
|
|
|
# @classmethod
|
|
# def VALIDATE_INPUTS(cls, audio, **kwargs):
|
|
# if not folder_paths.exists_annotated_filepath(audio):
|
|
# return "Invalid video file: {}".format(audio)
|
|
# return True
|