diff --git a/tests/test_category_mapping.py b/tests/test_category_mapping.py index d16abc2..b57cd1e 100644 --- a/tests/test_category_mapping.py +++ b/tests/test_category_mapping.py @@ -1,6 +1,8 @@ +import json import unittest +from pathlib import Path -from core.type_utils import normalize_download_category, resolve_model_category +from core.type_utils import CATEGORY_MAP, normalize_download_category, resolve_model_category class CategoryMappingTests(unittest.TestCase): @@ -36,6 +38,15 @@ class CategoryMappingTests(unittest.TestCase): self.assertEqual("checkpoints", normalize_download_category(None)) self.assertEqual("new_category", normalize_download_category("new 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() + + self.assertEqual(CATEGORY_MAP, json.loads(payload)) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index 42a9e76..c5c81d3 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -19,6 +19,7 @@ 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 { 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'; @@ -36,6 +37,7 @@ import { } from '../web/resolver/custom_nodes/registry.js'; void normalizeDownloadCategoryValue; +void CATEGORY_ALIASES; void normalizeSha256; const projectRoot = path.resolve(import.meta.dirname, '..'); @@ -5581,6 +5583,11 @@ test('download category normalization maps gguf folder keys to diffusion models' assert.equal(normalizeDownloadCategory('model_gguf'), 'diffusion_models'); }); +test('download category normalization uses the generated backend alias fallback', () => { + assert.equal(CATEGORY_ALIASES.control_net, 'controlnet'); + assert.equal(normalizeDownloadCategoryValue('control-net'), 'controlnet'); +}); + 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 new file mode 100644 index 0000000..3496e8a --- /dev/null +++ b/tools/generate_category_aliases.py @@ -0,0 +1,27 @@ +"""Generate the frontend category alias artifact from the backend source of truth.""" + +import json +import sys +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(PROJECT_ROOT)) + +from core.type_utils import CATEGORY_MAP + +OUTPUT_PATH = PROJECT_ROOT / "web" / "resolver" / "utils" / "category_aliases.generated.js" + + +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 const CATEGORY_ALIASES = Object.freeze(\n" + f"{payload}\n" + ");\n" + ) + OUTPUT_PATH.write_text(content, encoding="utf-8", newline="\n") + + +if __name__ == "__main__": + main() diff --git a/web/resolver/utils/category_aliases.generated.js b/web/resolver/utils/category_aliases.generated.js new file mode 100644 index 0000000..2f46b0c --- /dev/null +++ b/web/resolver/utils/category_aliases.generated.js @@ -0,0 +1,68 @@ +// Generated by tools/generate_category_aliases.py. Do not edit manually. +export const CATEGORY_ALIASES = Object.freeze( +{ + "audio_encoder": "audio_encoders", + "audio_encoders": "audio_encoders", + "background_removal": "background_removal", + "background_removal_model": "background_removal", + "checkpoint": "checkpoints", + "checkpoints": "checkpoints", + "clip": "text_encoders", + "clip_gguf": "text_encoders", + "clip_vision": "clip_vision", + "clips": "text_encoders", + "control_net": "controlnet", + "controlnet": "controlnet", + "default": "upscale_models", + "detection": "detection", + "diffusers": "diffusers", + "diffusion_model": "diffusion_models", + "diffusion_models": "diffusion_models", + "embedding": "embeddings", + "embeddings": "embeddings", + "frame_interpolation": "frame_interpolation", + "frame_interpolation_model": "frame_interpolation", + "geometry_estimation": "geometry_estimation", + "geometry_estimation_model": "geometry_estimation", + "gligen": "gligen", + "hypernetwork": "hypernetworks", + "hypernetworks": "hypernetworks", + "ip_adapter": "ipadapter", + "ipadapter": "ipadapter", + "latent_upscale_model": "latent_upscale_models", + "latent_upscale_models": "latent_upscale_models", + "lora": "loras", + "loras": "loras", + "model_gguf": "diffusion_models", + "model_patch": "model_patches", + "model_patches": "model_patches", + "optical_flow": "optical_flow", + "optical_flow_model": "optical_flow", + "photomaker": "photomaker", + "sam": "sams", + "sam_model": "sams", + "sam_models": "sams", + "sams": "sams", + "seedvr2": "seedvr2", + "seedvr2_dit": "seedvr2", + "seedvr2_vae": "seedvr2", + "select_safetensors": "diffusion_models", + "style_model": "style_models", + "style_models": "style_models", + "text_encoder": "text_encoders", + "text_encoders": "text_encoders", + "textual_inversion": "embeddings", + "textualinversion": "embeddings", + "ultralytics": "ultralytics", + "ultralytics_bbox": "ultralytics", + "ultralytics_segm": "ultralytics", + "unet": "diffusion_models", + "unet_gguf": "diffusion_models", + "upscale_model": "upscale_models", + "upscale_models": "upscale_models", + "upscaler": "upscale_models", + "vae": "vae", + "vae_approx": "vae_approx", + "yolo": "ultralytics" +} +); diff --git a/web/resolver/utils/category_utils.js b/web/resolver/utils/category_utils.js index cdc0a80..5cbee4d 100644 --- a/web/resolver/utils/category_utils.js +++ b/web/resolver/utils/category_utils.js @@ -1,39 +1,4 @@ -const FALLBACK_DOWNLOAD_CATEGORY_ALIASES = Object.freeze({ - checkpoint: 'checkpoints', - lora: 'loras', - embedding: 'embeddings', - textualinversion: 'embeddings', - textual_inversion: 'embeddings', - upscaler: 'upscale_models', - unet: 'diffusion_models', - unet_gguf: 'diffusion_models', - model_gguf: 'diffusion_models', - select_safetensors: 'diffusion_models', - diffusion_model: 'diffusion_models', - diffusion_models: 'diffusion_models', - clip: 'text_encoders', - clips: 'text_encoders', - clip_gguf: 'text_encoders', - text_encoder: 'text_encoders', - text_encoders: 'text_encoders', - ip_adapter: 'ipadapter', - upscale_model: 'upscale_models', - latent_upscale_model: 'latent_upscale_models', - style_model: 'style_models', - audio_encoder: 'audio_encoders', - model_patch: 'model_patches', - sam: 'sams', - sam_model: 'sams', - sam_models: 'sams', - ultralytics_bbox: 'ultralytics', - ultralytics_segm: 'ultralytics', - yolo: 'ultralytics', - background_removal_model: 'background_removal', - frame_interpolation_model: 'frame_interpolation', - geometry_estimation_model: 'geometry_estimation', - optical_flow_model: 'optical_flow', - default: 'upscale_models', -}); +import { CATEGORY_ALIASES } from './category_aliases.generated.js'; function normalizeCategoryToken(category) { return String(category || '') @@ -46,11 +11,11 @@ function normalizeCategoryToken(category) { export function normalizeDownloadCategoryValue( category = '', - aliases = FALLBACK_DOWNLOAD_CATEGORY_ALIASES, + aliases = CATEGORY_ALIASES, ) { const token = normalizeCategoryToken(category); const categoryAliases = aliases && typeof aliases === 'object' ? aliases - : FALLBACK_DOWNLOAD_CATEGORY_ALIASES; + : CATEGORY_ALIASES; return categoryAliases[token] || token || 'checkpoints'; }