This commit is contained in:
Andrew Pulley 2026-07-07 09:35:30 +08:00 committed by GitHub
commit 879396b457
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 1 deletions

View File

@ -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["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") output_directory = os.path.join(base_path, "output")
temp_directory = os.path.join(base_path, "temp") temp_directory = os.path.join(base_path, "temp")

View File

@ -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