diff --git a/README.md b/README.md index 934473b..604108f 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ candidate direct-download URLs from Civitai and Hugging Face. - Downloads selected Civitai, Hugging Face, or public Quark-share candidates into the matching configured ComfyUI model folder. - Shows candidate file sizes and loads successfully downloaded files into their workflow nodes. +- Shows live downloaded bytes, total size, and percentage while a model download is running. - Searches the configured Quark Netdisk libraries directly for missing model filenames. - Ignores serialized workflow metadata, URLs, free-text fields, and generic `model` fields on non-loader custom nodes to prevent false missing-model reports. diff --git a/find_models.py b/find_models.py index 39e6278..b49e031 100644 --- a/find_models.py +++ b/find_models.py @@ -5,6 +5,8 @@ import os import re import subprocess import tempfile +import time +import uuid from difflib import SequenceMatcher from pathlib import Path from typing import Any @@ -73,6 +75,8 @@ KNOWN_MODEL_SOURCES = { "confidence": 1.0, }, } +DOWNLOAD_JOBS: dict[str, dict[str, Any]] = {} +DOWNLOAD_TASKS: dict[str, asyncio.Task[Any]] = {} def _safe_query(name: str) -> str: @@ -458,9 +462,7 @@ async def _quark_download_url(session: aiohttp.ClientSession, payload: dict[str, raise web.HTTPBadGateway(text=f"夸克拒绝直链下载：{detail}") -@PromptServer.instance.routes.post("/findmodels/download") -async def download_model(request: web.Request) -> web.Response: - payload = await request.json() +async def _download_model_payload(payload: dict[str, Any], progress: dict[str, Any] | None = None) -> dict[str, Any]: url = str(payload.get("url", "")).strip() quark = payload.get("quark") category = str(payload.get("category", "")).strip() @@ -478,6 +480,15 @@ async def download_model(request: web.Request) -> web.Response: temp_path: Path | None = None try: + if progress is not None: + progress.update( + status="downloading", + filename=filename, + category=category, + downloaded=0, + total=expected_size, + updated_at=time.time(), + ) fd, temp_name = tempfile.mkstemp(prefix=f".{filename}.", suffix=".part", dir=target_dir) os.close(fd) temp_path = Path(temp_name) @@ -495,9 +506,21 @@ async def download_model(request: web.Request) -> web.Response: content_type = response.headers.get("Content-Type", "").lower() if "text/html" in content_type or "application/json" in content_type: raise web.HTTPBadGateway(text="下载地址返回了网页或错误信息，而不是模型文件") + response_size = _size_value( + response.headers.get("X-Linked-Size") + or response.headers.get("X-File-Size") + or response.headers.get("Content-Length") + ) + total = expected_size or response_size + downloaded = 0 + if progress is not None: + progress.update(total=total, updated_at=time.time()) with temp_path.open("wb") as output: async for chunk in response.content.iter_chunked(1024 * 1024): output.write(chunk) + downloaded += len(chunk) + if progress is not None: + progress.update(downloaded=downloaded, total=total, updated_at=time.time()) if not _is_model_payload(temp_path): raise web.HTTPBadGateway(text="下载结果不是有效模型文件，可能是登录页面、错误信息或 Git LFS 指针") temp_path.replace(target) @@ -514,13 +537,69 @@ async def download_model(request: web.Request) -> web.Response: if temp_path and temp_path.exists(): temp_path.unlink() - return web.json_response( - { - "downloaded": True, - "filename": filename, - "relative_name": filename, - "category": category, - "path": str(target), - "size": target.stat().st_size, - } - ) + return { + "downloaded": True, + "filename": filename, + "relative_name": filename, + "category": category, + "path": str(target), + "size": target.stat().st_size, + } + + +async def _run_download_job(job_id: str, payload: dict[str, Any]) -> None: + progress = DOWNLOAD_JOBS[job_id] + try: + result = await _download_model_payload(payload, progress) + progress.update(status="completed", result=result, downloaded=result["size"], total=result["size"]) + except Exception as error: + progress.update(status="failed", error=getattr(error, "text", None) or str(error)) + finally: + progress["updated_at"] = time.time() + DOWNLOAD_TASKS.pop(job_id, None) + + +def _purge_download_jobs() -> None: + cutoff = time.time() - 3600 + for job_id, job in list(DOWNLOAD_JOBS.items()): + if job.get("status") in {"completed", "failed"} and job.get("updated_at", 0) < cutoff: + DOWNLOAD_JOBS.pop(job_id, None) + + +@PromptServer.instance.routes.post("/findmodels/download/start") +async def start_download_model(request: web.Request) -> web.Response: + payload = await request.json() + _safe_filename(str(payload.get("filename", ""))) + category = str(payload.get("category", "")).strip() + _target_directory(category) + if not payload.get("quark") and not _allowed_download_url(str(payload.get("url", "")).strip()): + raise web.HTTPBadRequest(text="Only approved HTTPS model providers are allowed") + _purge_download_jobs() + job_id = uuid.uuid4().hex + DOWNLOAD_JOBS[job_id] = { + "id": job_id, + "status": "queued", + "filename": basename(str(payload.get("filename", ""))), + "category": category, + "downloaded": 0, + "total": _size_value(payload.get("size")), + "created_at": time.time(), + "updated_at": time.time(), + } + DOWNLOAD_TASKS[job_id] = asyncio.create_task(_run_download_job(job_id, payload)) + return web.json_response(DOWNLOAD_JOBS[job_id]) + + +@PromptServer.instance.routes.post("/findmodels/download/progress") +async def download_model_progress(request: web.Request) -> web.Response: + payload = await request.json() + job_id = str(payload.get("job_id", "")).strip() + job = DOWNLOAD_JOBS.get(job_id) + if not job: + raise web.HTTPNotFound(text="Download task not found") + return web.json_response(job) + + +@PromptServer.instance.routes.post("/findmodels/download") +async def download_model(request: web.Request) -> web.Response: + return web.json_response(await _download_model_payload(await request.json())) diff --git a/model_finder.py b/model_finder.py index fa49cfe..be871ce 100644 --- a/model_finder.py +++ b/model_finder.py @@ -63,7 +63,7 @@ EXPLICIT_MODEL_WIDGETS = ( ) CATEGORY_HINTS = ( - ("diffusion_models", ("fantasytalking", "fantasyportrait")), + ("diffusion_models", ("fantasytalking", "fantasyportrait", "wanvideo")), ("upscale_models", ("rife", "frame_interpolation", "film", "upscale", "esrgan", "super_resolution", "superresolution")), ("controlnet", ("controlnet", "control_net", "control")), ("embeddings", ("embedding", "textual_inversion")), @@ -168,7 +168,10 @@ def extract_references(payload: Any) -> list[Reference]: node_hint = re.sub(r"[^a-z0-9_]+", "", str(node_type or "").lower()) if node_id is not None: explicit_widget = widget_hint in EXPLICIT_MODEL_WIDGETS - loader_widget = "loader" in node_hint and any(term in widget_hint for term in MODEL_WIDGET_HINTS) + # Custom loaders frequently translate the widget label (for example + # "模型"), so the file extension plus the loader node type is the + # reliable signal. Non-loader custom nodes remain excluded. + loader_widget = "loader" in node_hint if not explicit_widget and not loader_widget and not model_selector: return ref = Reference( @@ -178,7 +181,8 @@ def extract_references(payload: Any) -> list[Reference]: widget=None if widget is None else str(widget), node_type=None if node_type is None else str(node_type), strict=model_selector - or any(term in re.sub(r"[^a-z0-9_]+", "", str(node_type or "").lower()) for term in STANDARD_LOADER_HINTS), + or "loader" in node_hint + or any(term in node_hint for term in STANDARD_LOADER_HINTS), ) found[(ref.name.lower(), ref.category, ref.node_id, ref.widget)] = ref diff --git a/pyproject.toml b/pyproject.toml index 3f64075..c8cc10f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "comfyui-findmodels" -version = "1.9.1" +version = "1.10.0" description = "Find and adapt missing models in the current ComfyUI workflow" readme = "README.md" requires-python = ">=3.10" diff --git a/tests/test_model_finder.py b/tests/test_model_finder.py index 18c84ca..d3303fc 100644 --- a/tests/test_model_finder.py +++ b/tests/test_model_finder.py @@ -180,6 +180,41 @@ class ModelFinderTests(unittest.TestCase): result = analyze(payload, lambda category: []) self.assertEqual(result["models"][0]["category"], "diffusion_models") + def test_detects_missing_model_in_custom_loader_with_translated_widget(self): + payload = { + "nodes": [{ + "id": 22, + "type": "WanVideo Model Loader", + "widgets": [{ + "name": "模型", + "value": "Wan/Wan14Bi2vFusioniX_fp8.safetensors", + "model_selector": False, + }], + }] + } + result = analyze(payload, lambda category: []) + self.assertEqual(result["summary"]["missing"], 1) + self.assertEqual(result["models"][0]["category"], "diffusion_models") + self.assertEqual(result["models"][0]["name"], "Wan/Wan14Bi2vFusioniX_fp8.safetensors") + + def test_hides_installed_model_in_custom_loader_with_translated_widget(self): + payload = { + "nodes": [{ + "id": 22, + "type": "WanVideo Model Loader", + "widgets": [{ + "name": "模型", + "value": "Wan/Wan14Bi2vFusioniX_fp8.safetensors", + "model_selector": False, + }], + }] + } + result = analyze( + payload, + lambda category: ["Wan/Wan14Bi2vFusioniX_fp8.safetensors"] if category == "diffusion_models" else [], + ) + self.assertEqual(result["models"], []) + if __name__ == "__main__": unittest.main() diff --git a/web/find_models.css b/web/find_models.css index 5d8b61c..b8db015 100644 --- a/web/find_models.css +++ b/web/find_models.css @@ -30,5 +30,8 @@ .fm-sources { display: grid; gap: 5px; margin-top: 7px; } .fm-node-candidates { display: grid; gap: 5px; margin-top: 7px; } .fm-source { display: grid; grid-template-columns: 1fr auto; gap: 6px; align-items: center; } +.fm-download-progress { + grid-column: 1 / -1; min-height: 16px; color: #8fd3a8; font-size: 12px; +} .fm-sources a { color: #7fb9ff; overflow-wrap: anywhere; } #find-models-panel button { padding: 5px 8px; color: #eee; background: #333; border: 1px solid #666; border-radius: 5px; cursor: pointer; } diff --git a/web/find_models.js b/web/find_models.js index 107d89d..42f386f 100644 --- a/web/find_models.js +++ b/web/find_models.js @@ -19,13 +19,39 @@ function formatSize(value) { return `${(size / (1024 ** index)).toFixed(index > 1 ? 2 : 0)} ${units[index]}`; } +function downloadProgressText(job) { + const downloaded = formatSize(job.downloaded); + const total = Number(job.total); + if (!Number.isFinite(total) || total <= 0) return `已下载 ${downloaded}`; + const percent = Math.min(100, Math.round((Number(job.downloaded || 0) / total) * 100)); + return `已下载 ${downloaded} / ${formatSize(total)} · ${percent}%`; +} + +async function waitForDownload(jobId, progressElement) { + while (true) { + const job = await post("/findmodels/download/progress", { job_id: jobId }); + progressElement.textContent = downloadProgressText(job); + if (job.status === "completed") return job.result; + if (job.status === "failed") throw new Error(job.error || "下载失败"); + await new Promise((resolve) => window.setTimeout(resolve, 500)); + } +} + function workflowSnapshot() { return { nodes: app.graph?._nodes?.map((node) => ({ id: node.id, type: node.type, widgets: (node.widgets || []).flatMap((widget) => { - const values = Array.isArray(widget.options?.values) ? widget.options.values : []; + let rawValues = widget.options?.values; + if (typeof rawValues === "function") { + try { + rawValues = rawValues(); + } catch (error) { + console.debug("[ComfyUI_FindModels] Unable to inspect combo values", widget.name, error); + } + } + const values = Array.isArray(rawValues) ? rawValues : []; const isModelSelector = values.some((value) => typeof value === "string" && /\.(bin|ckpt|gguf|onnx|pt|pth|safetensors)$/i.test(value), ); @@ -91,6 +117,7 @@ function sourceHtml(item, model) { data-size="${escapeHtml(item.size || "")}" data-category="${escapeHtml(model.category)}" data-node-id="${escapeHtml(model.node_id)}" data-widget="${escapeHtml(model.widget)}" data-original="${escapeHtml(model.name)}">下载到模型目录 +
`; } @@ -184,14 +211,16 @@ function ensurePanel() { if (!downloadButton) return; downloadButton.disabled = true; downloadButton.textContent = "正在下载…"; + const progressElement = downloadButton.closest(".fm-source").querySelector(".fm-download-progress"); try { - const downloaded = await post("/findmodels/download", { + const job = await post("/findmodels/download/start", { url: downloadButton.dataset.download, quark: downloadButton.dataset.quarkDownload ? JSON.parse(downloadButton.dataset.quarkDownload) : null, filename: downloadButton.dataset.filename, size: downloadButton.dataset.size || null, category: downloadButton.dataset.category, }); + const downloaded = await waitForDownload(job.id, progressElement); downloadButton.textContent = "下载完成，正在加载"; await applyMatch({ node_id: downloadButton.dataset.nodeId,