diff --git a/folder_paths.py b/folder_paths.py index ee048b0f2..5fdbacce7 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -58,7 +58,7 @@ folder_names_and_paths["geometry_estimation"] = ([os.path.join(models_dir, "geom folder_names_and_paths["optical_flow"] = ([os.path.join(models_dir, "optical_flow")], supported_pt_extensions) -folder_names_and_paths["detection"] = ([os.path.join(models_dir, "detection")], supported_pt_extensions) +folder_names_and_paths["detection"] = ([os.path.join(models_dir, "detection")], supported_pt_extensions | {".onnx"}) output_directory = os.path.join(base_path, "output") temp_directory = os.path.join(base_path, "temp") diff --git a/tests-unit/folder_paths_test/detection_onnx_extensions_test.py b/tests-unit/folder_paths_test/detection_onnx_extensions_test.py new file mode 100644 index 000000000..093a353d3 --- /dev/null +++ b/tests-unit/folder_paths_test/detection_onnx_extensions_test.py @@ -0,0 +1,23 @@ +import os +import tempfile + +import folder_paths + + +def test_detection_folder_allows_onnx_extension(): + _, extensions = folder_paths.folder_names_and_paths["detection"] + assert ".onnx" in extensions + + +def test_detection_folder_lists_onnx_files(): + with tempfile.TemporaryDirectory() as detection_dir: + onnx_path = os.path.join(detection_dir, "vitpose-l-wholebody.onnx") + with open(onnx_path, "wb") as onnx_file: + onnx_file.write(b"onnx") + + folder_paths.add_model_folder_path("detection", detection_dir, is_default=True) + folder_paths.filename_list_cache.pop("detection", None) + + filenames = folder_paths.get_filename_list("detection") + assert "vitpose-l-wholebody.onnx" in filenames + assert folder_paths.get_full_path("detection", "vitpose-l-wholebody.onnx") == onnx_path