From 06b2d4c1d029cdab89b860512cf9cee54c94990e Mon Sep 17 00:00:00 2001 From: Simon Pinfold Date: Thu, 2 Jul 2026 12:41:52 +1200 Subject: [PATCH] fix(assets): filter model_type tags by bucket extension sets Buckets sharing a base directory (e.g. diffusion_models and a custom unet_gguf) tagged every file in the directory regardless of whether the bucket could load it, so .safetensors files were tagged model_type:unet_gguf and vice versa. Carry each bucket's registered extension set through get_comfy_models_folders and only emit a model_type tag when the file extension matches, keeping the empty-set match-all convention from folder_paths.filter_files_extensions. Files under a model base matching no bucket now keep only the models tag instead of every directory-matching model_type tag. --- app/assets/scanner.py | 4 +- app/assets/services/path_utils.py | 42 ++++++++--- .../assets_test/services/test_bulk_ingest.py | 4 +- .../assets_test/services/test_path_utils.py | 75 ++++++++++++++++--- 4 files changed, 98 insertions(+), 27 deletions(-) diff --git a/app/assets/scanner.py b/app/assets/scanner.py index 2c1e97840..b13276b49 100644 --- a/app/assets/scanner.py +++ b/app/assets/scanner.py @@ -63,7 +63,7 @@ RootType = Literal["models", "input", "output"] def get_prefixes_for_root(root: RootType) -> list[str]: if root == "models": bases: list[str] = [] - for _bucket, paths in get_comfy_models_folders(): + for _bucket, paths, _exts in get_comfy_models_folders(): bases.extend(paths) return [os.path.abspath(p) for p in bases] if root == "input": @@ -81,7 +81,7 @@ def get_all_known_prefixes() -> list[str]: def collect_models_files() -> list[str]: out: list[str] = [] - for folder_name, bases in get_comfy_models_folders(): + for folder_name, bases, _exts in get_comfy_models_folders(): rel_files = folder_paths.get_filename_list(folder_name) or [] for rel_path in rel_files: if not all(is_visible(part) for part in Path(rel_path).parts): diff --git a/app/assets/services/path_utils.py b/app/assets/services/path_utils.py index 3430f5561..16277d90f 100644 --- a/app/assets/services/path_utils.py +++ b/app/assets/services/path_utils.py @@ -9,20 +9,23 @@ _NON_MODEL_FOLDER_NAMES = frozenset({"configs", "custom_nodes"}) _KNOWN_SUBFOLDER_TAGS = frozenset({"3d", "pasted", "painter", "threed", "webcam"}) -def get_comfy_models_folders() -> list[tuple[str, list[str]]]: - """Build list of (folder_name, base_paths[]) for all model locations. +def get_comfy_models_folders() -> list[tuple[str, list[str], set[str]]]: + """Build list of (folder_name, base_paths[], extensions) for all model locations. Includes every category registered in folder_names_and_paths, regardless of whether its paths are under the main models_dir, but excludes non-model entries like configs and custom_nodes. + + An empty extensions set means the category accepts any extension, + matching folder_paths.filter_files_extensions semantics. """ - targets: list[tuple[str, list[str]]] = [] + targets: list[tuple[str, list[str], set[str]]] = [] for name, values in folder_paths.folder_names_and_paths.items(): if name in _NON_MODEL_FOLDER_NAMES: continue - paths, _exts = values[0], values[1] + paths, exts = values[0], values[1] if paths: - targets.append((name, paths)) + targets.append((name, paths, set(exts))) return targets @@ -44,7 +47,9 @@ def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]: folder_name = model_type_tags[0].split(":", 1)[1] if not folder_name: raise ValueError("models uploads require exactly one model_type: tag") - model_folder_paths = dict(get_comfy_models_folders()) + model_folder_paths = { + name: paths for name, paths, _exts in get_comfy_models_folders() + } try: bases = model_folder_paths[folder_name] except KeyError: @@ -199,7 +204,7 @@ def get_asset_category_and_relative_path( # 4) models (check deepest matching base to avoid ambiguity) best: tuple[int, str, str] | None = None # (base_len, bucket, rel_inside_bucket) - for bucket, bases in get_comfy_models_folders(): + for bucket, bases, _exts in get_comfy_models_folders(): for b in bases: base_abs = os.path.abspath(b) if not _check_is_within(fp_abs, base_abs): @@ -225,6 +230,13 @@ def get_backend_system_tags_from_path(path: str) -> list[str]: The returned tags are only the backend-generated system tags: ``models``, ``model_type:``, ``input``, ``output``, and ``temp``. Model type tags are based on registered folder names, not path components. + + A ``model_type:`` tag is only emitted when the file's + extension is accepted by that folder's registered extension set, so + categories sharing a base directory (e.g. ``diffusion_models`` and a + custom ``unet_gguf``) tag only the files they can actually load. Files + under a model base whose extension matches no category still get the + ``models`` tag. """ fp_abs = os.path.abspath(path) fp_path = Path(fp_abs) @@ -242,17 +254,23 @@ def get_backend_system_tags_from_path(path: str) -> list[str]: if fp_path.is_relative_to(os.path.abspath(base)): _add(role) + ext = os.path.splitext(fp_abs)[1].lower() model_types: list[str] = [] - for folder_name, bases in get_comfy_models_folders(): + under_models_base = False + for folder_name, bases, extensions in get_comfy_models_folders(): for base in bases: if fp_path.is_relative_to(os.path.abspath(base)): - model_types.append(folder_name) + under_models_base = True + # Empty set accepts any extension, matching + # folder_paths.filter_files_extensions semantics. + if not extensions or ext in extensions: + model_types.append(folder_name) break - if model_types: + if under_models_base: _add("models") - for folder_name in model_types: - _add(f"model_type:{folder_name}") + for folder_name in model_types: + _add(f"model_type:{folder_name}") if not tags: raise ValueError( diff --git a/tests-unit/assets_test/services/test_bulk_ingest.py b/tests-unit/assets_test/services/test_bulk_ingest.py index 3754ece32..07cd5510d 100644 --- a/tests-unit/assets_test/services/test_bulk_ingest.py +++ b/tests-unit/assets_test/services/test_bulk_ingest.py @@ -215,8 +215,8 @@ class TestBatchInsertSeedAssets: patch( "app.assets.services.path_utils.get_comfy_models_folders", return_value=[ - ("checkpoints", [str(shared_root)]), - ("diffusion_models", [str(shared_root)]), + ("checkpoints", [str(shared_root)], {".safetensors"}), + ("diffusion_models", [str(shared_root)], {".safetensors"}), ], ), ): diff --git a/tests-unit/assets_test/services/test_path_utils.py b/tests-unit/assets_test/services/test_path_utils.py index af08afa30..3c4db067c 100644 --- a/tests-unit/assets_test/services/test_path_utils.py +++ b/tests-unit/assets_test/services/test_path_utils.py @@ -38,7 +38,7 @@ def fake_dirs(): with patch( "app.assets.services.path_utils.get_comfy_models_folders", - return_value=[("checkpoints", [str(models_dir)])], + return_value=[("checkpoints", [str(models_dir)], {".safetensors"})], ): yield { "input": input_dir, @@ -107,7 +107,7 @@ class TestGetAssetCategoryAndRelativePath: with patch( "app.assets.services.path_utils.get_comfy_models_folders", - return_value=[("LLM", [str(llm_dir)])], + return_value=[("LLM", [str(llm_dir)], {".safetensors"})], ): _name, tags = get_name_and_tags_from_asset_path(str(f)) @@ -136,8 +136,8 @@ class TestGetAssetCategoryAndRelativePath: with patch( "app.assets.services.path_utils.get_comfy_models_folders", return_value=[ - ("checkpoints", [str(shared_root)]), - ("loras", [str(shared_root)]), + ("checkpoints", [str(shared_root)], {".safetensors"}), + ("loras", [str(shared_root)], {".safetensors"}), ], ): _name, tags = get_name_and_tags_from_asset_path(str(f)) @@ -146,6 +146,55 @@ class TestGetAssetCategoryAndRelativePath: assert "model_type:checkpoints" in tags assert "model_type:loras" in tags + def test_shared_root_model_type_tags_respect_bucket_extensions(self, fake_dirs): + """Buckets sharing a base dir only tag files matching their extensions.""" + shared_root = fake_dirs["models"].parent / "unet" + shared_root.mkdir() + safetensors_file = shared_root / "wan.safetensors" + gguf_file = shared_root / "wan.gguf" + safetensors_file.touch() + gguf_file.touch() + + with patch( + "app.assets.services.path_utils.get_comfy_models_folders", + return_value=[ + ("diffusion_models", [str(shared_root)], {".safetensors"}), + ("unet_gguf", [str(shared_root)], {".gguf"}), + ], + ): + _name, safetensors_tags = get_name_and_tags_from_asset_path(str(safetensors_file)) + _name, gguf_tags = get_name_and_tags_from_asset_path(str(gguf_file)) + + assert "model_type:diffusion_models" in safetensors_tags + assert "model_type:unet_gguf" not in safetensors_tags + assert "model_type:unet_gguf" in gguf_tags + assert "model_type:diffusion_models" not in gguf_tags + + def test_empty_extension_set_tags_any_extension(self, fake_dirs): + """Custom buckets registered without extensions accept every file.""" + custom_root = fake_dirs["models"].parent / "custom_bucket" + custom_root.mkdir() + f = custom_root / "weights.bin" + f.touch() + + with patch( + "app.assets.services.path_utils.get_comfy_models_folders", + return_value=[("custom_bucket", [str(custom_root)], set())], + ): + _name, tags = get_name_and_tags_from_asset_path(str(f)) + + assert "models" in tags + assert "model_type:custom_bucket" in tags + + def test_no_extension_match_keeps_models_tag_without_model_type(self, fake_dirs): + f = fake_dirs["models"] / "notes.txt" + f.touch() + + _name, tags = get_name_and_tags_from_asset_path(str(f)) + + assert "models" in tags + assert not any(tag.startswith("model_type:") for tag in tags) + def test_output_backed_registered_folder_gets_model_and_output_tags(self, fake_dirs): output_checkpoints_dir = fake_dirs["output"] / "checkpoints" output_checkpoints_dir.mkdir() @@ -154,7 +203,7 @@ class TestGetAssetCategoryAndRelativePath: with patch( "app.assets.services.path_utils.get_comfy_models_folders", - return_value=[("checkpoints", [str(output_checkpoints_dir)])], + return_value=[("checkpoints", [str(output_checkpoints_dir)], {".safetensors"})], ): _name, tags = get_name_and_tags_from_asset_path(str(f)) @@ -277,7 +326,9 @@ class TestResponseStoragePaths: with patch( "app.assets.services.path_utils.get_comfy_models_folders", - return_value=[(folder_name, [str(default_model_dir), str(output_model_dir)])], + return_value=[ + (folder_name, [str(default_model_dir), str(output_model_dir)], {".safetensors"}) + ], ): assert compute_file_path(str(f)) == f"output/{folder_name}/saved.safetensors" assert compute_display_name(str(f)) == f"{folder_name}/saved.safetensors" @@ -299,7 +350,7 @@ class TestResponseStoragePaths: with patch( "app.assets.services.path_utils.get_comfy_models_folders", - return_value=[(folder_name, [str(output_model_dir)])], + return_value=[(folder_name, [str(output_model_dir)], {".safetensors"})], ): assert ( compute_file_path(str(f)) @@ -323,7 +374,7 @@ class TestResponseStoragePaths: with patch( "app.assets.services.path_utils.get_comfy_models_folders", - return_value=[("checkpoints", [str(external_checkpoints_dir)])], + return_value=[("checkpoints", [str(external_checkpoints_dir)], {".safetensors"})], ): assert compute_file_path(str(f)) is None assert compute_display_name(str(f)) is None @@ -347,7 +398,7 @@ class TestResponseStoragePaths: with patch( "app.assets.services.path_utils.get_comfy_models_folders", - return_value=[("checkpoints", [str(foo_dir), str(bar_dir)])], + return_value=[("checkpoints", [str(foo_dir), str(bar_dir)], {".safetensors"})], ): assert compute_file_path(str(foo_file)) is None assert compute_file_path(str(bar_file)) is None @@ -362,7 +413,7 @@ class TestResponseStoragePaths: with patch( "app.assets.services.path_utils.get_comfy_models_folders", - return_value=[("text_encoders", [str(output_clip_dir)])], + return_value=[("text_encoders", [str(output_clip_dir)], {".safetensors"})], ): assert compute_file_path(str(f)) == "output/clip/clip_l.safetensors" assert compute_display_name(str(f)) == "clip/clip_l.safetensors" @@ -384,7 +435,9 @@ class TestResponseStoragePaths: with patch( "app.assets.services.path_utils.get_comfy_models_folders", - return_value=[("diffusion_models", [str(unet_dir), str(diffusion_models_dir)])], + return_value=[ + ("diffusion_models", [str(unet_dir), str(diffusion_models_dir)], {".safetensors"}) + ], ): assert compute_file_path(str(f)) == "models/unet/wan.safetensors" assert compute_display_name(str(f)) == "unet/wan.safetensors"