diff --git a/app/model_manager.py b/app/model_manager.py index 8f6e34b33..ba300eff2 100644 --- a/app/model_manager.py +++ b/app/model_manager.py @@ -35,7 +35,17 @@ class ModelFileManager: for folder in model_types: if folder in folder_black_list: continue - output_folders.append({"name": folder, "folders": folder_paths.get_folder_paths(folder)}) + # Effective display filter: the folder's registered extension + # set, or the global supported_pt_extensions for match-all + # folders (empty set), resolved live so runtime registrations + # by custom nodes are reflected. + registered = folder_paths.folder_names_and_paths[folder][1] + effective = set(registered) if registered else set(folder_paths.supported_pt_extensions) + output_folders.append({ + "name": folder, + "folders": folder_paths.get_folder_paths(folder), + "extensions": sorted(effective), + }) return web.json_response(output_folders) # NOTE: This is an experiment to replace `/models/{folder}` diff --git a/openapi.yaml b/openapi.yaml index c37d543c7..02fa2b7d6 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -775,6 +775,14 @@ components: ModelFolder: description: Represents a folder containing models properties: + extensions: + description: 'Effective file-extension display filter for this folder: the registered extension set, or the global supported model extensions for folders registered without one (match-all). Resolved live, so runtime registrations by custom nodes are reflected.' + example: + - .ckpt + - .safetensors + items: + type: string + type: array folders: description: List of paths where models of this type are stored example: diff --git a/tests-unit/app_test/model_manager_test.py b/tests-unit/app_test/model_manager_test.py index ae59206f6..7111e4a5c 100644 --- a/tests-unit/app_test/model_manager_test.py +++ b/tests-unit/app_test/model_manager_test.py @@ -24,6 +24,28 @@ def app(model_manager): app.add_routes(routes) return app +async def test_get_model_folders_includes_effective_extensions(aiohttp_client, app, tmp_path): + """Folders expose their effective display filter: the registered extension + set, or the global supported_pt_extensions for match-all (empty) folders.""" + with patch('folder_paths.folder_names_and_paths', { + 'test_checkpoints': ([str(tmp_path)], {'.safetensors', '.ckpt'}), + 'test_configs': ([str(tmp_path)], ['.yaml']), + 'test_match_all': ([str(tmp_path)], set()), + 'configs': ([str(tmp_path)], ['.yaml']), + }), patch('folder_paths.supported_pt_extensions', {'.safetensors', '.bin'}): + client = await aiohttp_client(app) + response = await client.get('/experiment/models') + + assert response.status == 200 + folders = {f['name']: f for f in await response.json()} + + assert 'configs' not in folders # blocklisted + assert folders['test_checkpoints']['folders'] == [str(tmp_path)] + assert folders['test_checkpoints']['extensions'] == ['.ckpt', '.safetensors'] + assert folders['test_configs']['extensions'] == ['.yaml'] + # Match-all folders substitute the live global set. + assert folders['test_match_all']['extensions'] == ['.bin', '.safetensors'] + async def test_get_model_preview_safetensors(aiohttp_client, app, tmp_path): img = Image.new('RGB', (100, 100), 'white') img_byte_arr = BytesIO()