Add content-type filter method to folder_paths

This commit is contained in:
christian-byrne 2024-07-18 14:28:30 -07:00
parent 374e093e09
commit b696e28b7c
2 changed files with 25 additions and 8 deletions

View File

@ -182,17 +182,10 @@ class PreviewAudio(SaveAudio):
} }
class LoadAudio: class LoadAudio:
SUPPORTED_FORMATS = ('.wav', '.mp3', '.ogg', '.flac', '.aiff', '.aif')
@classmethod @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
input_dir = folder_paths.get_input_directory() input_dir = folder_paths.get_input_directory()
files = [ files = folder_paths.filter_files_content_types(os.listdir(input_dir), ["audio", "video"])
f for f in os.listdir(input_dir)
if (os.path.isfile(os.path.join(input_dir, f))
and f.endswith(LoadAudio.SUPPORTED_FORMATS)
)
]
return {"required": {"audio": (sorted(files), {"audio_upload": True})}} return {"required": {"audio": (sorted(files), {"audio_upload": True})}}
CATEGORY = "audio" CATEGORY = "audio"

View File

@ -1,5 +1,6 @@
import os import os
import time import time
import mimetypes
import logging import logging
from typing import Set, List, Dict, Tuple from typing import Set, List, Dict, Tuple
@ -43,6 +44,7 @@ input_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "inp
user_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "user") user_directory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "user")
filename_list_cache = {} filename_list_cache = {}
extension_mimetypes_cache = {}
if not os.path.exists(input_directory): if not os.path.exists(input_directory):
try: try:
@ -85,6 +87,28 @@ def get_directory_by_type(type_name):
return get_input_directory() return get_input_directory()
return None return None
def filter_files_content_types(files: List[str], content_types: List[str]) -> List[str]:
"""
Example:
files = os.listdir(folder_paths.get_input_directory())
filter_files_content_types(files, ["image", "audio", "video"])
"""
global extension_mimetypes_cache
result = []
for file in files:
extension = file.split('.')[-1]
if extension not in extension_mimetypes_cache:
mime_type, _ = mimetypes.guess_type(file, strict=False)
if not mime_type:
continue
content_type = mime_type.split('/')[0]
extension_mimetypes_cache[extension] = content_type
else:
content_type = extension_mimetypes_cache[extension]
if content_type in content_types:
result.append(file)
return result
# determine base_dir rely on annotation if name is 'filename.ext [annotation]' format # determine base_dir rely on annotation if name is 'filename.ext [annotation]' format
# otherwise use default_path as base_dir # otherwise use default_path as base_dir