From b696e28b7c906e85495ab61ec4c2e82e3122a1a5 Mon Sep 17 00:00:00 2001
From: christian-byrne <abolkonsky.rem@gmail.com>
Date: Thu, 18 Jul 2024 14:28:30 -0700
Subject: [PATCH] Add content-type filter method to folder_paths

---
 comfy_extras/nodes_audio.py |  9 +--------
 folder_paths.py             | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py
index 64d241e98..19bb37713 100644
--- a/comfy_extras/nodes_audio.py
+++ b/comfy_extras/nodes_audio.py
@@ -182,17 +182,10 @@ class PreviewAudio(SaveAudio):
                 }
 
 class LoadAudio:
-    SUPPORTED_FORMATS = ('.wav', '.mp3', '.ogg', '.flac', '.aiff', '.aif')
-
     @classmethod
     def INPUT_TYPES(s):
         input_dir = folder_paths.get_input_directory()
-        files = [
-            f for f in os.listdir(input_dir)
-            if (os.path.isfile(os.path.join(input_dir, f))
-                and f.endswith(LoadAudio.SUPPORTED_FORMATS)
-            )
-        ]
+        files = folder_paths.filter_files_content_types(os.listdir(input_dir), ["audio", "video"])
         return {"required": {"audio": (sorted(files), {"audio_upload": True})}}
 
     CATEGORY = "audio"
diff --git a/folder_paths.py b/folder_paths.py
index 2baf8ce1c..7658853ea 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -1,5 +1,6 @@
 import os
 import time
+import mimetypes
 import logging
 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")
 
 filename_list_cache = {}
+extension_mimetypes_cache = {}
 
 if not os.path.exists(input_directory):
     try:
@@ -85,6 +87,28 @@ def get_directory_by_type(type_name):
         return get_input_directory()
     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
 # otherwise use default_path as base_dir