diff --git a/core/settings.py b/core/settings.py index 2319b9a..5d4df26 100644 --- a/core/settings.py +++ b/core/settings.py @@ -73,7 +73,13 @@ TEMPLATE_KEY_ALIASES = { "loras": ("loras", "lora"), "checkpoints": ("checkpoints", "checkpoint"), "embeddings": ("embeddings", "embedding"), - "diffusion_models": ("diffusion_models", "diffusion_model", "unet", "unet_gguf"), + "diffusion_models": ( + "diffusion_models", + "diffusion_model", + "unet", + "unet_gguf", + "model_gguf", + ), "text_encoders": ("text_encoders", "text_encoder", "clip", "clips", "clip_gguf"), "controlnet": ("controlnet", "control_net"), "vae": ("vae",), diff --git a/core/type_utils.py b/core/type_utils.py index 773641f..1d0dcca 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -253,6 +253,7 @@ CATEGORY_MAP = { "diffusers": "diffusers", "unet": "diffusion_models", "unet_gguf": "diffusion_models", + "model_gguf": "diffusion_models", "select_safetensors": "diffusion_models", "diffusion_model": "diffusion_models", "diffusion_models": "diffusion_models", @@ -1130,6 +1131,7 @@ def get_category_folder_keys(category: str) -> list[str]: if folder_key == "diffusion_models": keys.append("unet") keys.append("unet_gguf") + keys.append("model_gguf") elif folder_key == "text_encoders": keys.append("clip") return keys diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 0de7da0..43a4d1d 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -22,7 +22,12 @@ except ImportError: log.warning("Model Resolver: folder_paths not available yet - will retry later") -from .type_utils import MODEL_EXTENSIONS, URN_TYPE_MAP, unique_ordered_strings +from .type_utils import ( + MODEL_EXTENSIONS, + URN_TYPE_MAP, + normalize_download_category, + unique_ordered_strings, +) URN_REGEX = re.compile(r"^urn:air:([^:]+):([^:]+):([^:]+):(\d+)@(\d+)$") @@ -1024,7 +1029,7 @@ def get_model_widget_category_hints( hints_list = [] for h in dynamic_category_hints: if h: - hints_list.append(h) + hints_list.append(normalize_download_category(h)) if dynamic_category_hints: return _ordered_unique_categories(hints_list) diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index b4c4ea9..270d4b0 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -34,6 +34,10 @@ const modelInfoMethodsSource = fs.readFileSync( path.join(projectRoot, 'web/resolver/views/model_info_methods.js'), 'utf8' ); +const missingBrowserMethodsSource = fs.readFileSync( + path.join(projectRoot, 'web/resolver/views/missing_browser_methods.js'), + 'utf8' +); const renderFormatMethodsSource = fs.readFileSync( path.join(projectRoot, 'web/resolver/utils/render_format_methods.js'), 'utf8' @@ -1827,7 +1831,7 @@ test('local match bad folder badge is omitted for supported folder', () => { assert.equal(html, ''); }); -test('local match bad folder badge treats unet gguf as diffusion models', () => { +test('local match bad folder badge treats gguf folder keys as diffusion models', () => { const getLocalMatchCategory = eval(`(${extractMethod(searchPanelMethodsSource, 'getLocalMatchCategory')})`); const getLocalMatchBadFolderWarning = eval(`(${extractMethod(searchPanelMethodsSource, 'getLocalMatchBadFolderWarning')})`); const renderLocalMatchBadFolderBadge = eval(`(${extractMethod(searchPanelMethodsSource, 'renderLocalMatchBadFolderBadge')})`); @@ -1837,7 +1841,7 @@ test('local match bad folder badge treats unet gguf as diffusion models', () => normalizeDownloadCategory(value = '') { return { unet_gguf: 'diffusion_models', - diffusion_models: 'diffusion_models', + model_gguf: 'diffusion_models', }[String(value || '').trim()] || String(value || '').trim(); }, getMissingSupportedDownloadCategories() { @@ -1851,7 +1855,7 @@ test('local match bad folder badge treats unet gguf as diffusion models', () => const html = renderLocalMatchBadFolderBadge.call( dialog, { category: 'diffusion_models' }, - { confidence: 100, model: { category: 'unet_gguf' } } + { confidence: 100, model: { category: 'model_gguf' } } ); assert.equal(html, ''); @@ -1925,11 +1929,12 @@ test('local match bad type badge is omitted for a file type accepted by node cat assert.equal(html, ''); }); -test('download category normalization maps unet gguf to diffusion models', () => { +test('download category normalization maps gguf folder keys to diffusion models', () => { const normalizeDownloadCategory = eval(`(${extractMethod(downloadTargetMethodsSource, 'normalizeDownloadCategory')})`); assert.equal(normalizeDownloadCategory('unet_gguf'), 'diffusion_models'); assert.equal(normalizeDownloadCategory('UNET GGUF'), 'diffusion_models'); + assert.equal(normalizeDownloadCategory('model_gguf'), 'diffusion_models'); }); test('download category normalization maps select safetensors alias to diffusion models', () => { @@ -1939,6 +1944,27 @@ test('download category normalization maps select safetensors alias to diffusion assert.equal(normalizeDownloadCategory('SELECT SAFETENSORS'), 'diffusion_models'); }); +test('missing model preview shows accepted GGUF format next to its category', () => { + const renderMissingModelFormatBadges = eval(`(${extractMethod(missingBrowserMethodsSource, 'renderMissingModelFormatBadges')})`); + const dialog = { + getMissingAcceptedModelFileTypes() { + return [{ extension: 'gguf', label: 'GGUF', display: 'GGUF (.gguf)' }]; + }, + escapeHtml(value) { + return String(value).replace(/[&<>\"]/g, char => ({ '&': '&', '<': '<', '>': '>', '"': '"' }[char])); + }, + }; + + const html = renderMissingModelFormatBadges.call(dialog, { + category: 'diffusion_models', + node_type: 'UnetLoaderGGUF', + }); + + assert.match(html, /mr-model-format-chip/); + assert.match(html, />GGUF<\/span>/); + assert.match(html, /Accepted file format for this node: GGUF \(\.gguf\)\./); +}); + test('hash match label map numbers distinct matched hashes', () => { const normalizeSearchResultSha256 = eval(`(${extractMethod(searchPanelMethodsSource, 'normalizeSearchResultSha256')})`); const getLocalMatchHash = eval(`(${extractMethod(searchPanelMethodsSource, 'getLocalMatchHash')})`); diff --git a/tests/test_refactoring_targets.py b/tests/test_refactoring_targets.py index f7c0d68..7f965db 100644 --- a/tests/test_refactoring_targets.py +++ b/tests/test_refactoring_targets.py @@ -246,8 +246,15 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): def test_category_folder_keys_mapping(self): from core.type_utils import get_category_folder_keys - self.assertEqual(get_category_folder_keys("diffusion_models"), ["diffusion_models", "unet", "unet_gguf"]) - self.assertEqual(get_category_folder_keys("unet_gguf"), ["diffusion_models", "unet", "unet_gguf"]) + expected_diffusion_keys = [ + "diffusion_models", + "unet", + "unet_gguf", + "model_gguf", + ] + self.assertEqual(get_category_folder_keys("diffusion_models"), expected_diffusion_keys) + self.assertEqual(get_category_folder_keys("model_gguf"), expected_diffusion_keys) + self.assertEqual(get_category_folder_keys("unet_gguf"), expected_diffusion_keys) self.assertEqual(get_category_folder_keys("text_encoders"), ["text_encoders", "clip"]) self.assertEqual(get_category_folder_keys("checkpoints"), ["checkpoints"]) @@ -298,5 +305,3 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): # Require download excludes main.safetensors because it has no download url self.assertEqual(select_primary_model_file(files, require_download=True)["id"], 1) - - diff --git a/tests/test_workflow_analyzer.py b/tests/test_workflow_analyzer.py index 96bcd15..840321c 100644 --- a/tests/test_workflow_analyzer.py +++ b/tests/test_workflow_analyzer.py @@ -1,6 +1,7 @@ import os import tempfile import unittest +from unittest.mock import patch from core.scanner import scan_directory from core.workflow_analyzer import analyze_workflow_models, identify_missing_models @@ -412,6 +413,40 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): self.assertEqual(4, len(refs)) self.assertEqual({"text_encoders"}, {ref["category"] for ref in refs}) + def test_gguf_unet_loader_keeps_diffusion_badge_and_raw_folder_hint(self): + for node_type in ( + "LoaderGGUF", + "LoaderGGUFAdvanced", + "UnetLoaderGGUF", + "UnetLoaderGGUFAdvanced", + ): + with self.subTest(node_type=node_type): + with patch( + "core.workflow_analyzer.get_dynamic_widget_category_hints", + return_value=["model_gguf"], + ): + refs = analyze_workflow_models( + { + "nodes": [ + { + "id": 345, + "type": node_type, + "widgets": [{"name": "gguf_name"}], + "widgets_values": [ + r"LTXV\LTX-2.3-22B-distilled-1.1-Q4_K_M.gguf" + ], + "outputs": [{"type": "MODEL", "links": [1]}], + } + ] + }, + available_models=[], + ) + + self.assertEqual(1, len(refs)) + self.assertEqual("diffusion_models", refs[0]["category"]) + self.assertEqual(["diffusion_models"], refs[0]["category_hints"]) + self.assertEqual(["model_gguf"], refs[0]["folder_key_hints"]) + class ScannerFolderModelTests(unittest.TestCase): def test_diffusers_folder_models_are_scanned_as_folder_entries(self): diff --git a/web/css/resolver-main.css b/web/css/resolver-main.css index 3308afa..7c1ab3d 100644 --- a/web/css/resolver-main.css +++ b/web/css/resolver-main.css @@ -1489,6 +1489,12 @@ font-weight: 700; } +.mr-model-format-chip { + background: rgba(67, 160, 255, 0.12); + border-color: rgba(104, 184, 255, 0.3); + color: #b9ddff; +} + /* ComfyUI node_slot-inspired type chips. */ .mr-type-chip, .mr-info-tag.mr-type-chip, diff --git a/web/resolver/search/download_target_methods.js b/web/resolver/search/download_target_methods.js index 6f7408b..76467f5 100644 --- a/web/resolver/search/download_target_methods.js +++ b/web/resolver/search/download_target_methods.js @@ -109,6 +109,7 @@ export const downloadTargetMethods = { 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', @@ -297,6 +298,10 @@ export const downloadTargetMethods = { CheckpointLoader: 'checkpoints', UNETLoader: 'diffusion_models', UNETLoaderAdvanced: 'diffusion_models', + LoaderGGUF: 'diffusion_models', + LoaderGGUFAdvanced: 'diffusion_models', + UnetLoaderGGUF: 'diffusion_models', + UnetLoaderGGUFAdvanced: 'diffusion_models', VAELoader: 'vae', VAELoaderKJ: 'vae', LoraLoader: 'loras', diff --git a/web/resolver/views/missing_browser_methods.js b/web/resolver/views/missing_browser_methods.js index 8de8238..86740b8 100644 --- a/web/resolver/views/missing_browser_methods.js +++ b/web/resolver/views/missing_browser_methods.js @@ -12,6 +12,17 @@ export const missingBrowserMethods = { return this.getFilenameFromPath(missing.original_path) || missing.name || 'Missing model'; }, + renderMissingModelFormatBadges(missing = {}) { + const acceptedTypes = this.getMissingAcceptedModelFileTypes?.(missing) || []; + return acceptedTypes.map(type => { + const label = String(type?.label || type?.extension || '').trim(); + if (!label) return ''; + const display = String(type?.display || label).trim(); + const tooltip = `Accepted file format for this node: ${display}.`; + return `${this.escapeHtml(label)}`; + }).join(''); + }, + getMissingLocateTarget(missing = {}) { const hasLocateNode = missing.locate_node_id !== undefined && missing.locate_node_id !== null @@ -2079,6 +2090,7 @@ export const missingBrowserMethods = { if (missing.category) { html += `${this.getCategoryDisplayName(missing.category)}`; } + html += this.renderMissingModelFormatBadges(missing); html += ``; if (nodeDisplay.canLocate) { html += this.getLocateIconHtml();