From c6cfbdb6060145bef53f4652bb8edc4e64dc548f Mon Sep 17 00:00:00 2001 From: rookiestar28 Date: Fri, 12 Jun 2026 16:35:44 +0800 Subject: [PATCH] feat(models): align managed folder type parity --- docs/release/api_contract.md | 5 +- services/model_manager.py | 36 +++++++++- services/model_manager_transfer.py | 7 ++ services/preflight.py | 9 +++ tests/e2e/specs/model_manager.spec.js | 16 +++++ tests/test_model_manager_service.py | 83 ++++++++++++++++++++++++ tests/test_preflight.py | 36 ++++++++++ web/tabs/model_manager_tab.js | 5 ++ web/tests/unit/model_manager_tab.test.js | 42 ++++++++++++ 9 files changed, 235 insertions(+), 4 deletions(-) diff --git a/docs/release/api_contract.md b/docs/release/api_contract.md index b0cfac1..bbf9b61 100644 --- a/docs/release/api_contract.md +++ b/docs/release/api_contract.md @@ -140,8 +140,9 @@ Connector diagnostics contract notes: Model-manager contract notes: - `/models/downloads` supports `since_seq` cursor polling and may return deterministic delta metadata (`requested_since_seq`, `effective_since_seq`, `next_since_seq`, truncation/reset hints) alongside the task list -- `model_type` values SHOULD use current ComfyUI folder keys where applicable, including `text_encoders`, `diffusion_models`, `clip_vision`, `style_models`, `upscale_models`, `vae_approx`, `audio_encoders`, `background_removal`, `frame_interpolation`, `geometry_estimation`, `optical_flow`, and `detection` -- legacy aliases such as `ckpt`, `checkpoints`, `loras`, `controlnets`, `clip`, `text_encoder`, `unet`, `diffusion_model`, `upscale_model`, and `audio_encoder` are normalized before filtering or import destination resolution +- `model_type` values SHOULD use current ComfyUI folder keys where applicable, including `text_encoders`, `diffusion_models`, `clip_vision`, `style_models`, `upscale_models`, `vae_approx`, `gligen`, `latent_upscale_models`, `hypernetworks`, `photomaker`, `model_patches`, `audio_encoders`, `background_removal`, `frame_interpolation`, `geometry_estimation`, `optical_flow`, and `detection` +- legacy aliases such as `ckpt`, `checkpoints`, `loras`, `controlnets`, `clip`, `text_encoder`, `unet`, `diffusion_model`, `upscale_model`, `latent_upscale_model`, `hypernetwork`, `model_patch`, and `audio_encoder` are normalized before filtering or import destination resolution +- current ComfyUI folder keys that are not managed model-file destinations fail closed for download creation: `configs` (configuration YAML), `diffusers` (folder-valued trees), `classifiers` (extensionless classifier artifacts), and `custom_nodes` (executable plugin code) - download creation requires structured provenance metadata (`publisher`, `license`, `source_url`) and a 64-char `expected_sha256` - import keeps fail-closed destination/filename validation and re-checks the staged file hash before activation diff --git a/services/model_manager.py b/services/model_manager.py index e1b759d..9b57c46 100644 --- a/services/model_manager.py +++ b/services/model_manager.py @@ -97,6 +97,11 @@ MODEL_TYPE_TO_SUBDIR = { "style_models": "style_models", "upscale_models": "upscale_models", "vae_approx": "vae_approx", + "gligen": "gligen", + "latent_upscale_models": "latent_upscale_models", + "hypernetworks": "hypernetworks", + "photomaker": "photomaker", + "model_patches": "model_patches", "audio_encoders": "audio_encoders", "background_removal": "background_removal", "frame_interpolation": "frame_interpolation", @@ -104,17 +109,30 @@ MODEL_TYPE_TO_SUBDIR = { "optical_flow": "optical_flow", "detection": "detection", } +MODEL_TYPE_EXCLUSION_REASONS = { + "configs": "configuration YAML is not a managed model-weight destination", + "diffusers": "diffusers is folder-valued and needs a directory-tree install design", + "classifiers": "classifiers are extensionless and need a dedicated content policy", + "custom_nodes": "custom_nodes are executable plugin code, not managed model files", +} MODEL_TYPE_ALIASES = { + "config": "configs", "ckpt": "checkpoint", "checkpoints": "checkpoint", "loras": "lora", "controlnets": "controlnet", "control_net": "controlnet", + "diffuser": "diffusers", "clip": "text_encoders", "text_encoder": "text_encoders", "unet": "diffusion_models", "diffusion_model": "diffusion_models", "upscale_model": "upscale_models", + "latent_upscale_model": "latent_upscale_models", + "hypernetwork": "hypernetworks", + "photo_maker": "photomaker", + "model_patch": "model_patches", + "classifier": "classifiers", "audio_encoder": "audio_encoders", "geometry": "geometry_estimation", "detector": "detection", @@ -249,11 +267,21 @@ def _truthy(value: str) -> bool: return str(value or "").strip().lower() in {"1", "true", "yes", "on"} -def _norm_model_type(model_type: str) -> str: +def _resolve_model_type_token(model_type: str) -> str: text = str(model_type or "").strip().lower() + if not text: + return "" + return MODEL_TYPE_ALIASES.get(text, text) + + +def _model_type_exclusion_reason(model_type: str) -> str: + return MODEL_TYPE_EXCLUSION_REASONS.get(_resolve_model_type_token(model_type), "") + + +def _norm_model_type(model_type: str) -> str: + text = _resolve_model_type_token(model_type) if not text: return DEFAULT_MODEL_TYPE - text = MODEL_TYPE_ALIASES.get(text, text) return text if text in MODEL_TYPE_TO_SUBDIR else "other" @@ -507,6 +535,10 @@ class ModelManager: def _norm_model_type(value: str) -> str: return _norm_model_type(value) + @staticmethod + def _model_type_exclusion_reason(value: str) -> str: + return _model_type_exclusion_reason(value) + @staticmethod def _norm_source(value: str) -> str: return _norm_source(value) diff --git a/services/model_manager_transfer.py b/services/model_manager_transfer.py index 587f124..ec95086 100644 --- a/services/model_manager_transfer.py +++ b/services/model_manager_transfer.py @@ -105,6 +105,13 @@ def create_download_task( raise manager._error( "validation_error", "expected_sha256 must be a 64-char hex string" ) + exclusion_reason = manager._model_type_exclusion_reason(model_type) + if exclusion_reason: + raise manager._error( + "unsupported_model_type", + f"model_type '{str(model_type or '').strip()}' is not supported " + f"for managed install/import: {exclusion_reason}", + ) manager._validate_url_policy(download_url) provenance = manager._validate_provenance(provenance) mtype = manager._norm_model_type(model_type) diff --git a/services/preflight.py b/services/preflight.py index b17f050..0ace3ed 100644 --- a/services/preflight.py +++ b/services/preflight.py @@ -48,6 +48,7 @@ _LEGACY_INVENTORY_CACHE_KEY = "inventory" _INVENTORY_LOCK = threading.RLock() _INVENTORY_SCAN_THREAD: threading.Thread | None = None _INVENTORY_ERROR_RETRY_SEC = 5 +_INVENTORY_EXCLUDED_MODEL_TYPES = {"custom_nodes"} # Heuristic mapping: input_key -> folder_paths type _INPUT_KEY_MAP = { @@ -83,6 +84,7 @@ def _get_node_class_mappings() -> Dict[str, Any]: def _resolve_inventory_model_types() -> List[str]: model_types = [ "checkpoints", + "configs", "loras", "vae", "embeddings", @@ -94,7 +96,12 @@ def _resolve_inventory_model_types() -> List[str]: "style_models", "diffusers", "vae_approx", + "gligen", + "latent_upscale_models", + "hypernetworks", "photomaker", + "classifiers", + "model_patches", "audio_encoders", "background_removal", "frame_interpolation", @@ -106,6 +113,8 @@ def _resolve_inventory_model_types() -> List[str]: ] if hasattr(folder_paths, "folder_names_and_paths"): for key in folder_paths.folder_names_and_paths.keys(): + if key in _INVENTORY_EXCLUDED_MODEL_TYPES: + continue if key not in model_types: model_types.append(key) return model_types diff --git a/tests/e2e/specs/model_manager.spec.js b/tests/e2e/specs/model_manager.spec.js index e15c7b6..b00cfae 100644 --- a/tests/e2e/specs/model_manager.spec.js +++ b/tests/e2e/specs/model_manager.spec.js @@ -145,6 +145,22 @@ test.describe('Model Manager Tab', () => { await waitForOpenClawReady(page); await clickTab(page, 'Model Manager'); + const modelTypeOptions = await page.locator('#mm-type option').evaluateAll((options) => options.map((option) => option.value)); + expect(modelTypeOptions).toEqual(expect.arrayContaining([ + 'gligen', + 'latent_upscale_models', + 'hypernetworks', + 'photomaker', + 'model_patches', + 'geometry_estimation', + 'optical_flow', + 'detection', + ])); + expect(modelTypeOptions).not.toContain('configs'); + expect(modelTypeOptions).not.toContain('diffusers'); + expect(modelTypeOptions).not.toContain('classifiers'); + expect(modelTypeOptions).not.toContain('custom_nodes'); + await expect(page.locator('#mm-search-results')).toContainText('Flux Test Model'); const queueButton = page.locator('#mm-search-results').getByRole('button', { name: 'Queue Download' }).first(); diff --git a/tests/test_model_manager_service.py b/tests/test_model_manager_service.py index 2501fa4..d3fea2f 100644 --- a/tests/test_model_manager_service.py +++ b/tests/test_model_manager_service.py @@ -10,10 +10,13 @@ from pathlib import Path, PurePosixPath from unittest.mock import patch from services.model_manager import ( + MODEL_TYPE_EXCLUSION_REASONS, + MODEL_TYPE_TO_SUBDIR, DownloadCancelled, DownloadTask, ModelManager, ModelManagerError, + _model_type_exclusion_reason, _norm_model_type, ) from services.model_manager_transfer import ( @@ -138,15 +141,95 @@ class TestModelManagerService(unittest.TestCase): self.assertEqual(_norm_model_type("audio_encoders"), "audio_encoders") self.assertEqual(_norm_model_type("background_removal"), "background_removal") self.assertEqual(_norm_model_type("frame_interpolation"), "frame_interpolation") + self.assertEqual(_norm_model_type("gligen"), "gligen") + self.assertEqual( + _norm_model_type("latent_upscale_models"), "latent_upscale_models" + ) + self.assertEqual(_norm_model_type("hypernetworks"), "hypernetworks") + self.assertEqual(_norm_model_type("photomaker"), "photomaker") + self.assertEqual(_norm_model_type("model_patches"), "model_patches") self.assertEqual(_norm_model_type("geometry_estimation"), "geometry_estimation") self.assertEqual(_norm_model_type("optical_flow"), "optical_flow") self.assertEqual(_norm_model_type("detection"), "detection") self.assertEqual(_norm_model_type("unet"), "diffusion_models") self.assertEqual(_norm_model_type("clip"), "text_encoders") + self.assertEqual( + _norm_model_type("latent_upscale_model"), "latent_upscale_models" + ) + self.assertEqual(_norm_model_type("hypernetwork"), "hypernetworks") + self.assertEqual(_norm_model_type("model_patch"), "model_patches") self.assertEqual(_norm_model_type("geometry"), "geometry_estimation") self.assertEqual(_norm_model_type("detector"), "detection") self.assertEqual(_norm_model_type("diffusers"), "other") + def test_current_comfyui_model_type_support_and_exclusions_are_explicit(self): + supported = { + "checkpoint": "checkpoints", + "lora": "loras", + "vae": "vae", + "controlnet": "controlnet", + "embedding": "embeddings", + "text_encoders": "text_encoders", + "diffusion_models": "diffusion_models", + "clip_vision": "clip_vision", + "style_models": "style_models", + "upscale_models": "upscale_models", + "vae_approx": "vae_approx", + "gligen": "gligen", + "latent_upscale_models": "latent_upscale_models", + "hypernetworks": "hypernetworks", + "photomaker": "photomaker", + "model_patches": "model_patches", + "audio_encoders": "audio_encoders", + "background_removal": "background_removal", + "frame_interpolation": "frame_interpolation", + "geometry_estimation": "geometry_estimation", + "optical_flow": "optical_flow", + "detection": "detection", + } + + for model_type, subdir in supported.items(): + with self.subTest(model_type=model_type): + self.assertEqual(MODEL_TYPE_TO_SUBDIR[model_type], subdir) + self.assertEqual(_norm_model_type(model_type), model_type) + + for excluded in ("configs", "diffusers", "classifiers", "custom_nodes"): + with self.subTest(excluded=excluded): + self.assertIn(excluded, MODEL_TYPE_EXCLUSION_REASONS) + self.assertEqual(_norm_model_type(excluded), "other") + self.assertTrue(_model_type_exclusion_reason(excluded)) + + @patch("services.model_manager.validate_outbound_url") + def test_create_download_task_rejects_known_excluded_current_folder_keys( + self, mock_validate + ): + payload = { + "model_id": "excluded-model", + "name": "Excluded Model", + "source": "catalog", + "source_label": "Catalog", + "download_url": "https://example.com/excluded.safetensors", + "expected_sha256": "a" * 64, + "provenance": { + "publisher": "OpenClaw", + "license": "OpenRAIL", + "source_url": "https://example.com/excluded", + }, + } + + for model_type in ("configs", "diffusers", "classifiers", "custom_nodes"): + with self.subTest(model_type=model_type): + with self.assertRaises(ModelManagerError) as ctx: + self.manager.create_download_task( + model_type=model_type, + **payload, + ) + self.assertEqual(ctx.exception.code, "unsupported_model_type") + self.assertIn(model_type, ctx.exception.detail) + + self.assertFalse(mock_validate.called) + self.assertEqual(self.manager._tasks, {}) + @patch( "services.model_manager.validate_outbound_url", return_value=("https", "example.com", 443, ["1.1.1.1"]), diff --git a/tests/test_preflight.py b/tests/test_preflight.py index 21c0a16..bea01d2 100644 --- a/tests/test_preflight.py +++ b/tests/test_preflight.py @@ -161,6 +161,42 @@ class TestPreflightBackend(AioHTTPTestCase): "detection", ) + def test_inventory_model_types_track_current_comfyui_keys_and_exclude_custom_nodes( + self, + ): + with patch.object( + services.preflight, "folder_paths", MagicMock(), create=True + ) as mock_folder_paths: + mock_folder_paths.folder_names_and_paths = { + "configs": [], + "gligen": [], + "latent_upscale_models": [], + "hypernetworks": [], + "photomaker": [], + "classifiers": [], + "model_patches": [], + "custom_nodes": [], + } + + model_types = services.preflight._resolve_inventory_model_types() + + for model_type in ( + "configs", + "diffusers", + "gligen", + "latent_upscale_models", + "hypernetworks", + "photomaker", + "classifiers", + "model_patches", + "geometry_estimation", + "optical_flow", + "detection", + ): + with self.subTest(model_type=model_type): + self.assertIn(model_type, model_types) + self.assertNotIn("custom_nodes", model_types) + @patch("api.preflight_handler.check_rate_limit") @patch("api.preflight_handler.require_admin_token") @unittest_run_loop diff --git a/web/tabs/model_manager_tab.js b/web/tabs/model_manager_tab.js index 9c18b1d..7e1ae24 100644 --- a/web/tabs/model_manager_tab.js +++ b/web/tabs/model_manager_tab.js @@ -188,6 +188,11 @@ export const ModelManagerTab = { + + + + + diff --git a/web/tests/unit/model_manager_tab.test.js b/web/tests/unit/model_manager_tab.test.js index ad2cc0b..f16415b 100644 --- a/web/tests/unit/model_manager_tab.test.js +++ b/web/tests/unit/model_manager_tab.test.js @@ -88,4 +88,46 @@ describe("model_manager_tab", () => { expect.objectContaining({ task_id: "task-1", state: "completed" }), ]); }); + + it("lists only managed-install supported current ComfyUI folder keys", async () => { + apiMock.searchModels.mockResolvedValue({ + ok: true, + data: { items: [] }, + }); + apiMock.listModelDownloadTasks.mockResolvedValue({ + ok: true, + data: { tasks: [] }, + }); + apiMock.listModelInstallations.mockResolvedValue({ + ok: true, + data: { installations: [] }, + }); + + const container = document.createElement("div"); + ModelManagerTab.render(container); + await vi.waitFor(() => { + expect(apiMock.searchModels).toHaveBeenCalled(); + }); + + const options = Array.from(container.querySelectorAll("#mm-type option")).map( + (option) => option.value + ); + + expect(options).toEqual(expect.arrayContaining([ + "text_encoders", + "diffusion_models", + "gligen", + "latent_upscale_models", + "hypernetworks", + "photomaker", + "model_patches", + "geometry_estimation", + "optical_flow", + "detection", + ])); + expect(options).not.toContain("configs"); + expect(options).not.toContain("diffusers"); + expect(options).not.toContain("classifiers"); + expect(options).not.toContain("custom_nodes"); + }); });