diff --git a/__init__.py b/__init__.py index 74ee0bf..7283148 100644 --- a/__init__.py +++ b/__init__.py @@ -824,6 +824,8 @@ class ModelLinkerExtension: "model_id": model_id, "version_id": version_id, "size": primary_file.get("size"), + "base_model": model_info.get("base_model"), + "tags": model_info.get("tags", []), } results["found"] = True elif category: @@ -846,6 +848,8 @@ class ModelLinkerExtension: ), "url": first_result.get("url"), "size": first_result.get("size"), + "base_model": first_result.get("base_model"), + "tags": first_result.get("tags", []), } results["found"] = True else: @@ -1006,6 +1010,56 @@ class ModelLinkerExtension: ) return web.json_response({"error": str(e)}, status=500) + @routes.get("/model_linker/subfolders/{category}") + async def get_subfolders(request): + """Get known subfolders for a category using ComfyUI folder_paths.""" + try: + import folder_paths + + raw_category = (request.match_info.get("category") or "").strip() + category_map = { + "checkpoint": "checkpoints", + "checkpoints": "checkpoints", + "lora": "loras", + "loras": "loras", + "vae": "vae", + "controlnet": "controlnet", + "clip": "clip", + "clip_vision": "clip_vision", + "upscaler": "upscale_models", + "upscale_models": "upscale_models", + "embeddings": "embeddings", + "embedding": "embeddings", + "diffusion_models": "diffusion_models", + "unet": "diffusion_models", + "text_encoders": "text_encoders", + "text_encoder": "text_encoders", + "ipadapter": "ipadapter", + "ip-adapter": "ipadapter", + "default": "upscale_models", + } + category = category_map.get(raw_category.lower(), raw_category.lower()) + + subfolders = set() + filenames = folder_paths.get_filename_list(category) or [] + for rel_path in filenames: + if not isinstance(rel_path, str): + continue + parts = [p for p in rel_path.replace("/", "\\").split("\\") if p] + if len(parts) <= 1: + continue + current = "" + for part in parts[:-1]: + current = f"{current}\\{part}" if current else part + subfolders.add(current) + + return web.json_response(sorted(subfolders)) + except Exception as e: + self.logger.error( + f"Model Linker subfolders error: {e}", exc_info=True + ) + return web.json_response({"error": str(e)}, status=500) + self.routes_setup = True self.logger.info("Model Linker: API routes registered successfully") return True diff --git a/core/downloader.py b/core/downloader.py index 89bae27..b2a07a9 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -96,11 +96,31 @@ def get_download_directory(category: str) -> Optional[str]: folder_key = category_map.get(category.lower(), category.lower()) + def _normalize(path_value: str) -> str: + return os.path.normcase(os.path.abspath(path_value)) + + def _is_within(path_value: str, root_value: str) -> bool: + try: + return os.path.commonpath([_normalize(path_value), _normalize(root_value)]) == _normalize(root_value) + except Exception: + return False + + def _choose_preferred_path(paths: List[str]) -> Optional[str]: + if not paths: + return None + + comfy_root = os.path.dirname(os.path.abspath(getattr(folder_paths, "__file__", ""))) + if comfy_root: + redirected_paths = [path for path in paths if not _is_within(path, comfy_root)] + if redirected_paths: + return redirected_paths[0] + + return paths[0] + try: paths = folder_paths.get_folder_paths(folder_key) if paths: - # Return the first path (usually the main models directory) - return paths[0] + return _choose_preferred_path(paths) # If category not found, try to get any models directory as fallback all_names = folder_paths.get_folder_names() @@ -108,7 +128,7 @@ def get_download_directory(category: str) -> Optional[str]: # Fall back to first available directory fallback_paths = folder_paths.get_folder_paths(all_names[0]) if fallback_paths: - return fallback_paths[0] + return _choose_preferred_path(fallback_paths) except Exception as e: log_debug(f"Could not get folder path for {folder_key}: {e}") diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 803be5a..9679dd7 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -148,6 +148,8 @@ def search_civitai_for_file( "url": f"https://civitai.com/models/{model_id}", "download_url": download_url, "size": file_info.get("sizeKB", 0) * 1024, + "base_model": version.get("baseModel"), + "tags": item.get("tags", []), "match_type": "exact", } _search_cache[cache_key] = result @@ -171,6 +173,8 @@ def search_civitai_for_file( "url": f"https://civitai.com/models/{model_id}", "download_url": download_url, "size": file_info.get("sizeKB", 0) * 1024, + "base_model": version.get("baseModel"), + "tags": item.get("tags", []), "match_type": "partial", } _search_cache[cache_key] = result @@ -261,6 +265,8 @@ def search_civitai( "url": f"https://civitai.com/models/{model_id}", "download_url": get_civitai_download_url(version_id, api_key), "downloads": model.get("stats", {}).get("downloadCount", 0), + "base_model": latest.get("baseModel"), + "tags": model.get("tags", []), } if primary_file: @@ -389,6 +395,8 @@ def resolve_urn( "model_name": data.get("name", "Unknown"), "version_name": target_version.get("name", "Unknown"), "expected_filename": primary_file.get("name", "Unknown"), + "base_model": target_version.get("baseModel"), + "tags": data.get("tags", []), "files": [ {"name": f.get("name"), "size": f.get("sizeKB", 0) * 1024} for f in files diff --git a/web/linker.js b/web/linker.js index 842a1f5..b9c9e6a 100644 --- a/web/linker.js +++ b/web/linker.js @@ -16,6 +16,8 @@ class LinkerManagerDialog extends ComfyDialog { this.currentWorkflow = null; this.missingModels = []; this.allModels = null; // list of all available models for dropdown + this.downloadDirectories = null; + this.downloadSubfolders = new Map(); this.pendingResolutions = []; this.pendingIndex = new Map(); // key -> index in pendingResolutions this.activeDownloads = {}; // Track active downloads @@ -254,6 +256,7 @@ class LinkerManagerDialog extends ComfyDialog { html += `
Direct CivitAI download can require an API key for some models.`; } html += ``; + html += this.renderDownloadTargetControls(missing, downloadSource.directory || downloadSource.category || 'checkpoints'); html += ``; return html; } @@ -1910,6 +1913,71 @@ class LinkerManagerDialog extends ComfyDialog { font-size: 12px; color: var(--ml-text-muted); } + .ml-download-target { + margin-top: 10px; + padding-top: 10px; + border-top: 1px solid rgba(255,255,255,0.06); + display: grid; + gap: 8px; + } + .ml-download-target-row { + display: grid; + gap: 6px; + } + .ml-download-target-label { + font-size: 11px; + font-weight: 700; + color: var(--ml-text-muted); + text-transform: uppercase; + letter-spacing: 0.06em; + } + .ml-download-target-input, + .ml-download-target-select { + width: 100%; + padding: 8px 10px; + background: rgba(255,255,255,0.04); + color: var(--ml-text); + border: 1px solid var(--ml-border); + border-radius: 10px; + font-size: 12px; + outline: none; + } + .ml-download-target-input:focus, + .ml-download-target-select:focus { + border-color: rgba(78,161,255,0.45); + box-shadow: 0 0 0 3px rgba(78,161,255,0.12); + } + .ml-download-target-wrap { + position: relative; + } + .ml-download-target-list { + position: absolute; + top: calc(100% + 4px); + left: 0; + right: 0; + max-height: 220px; + overflow: auto; + display: none; + z-index: 100000; + background: #22272e; + border: 1px solid var(--ml-border-strong); + border-radius: 12px; + box-shadow: var(--ml-shadow); + } + .ml-download-target-option { + padding: 8px 10px; + cursor: pointer; + border-bottom: 1px solid rgba(255,255,255,0.04); + color: var(--ml-text); + font-size: 12px; + } + .ml-download-target-option:last-child { + border-bottom: none; + } + .ml-download-target-option:hover, + .ml-download-target-option.is-highlighted { + background: rgba(78,161,255,0.16); + } `; document.head.appendChild(styles); @@ -2329,6 +2397,230 @@ class LinkerManagerDialog extends ComfyDialog { } } + async ensureDownloadDirectoriesLoaded() { + if (this.downloadDirectories) return; + try { + const resp = await api.fetchApi('/model_linker/directories'); + if (!resp.ok) throw new Error(`HTTP ${resp.status}`); + const directories = await resp.json(); + this.downloadDirectories = directories && typeof directories === 'object' ? directories : {}; + } catch (e) { + console.warn('Model Linker: could not load download directories', e); + this.downloadDirectories = {}; + } + } + + getDownloadCategoryOptions(defaultCategory = 'checkpoints') { + const directories = this.downloadDirectories || {}; + const keys = Object.keys(directories); + const preferred = defaultCategory || 'checkpoints'; + const ordered = [ + preferred, + ...keys.filter(key => key !== preferred) + ].filter((value, index, arr) => value && arr.indexOf(value) === index); + + return ordered.length > 0 ? ordered : [preferred]; + } + + getAvailableSubfolders(category = '') { + return this.downloadSubfolders.get((category || '').toLowerCase()) || []; + } + + normalizeFolderToken(value = '') { + return String(value || '') + .toLowerCase() + .replace(/[\/\\]+/g, ' ') + .replace(/[^a-z0-9]+/g, ''); + } + + getSuggestedCivitaiSubfolder(missing, category, folders = []) { + if ((category || '').toLowerCase() !== 'loras' || !folders.length) { + return ''; + } + + const civitaiData = { + ...(missing?.civitai_info || {}), + ...(missing?.download_source || {}) + }; + const baseModel = civitaiData.base_model || ''; + const tags = Array.isArray(civitaiData.tags) ? civitaiData.tags.filter(Boolean) : []; + if (!baseModel) return ''; + + const priorityTags = [ + 'concept', + 'style', + 'character', + 'clothing', + 'pose', + 'object', + 'vehicle', + 'artist', + 'celebrity' + ]; + const normalizedBase = this.normalizeFolderToken(baseModel); + if (!normalizedBase) return ''; + + const folderEntries = folders.map(folder => { + const segments = String(folder || '').split(/[\/\\]/).filter(Boolean); + return { + value: folder, + segments, + normalizedSegments: segments.map(segment => this.normalizeFolderToken(segment)) + }; + }); + + const baseMatches = folderEntries.filter(entry => entry.normalizedSegments[0] === normalizedBase); + if (!baseMatches.length) return ''; + + const exactBase = baseMatches.find(entry => entry.segments.length === 1); + const orderedTags = [ + ...priorityTags.filter(tag => tags.some(value => this.normalizeFolderToken(value) === this.normalizeFolderToken(tag))), + ...tags + ].filter((value, index, arr) => value && arr.findIndex(other => this.normalizeFolderToken(other) === this.normalizeFolderToken(value)) === index); + + for (const tag of orderedTags) { + const normalizedTag = this.normalizeFolderToken(tag); + if (!normalizedTag) continue; + const match = baseMatches.find(entry => entry.normalizedSegments[1] === normalizedTag); + if (match) { + return match.value; + } + } + + return exactBase?.value || ''; + } + + async applySuggestedCivitaiSubfolder(missing, categoryEl, subfolderEl) { + if (!categoryEl || !subfolderEl || subfolderEl.value.trim()) return; + + await this.ensureDownloadSubfoldersLoaded(categoryEl.value); + const folders = this.getAvailableSubfolders(categoryEl.value); + const suggestion = this.getSuggestedCivitaiSubfolder(missing, categoryEl.value, folders); + if (suggestion) { + subfolderEl.value = suggestion; + } + } + + async ensureDownloadSubfoldersLoaded(category = '') { + const key = (category || '').trim().toLowerCase(); + if (!key) return []; + if (this.downloadSubfolders.has(key)) { + return this.downloadSubfolders.get(key) || []; + } + + try { + const resp = await api.fetchApi(`/model_linker/subfolders/${encodeURIComponent(key)}`); + if (!resp.ok) throw new Error(`HTTP ${resp.status}`); + const subfolders = await resp.json(); + const list = Array.isArray(subfolders) ? subfolders : []; + this.downloadSubfolders.set(key, list); + return list; + } catch (e) { + console.warn(`Model Linker: could not load subfolders for ${key}`, e); + this.downloadSubfolders.set(key, []); + return []; + } + } + + renderDownloadTargetControls(missing, defaultCategory = 'checkpoints') { + const selectId = `download-category-${missing.node_id}-${missing.widget_index}`; + const subfolderId = `download-subfolder-${missing.node_id}-${missing.widget_index}`; + const subfolderListId = `download-subfolder-list-${missing.node_id}-${missing.widget_index}`; + const options = this.getDownloadCategoryOptions(defaultCategory) + .map(category => ``) + .join(''); + + let html = `
`; + html += `
`; + html += ``; + html += ``; + html += `
`; + html += `
`; + html += ``; + html += `
`; + html += ``; + html += `
`; + html += `
`; + html += `
`; + html += `
`; + return html; + } + + getDownloadTargetSelection(missing, fallbackCategory = 'checkpoints') { + const categoryEl = this.contentElement?.querySelector(`#download-category-${missing.node_id}-${missing.widget_index}`); + const subfolderEl = this.contentElement?.querySelector(`#download-subfolder-${missing.node_id}-${missing.widget_index}`); + return { + category: categoryEl?.value || fallbackCategory || 'checkpoints', + subfolder: (subfolderEl?.value || '').trim() + }; + } + + wireDownloadTargetAutocomplete(container, missing) { + const categoryEl = container.querySelector(`#download-category-${missing.node_id}-${missing.widget_index}`); + const subfolderEl = container.querySelector(`#download-subfolder-${missing.node_id}-${missing.widget_index}`); + const listEl = container.querySelector(`#download-subfolder-list-${missing.node_id}-${missing.widget_index}`); + if (!categoryEl || !subfolderEl || !listEl) return; + + const renderOptions = (targetEl, values, onSelect) => { + if (!values.length) { + targetEl.innerHTML = ''; + targetEl.style.display = 'none'; + return; + } + + targetEl.innerHTML = values + .slice(0, 50) + .map(value => `
${value}
`) + .join(''); + + targetEl.style.display = 'block'; + + targetEl.querySelectorAll('.ml-download-target-option').forEach(option => { + option.addEventListener('mousedown', (event) => { + event.preventDefault(); + const value = decodeURIComponent(option.dataset.value || ''); + onSelect(value); + targetEl.style.display = 'none'; + }); + }); + }; + + const populateSubfolderOptions = async (filterText = '') => { + const filter = (filterText || '').toLowerCase(); + await this.ensureDownloadSubfoldersLoaded(categoryEl.value); + const folders = this.getAvailableSubfolders(categoryEl.value); + const filtered = filter + ? folders.filter(folder => folder.toLowerCase().includes(filter)) + : folders; + + renderOptions(listEl, filtered, (value) => { + subfolderEl.value = value; + }); + }; + + const hideList = (targetEl) => { + setTimeout(() => { + targetEl.style.display = 'none'; + }, 150); + }; + + categoryEl.addEventListener('change', () => { + populateSubfolderOptions(subfolderEl.value); + this.applySuggestedCivitaiSubfolder(missing, categoryEl, subfolderEl); + }); + + subfolderEl.addEventListener('focus', () => { + populateSubfolderOptions(subfolderEl.value); + }); + + subfolderEl.addEventListener('input', () => { + populateSubfolderOptions(subfolderEl.value); + }); + + subfolderEl.addEventListener('blur', () => hideList(listEl)); + this.applySuggestedCivitaiSubfolder(missing, categoryEl, subfolderEl); + } + getStoredTokens() { return { civitai_key: localStorage.getItem('modelLinker.civitaiApiKey') || '', @@ -3237,6 +3529,7 @@ class LinkerManagerDialog extends ComfyDialog { // Ensure all models are loaded for dropdown await this.ensureAllModelsLoaded(); + await this.ensureDownloadDirectoriesLoaded(); // Always default to Missing Models tab when opening dialog if (this.activeTab !== 'missing') { @@ -3579,6 +3872,8 @@ class LinkerManagerDialog extends ComfyDialog { if (sourceButtons.length > 0) { this.syncSearchSourceUi(missing, container); } + + this.wireDownloadTargetAutocomplete(container, missing); // Wire Locate button (only for top-level nodes) const locateId = `locate-${missing.node_id}-${missing.widget_index}`; @@ -3818,6 +4113,7 @@ class LinkerManagerDialog extends ComfyDialog { html += `🔍 Search Online`; html += ``; html += `
Selected source: Everything
`; + html += this.renderDownloadTargetControls(missing, missing.category || 'checkpoints'); html += ``; html += `
`; } @@ -4306,7 +4602,9 @@ class LinkerManagerDialog extends ComfyDialog { // Use filename from download source if available (may be different from original) const originalFilename = missing.original_path?.split('/').pop()?.split('\\').pop() || 'model.safetensors'; const filename = source.filename || originalFilename; - const category = source.directory || missing.category || 'checkpoints'; + const targetSelection = this.getDownloadTargetSelection(missing, source.directory || missing.category || 'checkpoints'); + const category = targetSelection.category; + const subfolder = targetSelection.subfolder; const progressId = `download-progress-${missing.node_id}-${missing.widget_index}`; const progressDiv = this.contentElement?.querySelector(`#${progressId}`); const downloadBtn = this.contentElement?.querySelector(`#download-${missing.node_id}-${missing.widget_index}`); @@ -4346,6 +4644,7 @@ class LinkerManagerDialog extends ComfyDialog { url: source.url, filename: filename, category: category, + subfolder: subfolder, hf_token: tokens.hf_token, civitai_key: tokens.civitai_key }) @@ -4676,7 +4975,9 @@ class LinkerManagerDialog extends ComfyDialog { missing.civitai_info = { model_name: data.civitai.name, version_name: data.civitai.version_name, - expected_filename: data.civitai.filename + expected_filename: data.civitai.filename, + base_model: data.civitai.base_model, + tags: data.civitai.tags || [] }; missing.download_source = { source: 'civitai', @@ -4689,7 +4990,9 @@ class LinkerManagerDialog extends ComfyDialog { size: data.civitai.size, model_id: data.civitai.model_id || modelId, version_id: data.civitai.version_id || versionId, - model_url: data.civitai.url || `https://civitai.com/models/${modelId}?modelVersionId=${versionId}` + model_url: data.civitai.url || `https://civitai.com/models/${modelId}?modelVersionId=${versionId}`, + base_model: data.civitai.base_model, + tags: data.civitai.tags || [] }; downloadEl.outerHTML = this.renderDownloadSourceSection(missing, missing.download_source); @@ -4699,6 +5002,9 @@ class LinkerManagerDialog extends ComfyDialog { this.downloadModel(missing); }); } + if (this.contentElement) { + this.wireDownloadTargetAutocomplete(this.contentElement, missing); + } this.refreshUrnLocalMatches(missing); } } else if (downloadEl) { @@ -4855,6 +5161,7 @@ class LinkerManagerDialog extends ComfyDialog { const progressId = `download-progress-${missing.node_id}-${missing.widget_index}`; const progressDiv = this.contentElement?.querySelector(`#${progressId}`); const tokens = this.getStoredTokens(); + const targetSelection = this.getDownloadTargetSelection(missing, category || missing.category || 'checkpoints'); try { btn.disabled = true; @@ -4886,7 +5193,8 @@ class LinkerManagerDialog extends ComfyDialog { body: JSON.stringify({ url, filename, - category, + category: targetSelection.category, + subfolder: targetSelection.subfolder, hf_token: tokens.hf_token, civitai_key: tokens.civitai_key })