diff --git a/core/type_utils.py b/core/type_utils.py index 4bd1022..7bebdf8 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -358,6 +358,22 @@ CIVARCHIVE_API_TYPE_MAP = { from typing import Any, Callable, Dict, List, Optional, Tuple +def normalize_category_token(category: Any) -> str: + """Normalize a raw category into the shared alias lookup token format.""" + token = ( + str(category or "") + .strip() + .lower() + .replace("\\", "_") + .replace("/", "_") + .replace("-", "_") + .replace(" ", "_") + ) + while "__" in token: + token = token.replace("__", "_") + return token.strip("_") + + def resolve_model_category( category: Any, target_format: str = "folder", @@ -370,18 +386,7 @@ def resolve_model_category( - 'civarchive': CivArchive API type string (e.g. 'Checkpoint', 'LoCon') - 'urn': URN category string """ - token = ( - str(category or "") - .strip() - .lower() - .replace("\\", "_") - .replace("/", "_") - .replace("-", "_") - .replace(" ", "_") - ) - while "__" in token: - token = token.replace("__", "_") - token = token.strip("_") + token = normalize_category_token(category) if target_format == "civitai": canonical_folder = CATEGORY_MAP.get(token, token) diff --git a/tests/test_category_mapping.py b/tests/test_category_mapping.py index b57cd1e..846f865 100644 --- a/tests/test_category_mapping.py +++ b/tests/test_category_mapping.py @@ -2,7 +2,12 @@ import json import unittest from pathlib import Path -from core.type_utils import CATEGORY_MAP, normalize_download_category, resolve_model_category +from core.type_utils import ( + CATEGORY_MAP, + normalize_category_token, + normalize_download_category, + resolve_model_category, +) class CategoryMappingTests(unittest.TestCase): @@ -38,12 +43,32 @@ class CategoryMappingTests(unittest.TestCase): self.assertEqual("checkpoints", normalize_download_category(None)) self.assertEqual("new_category", normalize_download_category("new category")) + def test_category_token_contract_normalizes_common_separators(self): + expected = { + " TEXTUAL\\INVERSION ": ("textual_inversion", "embeddings"), + "unet / gguf": ("unet_gguf", "diffusion_models"), + "select---safetensors": ("select_safetensors", "diffusion_models"), + "clip__vision": ("clip_vision", "clip_vision"), + } + + for raw_category, (token, canonical_category) in expected.items(): + self.assertEqual( + token, + normalize_category_token(raw_category), + ) + self.assertEqual( + canonical_category, + normalize_download_category(raw_category), + ) + def test_frontend_category_alias_artifact_matches_backend_map(self): project_root = Path(__file__).resolve().parents[1] artifact = ( project_root / "web" / "resolver" / "utils" / "category_aliases.generated.js" ).read_text(encoding="utf-8") - payload = artifact.split("Object.freeze(", 1)[1].rsplit(");", 1)[0].strip() + payload = artifact.split( + "export const CATEGORY_ALIASES = Object.freeze(\n", 1 + )[1].split("\n);\n", 1)[0] self.assertEqual(CATEGORY_MAP, json.loads(payload)) diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index d26d821..115f466 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -20,7 +20,10 @@ import { } from '../web/resolver/node_context_menu.js'; import { startSplitterDrag } from '../web/resolver/utils/splitter_drag.js'; import { normalizeDownloadCategoryValue } from '../web/resolver/utils/category_utils.js'; -import { CATEGORY_ALIASES } from '../web/resolver/utils/category_aliases.generated.js'; +import { + CATEGORY_ALIASES, + normalizeCategoryToken, +} from '../web/resolver/utils/category_aliases.generated.js'; import { getSha256Field, normalizeSha256 } from '../web/resolver/utils/hash_utils.js'; import { getSourceDisplayLabel, normalizeSourceKey } from '../web/resolver/utils/source_labels.js'; import { baseModelAliasMethods } from '../web/resolver/search/base_model_alias_methods.js'; @@ -39,6 +42,7 @@ import { void normalizeDownloadCategoryValue; void CATEGORY_ALIASES; +void normalizeCategoryToken; void normalizeSha256; const projectRoot = path.resolve(import.meta.dirname, '..'); @@ -5993,6 +5997,25 @@ test('download category normalization uses the generated backend alias fallback' assert.equal(normalizeDownloadCategoryValue('control-net'), 'controlnet'); }); +test('download category normalization follows the shared separator contract', () => { + const expected = [ + [' TEXTUAL\\INVERSION ', 'embeddings'], + ['unet / gguf', 'diffusion_models'], + ['select---safetensors', 'diffusion_models'], + ['clip__vision', 'clip_vision'], + ]; + + for (const [rawCategory, canonicalCategory] of expected) { + assert.equal(normalizeCategoryToken(rawCategory), rawCategory + .trim() + .toLowerCase() + .replace(/[/\\\s-]+/g, '_') + .replace(/_+/g, '_') + .replace(/^_|_$/g, '')); + assert.equal(normalizeDownloadCategoryValue(rawCategory), canonicalCategory); + } +}); + test('missing model categories preserve node priority and fallback order', () => { const getMissingSupportedDownloadCategories = eval( `(${extractMethod(downloadTargetMethodsSource, 'getMissingSupportedDownloadCategories')})` diff --git a/tools/generate_category_aliases.py b/tools/generate_category_aliases.py index 3496e8a..95b52e4 100644 --- a/tools/generate_category_aliases.py +++ b/tools/generate_category_aliases.py @@ -16,6 +16,14 @@ def main() -> None: payload = json.dumps(CATEGORY_MAP, ensure_ascii=False, indent=4, sort_keys=True) content = ( "// Generated by tools/generate_category_aliases.py. Do not edit manually.\n" + "export function normalizeCategoryToken(category) {\n" + " return String(category || '')\n" + " .trim()\n" + " .toLowerCase()\n" + " .replace(/[/\\\\\\s-]+/g, '_')\n" + " .replace(/_+/g, '_')\n" + " .replace(/^_|_$/g, '');\n" + "}\n\n" "export const CATEGORY_ALIASES = Object.freeze(\n" f"{payload}\n" ");\n" diff --git a/web/resolver/utils/category_aliases.generated.js b/web/resolver/utils/category_aliases.generated.js index 2f46b0c..4ad1591 100644 --- a/web/resolver/utils/category_aliases.generated.js +++ b/web/resolver/utils/category_aliases.generated.js @@ -1,4 +1,13 @@ // Generated by tools/generate_category_aliases.py. Do not edit manually. +export function normalizeCategoryToken(category) { + return String(category || '') + .trim() + .toLowerCase() + .replace(/[/\\\s-]+/g, '_') + .replace(/_+/g, '_') + .replace(/^_|_$/g, ''); +} + export const CATEGORY_ALIASES = Object.freeze( { "audio_encoder": "audio_encoders", diff --git a/web/resolver/utils/category_utils.js b/web/resolver/utils/category_utils.js index 5cbee4d..63c0b97 100644 --- a/web/resolver/utils/category_utils.js +++ b/web/resolver/utils/category_utils.js @@ -1,13 +1,4 @@ -import { CATEGORY_ALIASES } from './category_aliases.generated.js'; - -function normalizeCategoryToken(category) { - return String(category || '') - .trim() - .toLowerCase() - .replace(/[/\\\s-]+/g, '_') - .replace(/_+/g, '_') - .replace(/^_|_$/g, ''); -} +import { CATEGORY_ALIASES, normalizeCategoryToken } from './category_aliases.generated.js'; export function normalizeDownloadCategoryValue( category = '',