From 494dabf11b9d177663640ad2da534b234bbcbda4 Mon Sep 17 00:00:00 2001 From: DrHepa Date: Fri, 21 Aug 2026 07:54:11 +0200 Subject: [PATCH] feat(models): support multiple Hugging Face sources per node --- README.md | 36 +++ api/routers/model.py | 164 ++++++++++- api/services/generator_registry.py | 27 +- api/services/model_sources.py | 259 ++++++++++++++++++ api/tests/test_generator_registry.py | 56 ++++ api/tests/test_model_router.py | 187 +++++++++++++ api/tests/test_model_sources.py | 103 +++++++ .../main/extension-install-utils.test.mjs | 49 ++++ electron/main/extension-install-utils.ts | 21 +- electron/main/ipc-handlers.ts | 152 ++++++++-- electron/main/model-download-plan.test.mjs | 92 +++++++ electron/main/model-download-plan.ts | 139 ++++++++++ electron/main/model-download-preload.test.mjs | 55 ++++ electron/main/model-downloader.ts | 30 +- electron/main/model-sources.test.mjs | 109 ++++++++ electron/main/model-sources.ts | 190 +++++++++++++ electron/preload/electron-api.ts | 6 +- src/areas/models/ModelsPage.tsx | 44 +-- .../models/components/ExtensionDrawer.tsx | 7 +- .../models/components/extensionShared.tsx | 6 +- src/areas/models/utils.test.mjs | 29 ++ src/areas/models/utils.ts | 24 ++ src/shared/types/electron.d.ts | 6 +- 23 files changed, 1731 insertions(+), 60 deletions(-) create mode 100644 api/services/model_sources.py create mode 100644 api/tests/test_model_router.py create mode 100644 api/tests/test_model_sources.py create mode 100644 electron/main/model-download-plan.test.mjs create mode 100644 electron/main/model-download-plan.ts create mode 100644 electron/main/model-download-preload.test.mjs create mode 100644 electron/main/model-sources.test.mjs create mode 100644 electron/main/model-sources.ts diff --git a/README.md b/README.md index 11184ddc..9414106e 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,42 @@ Modly supports external model and process extensions. Each extension is a GitHub ![Install models](docs/install-models.png) +### Multiple Hugging Face repositories per model node + +A model node whose weights are split across repositories can declare +`model_sources`. Modly validates every source, downloads them sequentially in +one Models-page action, and considers the node installed only when every +declared check exists. + +```json +{ + "id": "generate", + "model_sources": [ + { + "id": "primary", + "provider": "huggingface", + "repo_id": "org/main-model", + "destination": ".", + "checks": ["model.safetensors"] + }, + { + "id": "encoder", + "provider": "huggingface", + "repo_id": "org/encoder", + "revision": "v1.0", + "destination": "auxiliary/encoder", + "include_prefixes": ["config.json", "model.safetensors"], + "checks": ["config.json", "model.safetensors"] + } + ] +} +``` + +`destination`, filters, and checks use safe POSIX paths relative to the node's +model directory. The only supported provider is `huggingface`. Existing nodes +that use `hf_repo`, `download_check`, `hf_include_prefixes`, and +`hf_skip_prefixes` keep their original behavior. + --- ## Workflows diff --git a/api/routers/model.py b/api/routers/model.py index 4f04718b..509fadda 100644 --- a/api/routers/model.py +++ b/api/routers/model.py @@ -8,9 +8,16 @@ from typing import Optional from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen -from fastapi import APIRouter, HTTPException +from fastapi import APIRouter, HTTPException, Request as FastAPIRequest from fastapi.responses import StreamingResponse from services.generator_registry import generator_registry, MODELS_DIR +from services.model_sources import ( + normalize_model_sources, + resolve_download_path, + resolve_model_root, + resolve_source_destination, + validate_source_file_plan, +) router = APIRouter(tags=["model"]) @@ -97,7 +104,7 @@ async def unload_all_models(): return {"unloaded": True} -@router.post("/unload/{model_id}") +@router.post("/unload/{model_id:path}") async def unload_model(model_id: str): """Unloads a model from memory so its files can be safely deleted.""" try: @@ -122,6 +129,159 @@ async def cancel_hf_download(model_id: str): return {"cancelled": True} +@router.post("/hf-download-sources") +async def hf_download_sources(request: FastAPIRequest, model_id: str): + """Download all Hugging Face sources declared for one model node.""" + try: + body = await request.json() + if not isinstance(body, dict): + raise ValueError("Request body must be an object") + sources = normalize_model_sources({"model_sources": body.get("sources")}) + if sources is None: + raise ValueError("sources are required") + model_root = resolve_model_root(MODELS_DIR, model_id) + destinations = { + source["id"]: resolve_source_destination( + MODELS_DIR, model_id, source["destination"] + ) + for source in sources + } + except (TypeError, ValueError) as exc: + raise HTTPException(400, str(exc)) from exc + + authorization = request.headers.get("authorization", "") + hf_token = ( + authorization[7:].strip() + if authorization.lower().startswith("bearer ") + else os.environ.get("HUGGING_FACE_HUB_TOKEN") + or os.environ.get("HF_TOKEN") + or None + ) + control = _new_download_control(model_id) + + async def stream(): + loop = asyncio.get_running_loop() + + def _fmt(data: dict) -> str: + return f"data: {json.dumps(data)}\n\n" + + try: + yield _fmt({"percent": 0, "status": "Listing repository files..."}) + files_by_source: dict[str, list[str]] = {} + + for source in sources: + _check_download_control(control) + + def _list_files(current=source): + from huggingface_hub import list_repo_files + listed = list_repo_files( + current["repo_id"], + revision=current.get("revision"), + token=hf_token, + ) + include = current.get("include_prefixes", []) + skip = current.get("skip_prefixes", []) + return [ + filename for filename in listed + if (not include or any(filename.startswith(prefix) for prefix in include)) + if not any(filename.startswith(prefix) for prefix in skip) + ] + + files = await loop.run_in_executor(None, _list_files) + if not files: + raise RuntimeError( + f'No files found in Hugging Face repo: {source["repo_id"]}' + ) + destination = destinations[source["id"]] + for filename in files: + resolve_download_path(destination, filename) + files_by_source[source["id"]] = files + + validate_source_file_plan(sources, files_by_source) + planned_files = [ + (source, filename) + for source in sources + for filename in files_by_source[source["id"]] + ] + total = len(planned_files) + yield _fmt({"percent": 1, "status": f"Downloading {total} files..."}) + + from huggingface_hub import hf_hub_url + + for index, (source, filename) in enumerate(planned_files): + _check_download_control(control) + display_file = f'{source["id"]}/{filename}' + base_percent = 1 + round(index / total * 94) + yield _fmt({ + "percent": base_percent, + "file": display_file, + "fileIndex": index + 1, + "totalFiles": total, + "status": f"Starting {display_file}", + "bytesDownloaded": 0, + "stalledSeconds": 0, + }) + + queue: asyncio.Queue[dict] = asyncio.Queue() + + def _progress(message: dict) -> None: + message["file"] = display_file + loop.call_soon_threadsafe(queue.put_nowait, message) + + url = hf_hub_url( + repo_id=source["repo_id"], + filename=filename, + revision=source.get("revision"), + ) + future = loop.run_in_executor( + None, + lambda: _download_file_streamed( + url=url, + filename=filename, + dest_dir=str(destinations[source["id"]]), + file_index=index + 1, + total_files=total, + base_percent=base_percent, + progress_cb=_progress, + control=control, + token=hf_token, + ), + ) + while not future.done(): + try: + message = await asyncio.wait_for(queue.get(), timeout=2.0) + except asyncio.TimeoutError: + continue + yield _fmt(message) + + final_size = await future + _check_download_control(control) + yield _fmt({ + "percent": 1 + round((index + 1) / total * 94), + "file": display_file, + "fileIndex": index + 1, + "totalFiles": total, + "status": "Downloaded", + "bytesDownloaded": final_size, + "stalledSeconds": 0, + }) + + yield _fmt({"percent": 100, "status": "done"}) + except DownloadPaused: + yield _fmt({"paused": True, "status": "paused"}) + except DownloadCancelled: + for part in model_root.rglob("*.part"): + part.unlink(missing_ok=True) + yield _fmt({"cancelled": True, "status": "cancelled"}) + except Exception as exc: + yield _fmt({"error": str(exc)}) + finally: + if _download_controls.get(model_id) is control: + _download_controls.pop(model_id, None) + + return StreamingResponse(stream(), media_type="text/event-stream") + + @router.get("/hf-download") async def hf_download( repo_id: str, diff --git a/api/services/generator_registry.py b/api/services/generator_registry.py index ecb6bb6c..348a42cb 100644 --- a/api/services/generator_registry.py +++ b/api/services/generator_registry.py @@ -25,6 +25,7 @@ from services.generators.base import BaseGenerator from services.extension_process import ExtensionProcess, _venv_python +from services.model_sources import model_sources_are_downloaded, normalize_model_sources # ------------------------------------------------------------------ # # Global paths @@ -428,6 +429,9 @@ def _discover_extensions( ext_id = manifest["id"] class_name = manifest["generator_class"] + if "model_sources" in manifest: + raise ValueError("model_sources must be declared on a model node") + if ext_id != ext_dir.name: message = ( f"Extension folder '{ext_dir.name}' declares mismatched " @@ -523,6 +527,7 @@ def _discover_extensions( if nodes: for node in nodes: + model_sources = normalize_model_sources(node) node_manifest = { **manifest, "id": f"{ext_id}/{node['id']}", @@ -537,6 +542,8 @@ def _discover_extensions( "input": node.get("input", "image"), "output": node.get("output", "mesh"), } + if model_sources is not None: + node_manifest["model_sources"] = model_sources full_id = f"{ext_id}/{node['id']}" result[full_id] = (cls_or_None, node_manifest, ext_dir, legacy_context) if subprocess_mode: @@ -699,8 +706,14 @@ def get_active(self) -> BaseGenerator: """Returns the active generator. Downloads and loads if necessary.""" self._assert_not_quarantined(self._active_id) gen = self._generators[self._active_id] + downloaded = self._is_downloaded(self._active_id, gen) + if "model_sources" in self._manifests[self._active_id] and not downloaded: + raise RuntimeError( + "Model sources are incomplete. Download this node's weights " + "from the Modly Models page before generation." + ) if not gen.is_loaded(): - if not gen.is_downloaded(): + if not downloaded: if isinstance(gen, ExtensionProcess): # Let the subprocess handle its own download logic during # load() — some extensions (e.g. mv-adapter) need custom @@ -743,13 +756,21 @@ def switch_model(self, model_id: str) -> None: # Status # ------------------------------------------------------------------ # + def _is_downloaded(self, model_id: str, gen: BaseGenerator) -> bool: + manifest = self._manifests[model_id] + if "model_sources" in manifest: + return model_sources_are_downloaded( + MODELS_DIR, model_id, manifest["model_sources"] + ) + return gen.is_downloaded() + def active_status(self) -> dict: gen = self._generators[self._active_id] manifest = self._manifests[self._active_id] return { "id": self._active_id, "name": manifest.get("name", gen.DISPLAY_NAME), - "downloaded": gen.is_downloaded(), + "downloaded": self._is_downloaded(self._active_id, gen), "loaded": gen.is_loaded(), } @@ -765,7 +786,7 @@ def all_status(self) -> list: "vram_gb": manifest.get("vram_gb", gen.VRAM_GB), "hf_repo": manifest.get("hf_repo", ""), "tags": manifest.get("tags", []), - "downloaded": gen.is_downloaded(), + "downloaded": self._is_downloaded(model_id, gen), "loaded": gen.is_loaded(), "active": model_id == self._active_id, }) diff --git a/api/services/model_sources.py b/api/services/model_sources.py new file mode 100644 index 00000000..5988ae02 --- /dev/null +++ b/api/services/model_sources.py @@ -0,0 +1,259 @@ +"""Validation and readiness helpers for manifest-declared Hugging Face sources.""" + +from __future__ import annotations + +import re +import unicodedata +from pathlib import Path +from typing import Any + + +_SAFE_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") +_WINDOWS_DEVICE = re.compile( + r"^(?:con|prn|aux|nul|com[1-9]|lpt[1-9])(?:\..*)?$", re.IGNORECASE +) +_WINDOWS_UNSAFE = re.compile(r'[<>"|?*\x00-\x1f]') + + +def _portable_segment(value: str, field: str) -> str: + if ( + not value + or value in {".", ".."} + or value.endswith((".", " ")) + or ":" in value + or _WINDOWS_UNSAFE.search(value) + or _WINDOWS_DEVICE.fullmatch(value) + ): + raise ValueError(f'{field} contains unsafe path segment "{value}"') + return value + + +def safe_source_id(value: Any, field: str = "model source id") -> str: + if ( + not isinstance(value, str) + or not value + or value != value.strip() + or _SAFE_ID.fullmatch(value) is None + ): + raise ValueError(f"{field} must be a safe non-empty identifier") + return _portable_segment(value, field) + + +def safe_relative_path(value: Any, field: str, *, allow_dot: bool = False) -> str: + if not isinstance(value, str) or not value or value != value.strip(): + raise ValueError(f"{field} must be a non-empty relative path") + if allow_dot and value == ".": + return value + if value == "." or value.startswith("/") or "\\" in value: + raise ValueError(f"{field} must be a safe relative POSIX path") + for part in value.split("/"): + _portable_segment(part, field) + return value + + +def _safe_prefix(value: Any, field: str) -> str: + path = value[:-1] if isinstance(value, str) and value.endswith("/") else value + safe_relative_path(path, field) + return value + + +def _prefixes(value: Any, field: str) -> list[str] | None: + if not isinstance(value, list): + raise ValueError(f"{field} must be an array") + return [_safe_prefix(entry, f"{field}[{index}]") for index, entry in enumerate(value)] + + +def _safe_repo_id(value: Any, field: str) -> str: + if not isinstance(value, str) or not value or value != value.strip() or "\\" in value: + raise ValueError(f"{field} must be a non-empty Hugging Face repository id") + parts = value.split("/") + if len(parts) > 2 or any( + part in {"", ".", ".."} or _SAFE_ID.fullmatch(part) is None for part in parts + ): + raise ValueError(f"{field} is not a safe Hugging Face repository id") + return value + + +def _safe_revision(value: Any, field: str) -> str | None: + if value is None: + return None + if ( + not isinstance(value, str) + or not value + or value != value.strip() + or value.startswith("/") + or "\\" in value + or "\0" in value + or any(part in {"", ".", ".."} for part in value.split("/")) + ): + raise ValueError(f"{field} must be a safe non-empty revision") + return value + + +def normalize_model_sources(node: dict[str, Any]) -> list[dict[str, Any]] | None: + """Validate only the new contract; legacy fields remain untouched.""" + if "model_sources" not in node: + return None + raw_sources = node["model_sources"] + if not isinstance(raw_sources, list) or not raw_sources: + raise ValueError("model_sources must be a non-empty array") + + aliases: dict[str, str] = {} + sources: list[dict[str, Any]] = [] + for index, raw in enumerate(raw_sources): + field = f"model_sources[{index}]" + if not isinstance(raw, dict): + raise ValueError(f"{field} must be an object") + source_id = safe_source_id(raw.get("id"), f"{field}.id") + alias = unicodedata.normalize("NFC", source_id).casefold() + if alias in aliases: + raise ValueError( + f'model source ids "{aliases[alias]}" and "{source_id}" are not portable-unique' + ) + aliases[alias] = source_id + if raw.get("provider") != "huggingface": + raise ValueError(f'{field}.provider must be "huggingface"') + checks = raw.get("checks") + if not isinstance(checks, list) or not checks: + raise ValueError(f"{field}.checks must be a non-empty array") + + source: dict[str, Any] = { + "id": source_id, + "provider": "huggingface", + "repo_id": _safe_repo_id(raw.get("repo_id"), f"{field}.repo_id"), + "destination": safe_relative_path( + raw.get("destination"), f"{field}.destination", allow_dot=True + ), + "checks": [ + safe_relative_path(check, f"{field}.checks[{check_index}]") + for check_index, check in enumerate(checks) + ], + } + revision = ( + _safe_revision(raw["revision"], f"{field}.revision") + if "revision" in raw + else None + ) + if "revision" in raw and revision is None: + raise ValueError(f"{field}.revision must be a safe non-empty revision") + include = ( + _prefixes(raw["include_prefixes"], f"{field}.include_prefixes") + if "include_prefixes" in raw + else None + ) + skip = ( + _prefixes(raw["skip_prefixes"], f"{field}.skip_prefixes") + if "skip_prefixes" in raw + else None + ) + if revision is not None: + source["revision"] = revision + if include is not None: + source["include_prefixes"] = include + if skip is not None: + source["skip_prefixes"] = skip + sources.append(source) + return sources + + +def _path_has_symlink(root: Path, candidate: Path) -> bool: + root = root.absolute() + candidate = candidate.absolute() + try: + relative = candidate.relative_to(root) + except ValueError: + return True + current = root + if current.exists() and current.is_symlink(): + return True + for part in relative.parts: + current /= part + if current.exists() and current.is_symlink(): + return True + return False + + +def resolve_model_root(models_dir: Path, model_id: str) -> Path: + if not isinstance(model_id, str): + raise ValueError("Model id must be a string") + parts = model_id.split("/") + if len(parts) != 2: + raise ValueError("Model id must identify one extension node") + extension_id = safe_source_id(parts[0], "extension id") + node_id = safe_source_id(parts[1], "model node id") + root = models_dir.absolute() + candidate = root / extension_id / node_id + if _path_has_symlink(root, candidate): + raise ValueError("Model path resolves through a symlink") + try: + candidate.resolve().relative_to(root.resolve()) + except ValueError as exc: + raise ValueError("Model path escapes the models directory") from exc + return candidate + + +def resolve_source_destination(models_dir: Path, model_id: str, destination: str) -> Path: + model_root = resolve_model_root(models_dir, model_id) + safe_destination = safe_relative_path(destination, "destination", allow_dot=True) + candidate = model_root if safe_destination == "." else model_root.joinpath(*safe_destination.split("/")) + if _path_has_symlink(model_root, candidate): + raise ValueError("Source destination resolves through a symlink") + return candidate + + +def resolve_download_path(destination: Path, filename: str) -> Path: + safe_filename = safe_relative_path(filename, "Hugging Face repository file") + candidate = destination.joinpath(*safe_filename.split("/")) + if _path_has_symlink(destination, candidate): + raise ValueError("Download target resolves through a symlink") + return candidate + + +def model_sources_are_downloaded( + models_dir: Path, model_id: str, sources: list[dict[str, Any]] +) -> bool: + try: + model_root = resolve_model_root(models_dir, model_id) + if not model_root.is_dir(): + return False + for source in sources: + destination = resolve_source_destination( + models_dir, model_id, source["destination"] + ) + if not destination.is_dir(): + return False + for check in source["checks"]: + candidate = resolve_download_path(destination, check) + if not candidate.exists() or _path_has_symlink(model_root, candidate): + return False + return bool(sources) + except (KeyError, OSError, TypeError, ValueError): + return False + + +def validate_source_file_plan( + sources: list[dict[str, Any]], files_by_source: dict[str, list[str]] +) -> None: + """Reject cross-source aliases before the first file is written.""" + aliases: dict[str, tuple[str, str]] = {} + for source in sources: + source_id = source["id"] + destination = source["destination"] + for filename in files_by_source[source_id]: + safe_filename = safe_relative_path(filename, f'model source "{source_id}" file') + target = safe_filename if destination == "." else f"{destination}/{safe_filename}" + for value in (target, f"{target}.part"): + alias = unicodedata.normalize("NFC", value).casefold() + for previous_alias, (previous_source, previous_target) in aliases.items(): + if previous_source == source_id: + continue + if ( + alias == previous_alias + or alias.startswith(f"{previous_alias}/") + or previous_alias.startswith(f"{alias}/") + ): + raise ValueError( + "Model sources have a portable target collision: " + f'"{previous_source}:{previous_target}" and "{source_id}:{value}"' + ) + aliases[alias] = (source_id, value) diff --git a/api/tests/test_generator_registry.py b/api/tests/test_generator_registry.py index 7340b72c..ff9d090c 100644 --- a/api/tests/test_generator_registry.py +++ b/api/tests/test_generator_registry.py @@ -155,6 +155,62 @@ def test_legacy_generator_supports_eager_and_lazy_sibling_imports(self) -> None: self.registry.reload() self.assertNotIn(str(extension.resolve()), sys.path) + def test_declared_sources_block_generation_even_when_generator_overrides_readiness(self) -> None: + extension = self._make_extension("multi-source") + manifest = { + "id": "multi-source", + "name": "multi-source", + "type": "model", + "generator_class": "TestGenerator", + "nodes": [{ + "id": "generate", + "model_sources": [ + { + "id": "primary", + "provider": "huggingface", + "repo_id": "org/main", + "destination": ".", + "checks": ["main.bin"], + }, + { + "id": "encoder", + "provider": "huggingface", + "repo_id": "org/encoder", + "destination": "auxiliary/encoder", + "checks": ["encoder.bin"], + }, + ], + }], + } + (extension / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8") + (extension / "generator.py").write_text( + "\n".join([ + "from services.generators.base import BaseGenerator", + "class TestGenerator(BaseGenerator):", + " def is_downloaded(self): return True", + " def load(self): self._model = object()", + " def generate(self, image_bytes, params, progress_cb=None, cancel_event=None):", + " return self.outputs_dir / 'result.glb'", + ]), + encoding="utf-8", + ) + + self.registry.initialize() + self.registry._active_id = "multi-source/generate" + with self.assertRaisesRegex(RuntimeError, "Model sources are incomplete"): + self.registry.get_active() + self.assertFalse(self.registry.all_status()[0]["downloaded"]) + + model_root = self.models_dir / "multi-source" / "generate" + (model_root / "auxiliary" / "encoder").mkdir(parents=True) + (model_root / "main.bin").write_bytes(b"main") + (model_root / "auxiliary" / "encoder" / "encoder.bin").write_bytes(b"encoder") + self.assertIsNotNone(self.registry.get_active()) + self.assertTrue(self.registry.all_status()[0]["downloaded"]) + (model_root / "main.bin").unlink() + with self.assertRaisesRegex(RuntimeError, "Model sources are incomplete"): + self.registry.get_active() + def test_reload_preserves_legacy_path_owned_by_the_host(self) -> None: extension = self._make_extension("host-owned-path") self._write_manifest(extension, extension_id="host-owned-path") diff --git a/api/tests/test_model_router.py b/api/tests/test_model_router.py new file mode 100644 index 00000000..d0fe1279 --- /dev/null +++ b/api/tests/test_model_router.py @@ -0,0 +1,187 @@ +import asyncio +import json +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +from starlette.requests import Request + +import routers.model as model_router + + +SOURCES = [ + { + "id": "primary", + "provider": "huggingface", + "repo_id": "org/main", + "destination": ".", + "checks": ["main.bin"], + }, + { + "id": "encoder", + "provider": "huggingface", + "repo_id": "org/encoder", + "destination": "auxiliary/encoder", + "checks": ["encoder.bin"], + }, +] + + +def request_for(sources: list[dict]) -> Request: + body = json.dumps({"sources": sources}).encode() + sent = False + + async def receive(): + nonlocal sent + if sent: + return {"type": "http.disconnect"} + sent = True + return {"type": "http.request", "body": body, "more_body": False} + + return Request({ + "type": "http", + "method": "POST", + "path": "/model/hf-download-sources", + "headers": [(b"authorization", b"Bearer test-token")], + "query_string": b"", + "server": ("test", 80), + "client": ("test", 1), + "scheme": "http", + }, receive) + + +async def collect_events(response) -> list[dict]: + payload = "" + async for chunk in response.body_iterator: + payload += chunk.decode() if isinstance(chunk, bytes) else chunk + return [ + json.loads(block[6:]) + for block in payload.strip().split("\n\n") + if block.startswith("data: ") + ] + + +class MultiSourceRouterTests(unittest.TestCase): + def setUp(self) -> None: + self.tempdir = tempfile.TemporaryDirectory(prefix="modly-model-router-") + self.models_dir = Path(self.tempdir.name) / "models" + self.models_dir.mkdir() + self.old_models_dir = model_router.MODELS_DIR + model_router.MODELS_DIR = self.models_dir + self.old_hf_module = sys.modules.get("huggingface_hub") + + def tearDown(self) -> None: + model_router.MODELS_DIR = self.old_models_dir + model_router._download_controls.clear() + if self.old_hf_module is None: + sys.modules.pop("huggingface_hub", None) + else: + sys.modules["huggingface_hub"] = self.old_hf_module + self.tempdir.cleanup() + + def install_hf_stub(self, files: dict[str, list[str]], calls: list[str]) -> None: + module = types.ModuleType("huggingface_hub") + + def list_repo_files(repo_id, revision=None, token=None): + calls.append(f"list:{repo_id}:{revision}:{token}") + return files[repo_id] + + def hf_hub_url(repo_id, filename, revision=None): + return f"https://example.invalid/{repo_id}/{revision or 'main'}/{filename}" + + module.list_repo_files = list_repo_files + module.hf_hub_url = hf_hub_url + sys.modules["huggingface_hub"] = module + + def test_lists_every_source_before_sequential_download_with_monotonic_progress(self) -> None: + calls: list[str] = [] + controls: list[int] = [] + self.install_hf_stub({"org/main": ["main.bin"], "org/encoder": ["encoder.bin"]}, calls) + + def fake_download(**kwargs): + calls.append(f"download:{kwargs['dest_dir']}:{kwargs['filename']}") + controls.append(id(kwargs["control"])) + target = Path(kwargs["dest_dir"]) / kwargs["filename"] + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes(b"data") + kwargs["progress_cb"]({ + "percent": kwargs["base_percent"], + "file": kwargs["filename"], + "fileIndex": kwargs["file_index"], + "totalFiles": kwargs["total_files"], + "status": "Downloading...", + "bytesDownloaded": 4, + "stalledSeconds": 0, + }) + return 4 + + async def run(): + with patch.object(model_router, "_download_file_streamed", fake_download): + response = await model_router.hf_download_sources( + request_for(SOURCES), "pixal3d/generate" + ) + return await collect_events(response) + + events = asyncio.run(run()) + first_download = next(index for index, value in enumerate(calls) if value.startswith("download:")) + self.assertTrue(all(value.startswith("list:") for value in calls[:first_download])) + self.assertEqual(len(set(controls)), 1) + self.assertEqual([event["percent"] for event in events if "percent" in event], sorted( + event["percent"] for event in events if "percent" in event + )) + self.assertEqual(events[-1], {"percent": 100, "status": "done"}) + self.assertTrue((self.models_dir / "pixal3d/generate/main.bin").is_file()) + self.assertTrue((self.models_dir / "pixal3d/generate/auxiliary/encoder/encoder.bin").is_file()) + + def test_pause_cancel_and_resume_reuse_one_model_control(self) -> None: + calls: list[str] = [] + self.install_hf_stub({"org/main": ["main.bin"]}, calls) + source = [SOURCES[0]] + mode = "pause" + + def controlled_download(**kwargs): + target = Path(kwargs["dest_dir"]) / kwargs["filename"] + target.parent.mkdir(parents=True, exist_ok=True) + part = Path(f"{target}.part") + part.write_bytes(b"partial") + if mode == "pause": + kwargs["control"]["pause"].set() + model_router._check_download_control(kwargs["control"]) + if mode == "cancel": + kwargs["control"]["cancel"].set() + model_router._check_download_control(kwargs["control"]) + part.replace(target) + return target.stat().st_size + + async def one_run(): + with patch.object(model_router, "_download_file_streamed", controlled_download): + response = await model_router.hf_download_sources( + request_for(source), "pixal3d/generate" + ) + return await collect_events(response) + + paused = asyncio.run(one_run()) + self.assertTrue(paused[-1]["paused"]) + self.assertTrue((self.models_dir / "pixal3d/generate/main.bin.part").is_file()) + + mode = "cancel" + cancelled = asyncio.run(one_run()) + self.assertTrue(cancelled[-1]["cancelled"]) + self.assertFalse((self.models_dir / "pixal3d/generate/main.bin.part").exists()) + + mode = "resume" + resumed = asyncio.run(one_run()) + self.assertEqual(resumed[-1], {"percent": 100, "status": "done"}) + self.assertTrue((self.models_dir / "pixal3d/generate/main.bin").is_file()) + + def test_composite_model_unload_route_uses_path_converter(self) -> None: + paths = {route.path for route in model_router.router.routes} + self.assertIn("/unload/{model_id:path}", paths) + self.assertEqual(model_router.Request.__module__, "urllib.request") + + +if __name__ == "__main__": + unittest.main() diff --git a/api/tests/test_model_sources.py b/api/tests/test_model_sources.py new file mode 100644 index 00000000..72a68c71 --- /dev/null +++ b/api/tests/test_model_sources.py @@ -0,0 +1,103 @@ +import os +import tempfile +import unittest +from pathlib import Path + +from services.model_sources import ( + model_sources_are_downloaded, + normalize_model_sources, + resolve_model_root, + validate_source_file_plan, +) + + +def valid_node() -> dict: + return { + "model_sources": [ + { + "id": "primary", + "provider": "huggingface", + "repo_id": "org/main", + "destination": ".", + "checks": ["pipeline.json"], + }, + { + "id": "encoder", + "provider": "huggingface", + "repo_id": "org/encoder", + "revision": "refs/pr/1", + "destination": "auxiliary/encoder", + "include_prefixes": ["config.json", "weights/"], + "checks": ["config.json", "model.safetensors"], + }, + ] + } + + +class ModelSourcesTests(unittest.TestCase): + def test_validates_new_sources_without_reinterpreting_legacy_fields(self) -> None: + sources = normalize_model_sources(valid_node()) + self.assertEqual([source["id"] for source in sources or []], ["primary", "encoder"]) + self.assertIsNone(normalize_model_sources({ + "hf_repo": "legacy/repo", + "download_check": "../generate/model.safetensors", + "hf_skip_prefixes": ["weights/**"], + })) + + def test_rejects_unsafe_and_non_portable_declarations(self) -> None: + source = valid_node()["model_sources"][0] + for destination in ("../outside", "aux/CON", "aux/name.", "C:/models"): + with self.subTest(destination=destination), self.assertRaises(ValueError): + normalize_model_sources({ + "model_sources": [{**source, "destination": destination}] + }) + with self.assertRaisesRegex(ValueError, "provider"): + normalize_model_sources({ + "model_sources": [{**source, "provider": "url"}] + }) + with self.assertRaisesRegex(ValueError, "portable-unique"): + normalize_model_sources({ + "model_sources": [source, {**source, "id": "PRIMARY"}] + }) + with self.assertRaisesRegex(ValueError, "checks"): + normalize_model_sources({ + "model_sources": [{**source, "checks": []}] + }) + + def test_rejects_portable_cross_source_file_collisions(self) -> None: + sources = normalize_model_sources(valid_node()) or [] + with self.assertRaisesRegex(ValueError, "portable target collision"): + validate_source_file_plan(sources, { + "primary": ["Auxiliary/Encoder/model.safetensors"], + "encoder": ["model.safetensors"], + }) + + def test_requires_all_checks_and_rejects_symlinked_extension_ancestry(self) -> None: + sources = normalize_model_sources(valid_node()) or [] + with tempfile.TemporaryDirectory(prefix="modly-model-sources-") as tmp: + models = Path(tmp) / "models" + model_root = models / "pixal3d" / "generate" + encoder = model_root / "auxiliary" / "encoder" + encoder.mkdir(parents=True) + (model_root / "pipeline.json").write_text("{}", encoding="utf-8") + (encoder / "config.json").write_text("{}", encoding="utf-8") + self.assertFalse(model_sources_are_downloaded(models, "pixal3d/generate", sources)) + (encoder / "model.safetensors").write_bytes(b"x") + self.assertTrue(model_sources_are_downloaded(models, "pixal3d/generate", sources)) + + for child in sorted((models / "pixal3d").rglob("*"), reverse=True): + child.unlink() if child.is_file() else child.rmdir() + (models / "pixal3d").rmdir() + outside = Path(tmp) / "outside" + (outside / "generate").mkdir(parents=True) + try: + os.symlink(outside, models / "pixal3d", target_is_directory=True) + except (NotImplementedError, OSError) as exc: + self.skipTest(f"Symlinks unavailable: {exc}") + with self.assertRaisesRegex(ValueError, "symlink"): + resolve_model_root(models, "pixal3d/generate") + self.assertFalse(model_sources_are_downloaded(models, "pixal3d/generate", sources)) + + +if __name__ == "__main__": + unittest.main() diff --git a/electron/main/extension-install-utils.test.mjs b/electron/main/extension-install-utils.test.mjs index d5d1a389..84139f9a 100644 --- a/electron/main/extension-install-utils.test.mjs +++ b/electron/main/extension-install-utils.test.mjs @@ -53,6 +53,55 @@ test('validateInstallManifest still rejects missing process entry files', () => ) }) +test('validateInstallManifest accepts multi-source nodes and preserves legacy shapes', () => { + const mod = loadModule() + assert.doesNotThrow(() => mod.validateInstallManifest({ + id: 'multi-model', + generator_class: 'Generator', + nodes: [{ + id: 'generate', + model_sources: [ + { + id: 'primary', provider: 'huggingface', repo_id: 'org/main', + destination: '.', checks: ['pipeline.json'], + }, + { + id: 'encoder', provider: 'huggingface', repo_id: 'org/encoder', + destination: 'auxiliary/encoder', checks: ['model.safetensors'], + }, + ], + }], + }, { hasEntryFile: () => false, hasGeneratorFile: () => true }, 'repository')) + + assert.doesNotThrow(() => mod.validateInstallManifest({ + id: 'legacy', + generator_class: 'Generator', + nodes: [{ + id: 'projection', + hf_repo: 'org/legacy', + download_check: '../generate/model.safetensors', + hf_skip_prefixes: ['weights/**'], + }], + }, { hasEntryFile: () => false, hasGeneratorFile: () => true }, 'repository')) +}) + +test('validateInstallManifest rejects malformed or process model_sources', () => { + const mod = loadModule() + const source = { + id: 'weights', provider: 'huggingface', repo_id: 'org/model', + destination: '../outside', checks: ['model.safetensors'], + } + assert.throws(() => mod.validateInstallManifest({ + id: 'unsafe', generator_class: 'Generator', + nodes: [{ id: 'generate', model_sources: [source] }], + }, { hasEntryFile: () => false, hasGeneratorFile: () => true }, 'repository'), /destination/i) + + assert.throws(() => mod.validateInstallManifest({ + id: 'process', type: 'process', entry: 'processor.js', + nodes: [{ id: 'run', model_sources: [{ ...source, destination: '.' }] }], + }, { hasEntryFile: () => true, hasGeneratorFile: () => false }, 'repository'), /only for model nodes/i) +}) + test('python process setup failures are treated as fatal', () => { const mod = loadModule() diff --git a/electron/main/extension-install-utils.ts b/electron/main/extension-install-utils.ts index 4195cdca..05b965b0 100644 --- a/electron/main/extension-install-utils.ts +++ b/electron/main/extension-install-utils.ts @@ -1,9 +1,16 @@ +import { + normalizeModelSources, + safeModelSourceId, + type ModelSourceNode, +} from './model-sources' + export interface InstallManifest { id?: string type?: 'model' | 'process' entry?: string generator_class?: string - nodes?: Array<{ id?: string }> + model_sources?: unknown + nodes?: Array<{ id?: string; model_sources?: unknown } & ModelSourceNode> } export interface ValidatedInstallManifest { @@ -39,6 +46,18 @@ export function validateInstallManifest( const entryFile = manifest.entry ?? 'processor.js' const nodes = Array.isArray(manifest.nodes) ? manifest.nodes.filter((node) => node?.id) : [] + if (manifest.model_sources !== undefined) { + throw new Error('manifest.json: model_sources must be declared on a model node') + } + for (const node of Array.isArray(manifest.nodes) ? manifest.nodes : []) { + if (node.model_sources === undefined) continue + if (isProcess) { + throw new Error('manifest.json: model_sources is supported only for model nodes') + } + safeModelSourceId(node.id, 'model node id') + normalizeModelSources(node) + } + if (isProcess) { if (!opts.hasEntryFile(entryFile)) { throw new Error(`manifest.json: entry file "${entryFile}" missing from ${sourceLabel}`) diff --git a/electron/main/ipc-handlers.ts b/electron/main/ipc-handlers.ts index a0b2c136..a504d962 100644 --- a/electron/main/ipc-handlers.ts +++ b/electron/main/ipc-handlers.ts @@ -13,7 +13,15 @@ import { isModelDownloaded, listDownloadedModels, downloadModelFromHF, + downloadModelSourcesFromHF, } from './model-downloader' +import { resolveInstalledModelDownloadPlan } from './model-download-plan' +import { + areModelSourcesDownloaded, + modelHasLocalData, + normalizeModelSources, + resolveModelRoot, +} from './model-sources' import { getSettings, setSettings } from './settings-store' import { checkSetupNeeded, markSetupDone, runFullSetup, getVenvPythonExe, ensureSslPatch } from './python-setup' import { logger } from './logger' @@ -265,7 +273,12 @@ const renameWithRetry = (from: string, to: string, label: string) => renameExtensionWithRetry(from, to, label, logger) export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGetter): void { - const activeDownloads = new Map() + type ActiveDownload = { + progress: { percent: number; file?: string; fileIndex?: number; totalFiles?: number } + done: Promise + finish: () => void + } + const activeDownloads = new Map() // Logging from renderer ipcMain.on('log:error', (_event, message: string) => logger.error(`[Renderer] ${message}`)) ipcMain.handle('log:getPath', () => join(app.getPath('userData'), 'logs', 'modly.log')) @@ -423,7 +436,21 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe }) ipcMain.handle('model:delete', async (_, modelId: string): Promise<{ success: boolean; error?: string }> => { - const modelDir = join(getSettings(app.getPath('userData')).modelsDir, modelId) + if (activeDownloads.has(modelId)) { + return { success: false, error: 'Cannot remove model weights while their download is active' } + } + let modelDir: string + try { + await resolveInstalledModelDownloadPlan({ + modelId, + userExtensionsDir: getSettings(app.getPath('userData')).extensionsDir, + builtinExtensionsDir: getBuiltinExtensionsDir(), + blockedExtensionIds: activeExtensionInstalls, + }) + modelDir = resolveModelRoot(getSettings(app.getPath('userData')).modelsDir, modelId) + } catch (err) { + return { success: false, error: String(err) } + } // Unload the model and wait for confirmation so file handles are released try { @@ -475,28 +502,74 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe return listDownloadedModels(modelsDir) }) - ipcMain.handle('model:isDownloaded', (_, modelId: string, downloadCheck?: string): boolean => { + ipcMain.handle('model:isDownloaded', async (_, modelId: string): Promise => { const modelsDir = getSettings(app.getPath('userData')).modelsDir - return isModelDownloaded(modelsDir, modelId, downloadCheck) + try { + const plan = await resolveInstalledModelDownloadPlan({ + modelId, + userExtensionsDir: getSettings(app.getPath('userData')).extensionsDir, + builtinExtensionsDir: getBuiltinExtensionsDir(), + blockedExtensionIds: activeExtensionInstalls, + }) + return plan.kind === 'multi-source' + ? areModelSourcesDownloaded(modelsDir, modelId, plan.sources) + : isModelDownloaded(modelsDir, modelId, plan.downloadCheck) + } catch { + return false + } + }) + + ipcMain.handle('model:hasLocalData', async (_, modelId: string): Promise => { + try { + await resolveInstalledModelDownloadPlan({ + modelId, + userExtensionsDir: getSettings(app.getPath('userData')).extensionsDir, + builtinExtensionsDir: getBuiltinExtensionsDir(), + blockedExtensionIds: activeExtensionInstalls, + }) + return modelHasLocalData(getSettings(app.getPath('userData')).modelsDir, modelId) + } catch { + return false + } }) ipcMain.handle('model:activeDownloads', () => - [...activeDownloads.entries()].map(([modelId, progress]) => ({ modelId, ...progress })) + [...activeDownloads.entries()].map(([modelId, active]) => ({ modelId, ...active.progress })) ) ipcMain.handle('model:download', async ( event, - { repoId, modelId, skipPrefixes, includePrefixes }: { repoId: string; modelId: string; skipPrefixes?: string[]; includePrefixes?: string[] }, + modelId: string, ) => { if (activeDownloads.has(modelId)) { return { success: false, error: 'Download already in progress' } } - activeDownloads.set(modelId, { percent: 0 }) + let finish!: () => void + const done = new Promise((resolveDone) => { finish = resolveDone }) + const active: ActiveDownload = { progress: { percent: 0 }, done, finish } + activeDownloads.set(modelId, active) try { - await downloadModelFromHF(repoId, modelId, (progress) => { - activeDownloads.set(modelId, progress) + const plan = await resolveInstalledModelDownloadPlan({ + modelId, + userExtensionsDir: getSettings(app.getPath('userData')).extensionsDir, + builtinExtensionsDir: getBuiltinExtensionsDir(), + blockedExtensionIds: activeExtensionInstalls, + }) + const onProgress = (progress: typeof active.progress) => { + active.progress = progress event.sender.send('model:downloadProgress', { modelId, ...progress }) - }, skipPrefixes, includePrefixes) + } + if (plan.kind === 'multi-source') { + await downloadModelSourcesFromHF(modelId, plan.sources, onProgress) + } else { + await downloadModelFromHF( + plan.repoId, + modelId, + onProgress, + plan.skipPrefixes, + plan.includePrefixes, + ) + } return { success: true } } catch (err: any) { const message = err?.message ?? String(err) @@ -510,7 +583,8 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe } return { success: false, error: String(err) } } finally { - activeDownloads.delete(modelId) + if (activeDownloads.get(modelId) === active) activeDownloads.delete(modelId) + active.finish() } }) @@ -528,17 +602,24 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe ipcMain.handle('model:cancelDownload', async (_, modelId: string): Promise<{ success: boolean; error?: string }> => { try { + const active = activeDownloads.get(modelId) await axios.post(`${API_BASE_URL}/model/hf-download/cancel`, null, { params: { model_id: modelId }, timeout: 5000, }) - const modelDir = join(getSettings(app.getPath('userData')).modelsDir, modelId) + if (active) { + await Promise.race([ + active.done, + new Promise((_, reject) => { + setTimeout(() => reject(new Error('Timed out waiting for the download to stop')), 30_000) + }), + ]) + } + const modelDir = resolveModelRoot(getSettings(app.getPath('userData')).modelsDir, modelId) await rmAsync(modelDir, { recursive: true, force: true }) return { success: true } } catch (err) { return { success: false, error: String(err) } - } finally { - activeDownloads.delete(modelId) } }) @@ -835,6 +916,7 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe // extension type type?: 'model' | 'process' entry?: string + model_sources?: unknown // Optional top-level fallbacks — applied to each node if not set on the node params_schema?: unknown[] param_defaults?: Record @@ -851,6 +933,7 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe download_check?: string hf_skip_prefixes?: string[] hf_include_prefixes?: string[] + model_sources?: unknown }[] } @@ -866,20 +949,30 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe builtin, } - const nodes = (parsed.nodes ?? []).map(n => ({ - id: n.id, - name: n.name ?? n.id, - input: n.input ?? 'image' as const, - inputs: n.inputs, - inputLabels: n.input_labels, - output: n.output ?? 'mesh' as const, - paramsSchema: n.params_schema ?? parsed.params_schema ?? [], - paramDefaults: { ...(parsed.param_defaults ?? {}), ...(n.param_defaults ?? {}) }, - hfRepo: n.hf_repo, - downloadCheck: n.download_check, - hfSkipPrefixes: n.hf_skip_prefixes, - hfIncludePrefixes: n.hf_include_prefixes, - })) + if (parsed.model_sources !== undefined) { + throw new Error('manifest.json: model_sources must be declared on a model node') + } + const nodes = (parsed.nodes ?? []).map(n => { + if (parsed.type === 'process' && n.model_sources !== undefined) { + throw new Error('manifest.json: model_sources is supported only for model nodes') + } + const modelSources = normalizeModelSources(n) + return { + id: n.id, + name: n.name ?? n.id, + input: n.input ?? 'image' as const, + inputs: n.inputs, + inputLabels: n.input_labels, + output: n.output ?? 'mesh' as const, + paramsSchema: n.params_schema ?? parsed.params_schema ?? [], + paramDefaults: { ...(parsed.param_defaults ?? {}), ...(n.param_defaults ?? {}) }, + hfRepo: n.hf_repo, + downloadCheck: n.download_check, + hfSkipPrefixes: n.hf_skip_prefixes, + hfIncludePrefixes: n.hf_include_prefixes, + hasModelSources: modelSources !== undefined, + } + }) if (parsed.type === 'process') { return { ...common, type: 'process' as const, entry: parsed.entry ?? 'processor.js', nodes } @@ -1450,6 +1543,9 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe // Uninstall an extension — built-ins cannot be uninstalled ipcMain.handle('extensions:uninstall', async (_, extensionId: string) => { try { + if ([...activeDownloads.keys()].some((modelId) => modelId.split('/', 1)[0] === extensionId)) { + return { success: false, error: 'Cannot uninstall an extension while its model download is active' } + } // Corrupted folders can carry arbitrary names (manual copies, failed // unzips), so only enforce root confinement for the deletion path. The // strict id pattern still guards the built-in check — a non-conforming diff --git a/electron/main/model-download-plan.test.mjs b/electron/main/model-download-plan.test.mjs new file mode 100644 index 00000000..e5d037a0 --- /dev/null +++ b/electron/main/model-download-plan.test.mjs @@ -0,0 +1,92 @@ +import test from 'node:test' +import assert from 'node:assert/strict' +import { buildSync } from 'esbuild' +import { createRequire } from 'node:module' +import { mkdtempSync, mkdirSync, rmSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' + +function loadModule() { + const outfile = join(mkdtempSync(join(tmpdir(), 'modly-model-plan-module-')), 'model-plan.cjs') + const require = createRequire(import.meta.url) + const result = buildSync({ + entryPoints: [resolve('electron/main/model-download-plan.ts')], + bundle: true, + platform: 'node', + format: 'cjs', + write: false, + }) + writeFileSync(outfile, result.outputFiles[0].text, 'utf8') + return require(outfile) +} + +function setupExtension(manifest) { + const root = mkdtempSync(join(tmpdir(), 'modly-action-plan-')) + const user = join(root, 'user') + const builtin = join(root, 'builtin') + const extension = join(user, manifest.id) + mkdirSync(extension, { recursive: true }) + mkdirSync(builtin) + const manifestPath = join(extension, 'manifest.json') + writeFileSync(manifestPath, JSON.stringify(manifest)) + return { root, user, builtin, manifestPath } +} + +test('re-reads the installed manifest for each action and resolves only node-owned sources', async () => { + const { resolveInstalledModelDownloadPlan } = loadModule() + const manifest = { + id: 'pixal3d', + type: 'model', + nodes: [{ + id: 'generate', + model_sources: [{ + id: 'primary', provider: 'huggingface', repo_id: 'org/old', + destination: '.', checks: ['model.safetensors'], + }], + }], + } + const fixture = setupExtension(manifest) + const args = { + modelId: 'pixal3d/generate', + userExtensionsDir: fixture.user, + builtinExtensionsDir: fixture.builtin, + } + try { + const first = await resolveInstalledModelDownloadPlan(args) + assert.equal(first.kind, 'multi-source') + assert.equal(first.sources[0].repo_id, 'org/old') + + manifest.nodes[0].model_sources[0].repo_id = 'org/new' + writeFileSync(fixture.manifestPath, JSON.stringify(manifest)) + const second = await resolveInstalledModelDownloadPlan(args) + assert.equal(second.sources[0].repo_id, 'org/new') + } finally { + rmSync(fixture.root, { recursive: true, force: true }) + } +}) + +test('keeps legacy sibling checks and wildcard filters unchanged', async () => { + const { resolveInstalledModelDownloadPlan } = loadModule() + const fixture = setupExtension({ + id: 'triposplat', + type: 'model', + nodes: [{ + id: 'projection', + hf_repo: 'VAST-AI/TripoSplat', + download_check: '../generate/diffusion_models/triposplat_fp16.safetensors', + hf_skip_prefixes: ['weights/**', 'assets/*'], + }], + }) + try { + const plan = await resolveInstalledModelDownloadPlan({ + modelId: 'triposplat/projection', + userExtensionsDir: fixture.user, + builtinExtensionsDir: fixture.builtin, + }) + assert.equal(plan.kind, 'legacy') + assert.equal(plan.downloadCheck, '../generate/diffusion_models/triposplat_fp16.safetensors') + assert.deepEqual(plan.skipPrefixes, ['weights/**', 'assets/*']) + } finally { + rmSync(fixture.root, { recursive: true, force: true }) + } +}) diff --git a/electron/main/model-download-plan.ts b/electron/main/model-download-plan.ts new file mode 100644 index 00000000..27d2dba6 --- /dev/null +++ b/electron/main/model-download-plan.ts @@ -0,0 +1,139 @@ +import { existsSync } from 'node:fs' +import { readFile, readdir } from 'node:fs/promises' +import { join } from 'node:path' + +import { + EXT_INCOMPLETE_MARKER, + EXT_REGISTRATION_PENDING_MARKER, + assertSafeExtensionId, + resolveExtensionPathWithinRoot, +} from './extension-path-guard' +import { normalizeModelSources, safeModelSourceId, type ModelSource } from './model-sources' + +interface InstalledNode { + id?: unknown + hf_repo?: unknown + download_check?: unknown + hf_skip_prefixes?: unknown + hf_include_prefixes?: unknown + model_sources?: unknown +} + +interface InstalledManifest { + id?: unknown + type?: unknown + model_sources?: unknown + nodes?: unknown +} + +export type InstalledModelDownloadPlan = { + kind: 'legacy' + modelId: string + extensionId: string + nodeId: string + repoId: string + downloadCheck?: string + skipPrefixes?: string[] + includePrefixes?: string[] +} | { + kind: 'multi-source' + modelId: string + extensionId: string + nodeId: string + sources: ModelSource[] +} + +async function hasPendingRegistration(root: string, extensionId: string): Promise { + try { + const prefix = `${EXT_REGISTRATION_PENDING_MARKER}-${extensionId}-` + return (await readdir(root)).some((name) => ( + name.startsWith(prefix) && /^\d+$/.test(name.slice(prefix.length)) + )) + } catch { + return false + } +} + +function parseManifest(raw: string, extensionId: string, nodeId: string): InstalledModelDownloadPlan { + let parsed: unknown + try { + parsed = JSON.parse(raw) + } catch { + throw new Error(`Extension "${extensionId}" has an invalid manifest.json`) + } + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) { + throw new Error(`Extension "${extensionId}" has an invalid manifest.json`) + } + const manifest = parsed as InstalledManifest + if (manifest.id !== extensionId) throw new Error(`Installed manifest id does not match extension "${extensionId}"`) + if (manifest.type !== undefined && manifest.type !== 'model') { + throw new Error(`Extension "${extensionId}" is not a model extension`) + } + if (manifest.model_sources !== undefined) { + throw new Error('manifest.json: model_sources must be declared on a model node') + } + if (!Array.isArray(manifest.nodes)) throw new Error(`Extension "${extensionId}" does not declare model nodes`) + + const matches = manifest.nodes.filter((candidate): candidate is InstalledNode => ( + typeof candidate === 'object' + && candidate !== null + && !Array.isArray(candidate) + && (candidate as InstalledNode).id === nodeId + )) + if (matches.length !== 1) { + throw new Error(`Installed manifest must declare model node "${nodeId}" exactly once`) + } + + const node = matches[0] + const modelId = `${extensionId}/${nodeId}` + const sources = normalizeModelSources(node) + if (sources) return { kind: 'multi-source', modelId, extensionId, nodeId, sources } + + if (typeof node.hf_repo !== 'string' || !node.hf_repo) { + throw new Error(`Model node "${modelId}" has no Hugging Face download source`) + } + return { + kind: 'legacy', + modelId, + extensionId, + nodeId, + repoId: node.hf_repo, + downloadCheck: typeof node.download_check === 'string' ? node.download_check : undefined, + skipPrefixes: node.hf_skip_prefixes as string[] | undefined, + includePrefixes: node.hf_include_prefixes as string[] | undefined, + } +} + +/** Re-read the installed manifest for every model action; renderer metadata is never trusted. */ +export async function resolveInstalledModelDownloadPlan(args: { + modelId: unknown + userExtensionsDir: string + builtinExtensionsDir: string + blockedExtensionIds?: ReadonlySet +}): Promise { + if (typeof args.modelId !== 'string') throw new Error('Model id must be a string') + const parts = args.modelId.split('/') + if (parts.length !== 2) throw new Error('Model id must identify one extension node') + const extensionId = assertSafeExtensionId(parts[0]) + const nodeId = safeModelSourceId(parts[1], 'model node id') + if (args.blockedExtensionIds?.has(extensionId)) { + throw new Error(`Extension "${extensionId}" is being installed or repaired`) + } + + const userPath = resolveExtensionPathWithinRoot(args.userExtensionsDir, extensionId) + const builtinPath = resolveExtensionPathWithinRoot(args.builtinExtensionsDir, extensionId) + const extensionPath = existsSync(userPath) ? userPath : existsSync(builtinPath) ? builtinPath : undefined + if (!extensionPath) throw new Error(`Extension "${extensionId}" is not installed`) + const extensionRoot = extensionPath === userPath ? args.userExtensionsDir : args.builtinExtensionsDir + if ( + existsSync(join(extensionPath, EXT_INCOMPLETE_MARKER)) + || existsSync(join(extensionPath, EXT_REGISTRATION_PENDING_MARKER)) + || await hasPendingRegistration(extensionRoot, extensionId) + ) { + throw new Error(`Extension "${extensionId}" has an incomplete installation`) + } + + const manifestPath = join(extensionPath, 'manifest.json') + if (!existsSync(manifestPath)) throw new Error(`Extension "${extensionId}" has no manifest.json`) + return parseManifest(await readFile(manifestPath, 'utf-8'), extensionId, nodeId) +} diff --git a/electron/main/model-download-preload.test.mjs b/electron/main/model-download-preload.test.mjs new file mode 100644 index 00000000..8d9a16b7 --- /dev/null +++ b/electron/main/model-download-preload.test.mjs @@ -0,0 +1,55 @@ +import test from 'node:test' +import assert from 'node:assert/strict' +import { buildSync } from 'esbuild' +import { createRequire } from 'node:module' +import { mkdtempSync, readFileSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' + +function loadModule() { + const outfile = join(mkdtempSync(join(tmpdir(), 'modly-preload-download-')), 'electron-api.cjs') + const require = createRequire(import.meta.url) + const result = buildSync({ + entryPoints: [resolve('electron/preload/electron-api.ts')], + bundle: true, + platform: 'node', + format: 'cjs', + write: false, + }) + writeFileSync(outfile, result.outputFiles[0].text, 'utf8') + return require(outfile) +} + +test('renderer model actions send only the model node id', async () => { + const { createElectronApi } = loadModule() + const calls = [] + const ipc = { + invoke: async (...args) => { calls.push(args); return { success: true } }, + send: () => {}, + on: () => {}, + removeAllListeners: () => {}, + } + const api = createElectronApi(ipc, { setZoomFactor: () => {} }) + + await api.model.isDownloaded('pixal3d/generate') + await api.model.hasLocalData('pixal3d/generate') + await api.model.download('pixal3d/generate') + + assert.deepEqual(calls, [ + ['model:isDownloaded', 'pixal3d/generate'], + ['model:hasLocalData', 'pixal3d/generate'], + ['model:download', 'pixal3d/generate'], + ]) +}) + +test('declared partial data is removable and active downloads block destructive actions', () => { + const main = readFileSync(resolve('electron/main/ipc-handlers.ts'), 'utf8') + const page = readFileSync(resolve('src/areas/models/ModelsPage.tsx'), 'utf8') + const drawer = readFileSync(resolve('src/areas/models/components/ExtensionDrawer.tsx'), 'utf8') + + assert.match(main, /model:delete[\s\S]*activeDownloads\.has\(modelId\)/) + assert.match(main, /extensions:uninstall[\s\S]*activeDownloads\.keys\(\)/) + assert.match(page, /window\.electron\.model\.hasLocalData\(fullId\)/) + assert.match(drawer, /localDataIds\.includes\(fullId\) && state\.kind !== 'downloading'/) + assert.match(drawer, /Remove partial model data/) +}) diff --git a/electron/main/model-downloader.ts b/electron/main/model-downloader.ts index 8c5571d9..174b474a 100644 --- a/electron/main/model-downloader.ts +++ b/electron/main/model-downloader.ts @@ -6,6 +6,7 @@ import { existsSync, readdirSync, statSync, readFileSync } from 'fs' import { join } from 'path' import { getSettings } from './settings-store' import { app } from 'electron' +import type { ModelSource } from './model-sources' export interface DownloadProgress { percent: number @@ -121,7 +122,6 @@ export async function downloadModelFromHF( includePrefixes?: string[], ): Promise { const { net } = require('electron') - const STALL_TIMEOUT_MS = 120_000 let url = `${PYTHON_API_URL}/model/hf-download?repo_id=${encodeURIComponent(repoId)}&model_id=${encodeURIComponent(modelId)}` if (skipPrefixes && skipPrefixes.length > 0) { url += `&skip_prefixes=${encodeURIComponent(JSON.stringify(skipPrefixes))}` @@ -136,11 +136,39 @@ export async function downloadModelFromHF( const res = await net.fetch(url) if (!res.ok) throw new Error(`HuggingFace download failed: HTTP ${res.status}`) + await consumeDownloadStream(res, onProgress) +} + +/** Download a validated node-level source plan through one aggregate SSE stream. */ +export async function downloadModelSourcesFromHF( + modelId: string, + sources: ModelSource[], + onProgress: ProgressCallback, +): Promise { + const { net } = require('electron') + const headers: Record = { 'Content-Type': 'application/json' } + const hfToken = getSettings(app.getPath('userData')).hfToken + if (hfToken) headers.Authorization = `Bearer ${hfToken}` + const url = `${PYTHON_API_URL}/model/hf-download-sources?model_id=${encodeURIComponent(modelId)}` + const res = await net.fetch(url, { + method: 'POST', + headers, + body: JSON.stringify({ sources }), + }) + if (!res.ok) throw new Error(`HuggingFace multi-source download failed: HTTP ${res.status}`) + await consumeDownloadStream(res, onProgress) +} + +async function consumeDownloadStream( + res: Response, + onProgress: ProgressCallback, +): Promise { if (!res.body) throw new Error('No response body from HF download stream') const decoder = new TextDecoder() const reader = res.body.getReader() let buffer = '' + const STALL_TIMEOUT_MS = 120_000 async function readWithTimeout() { return await Promise.race([ diff --git a/electron/main/model-sources.test.mjs b/electron/main/model-sources.test.mjs new file mode 100644 index 00000000..b29f0dc5 --- /dev/null +++ b/electron/main/model-sources.test.mjs @@ -0,0 +1,109 @@ +import test from 'node:test' +import assert from 'node:assert/strict' +import { buildSync } from 'esbuild' +import { createRequire } from 'node:module' +import { mkdtempSync, mkdirSync, rmSync, symlinkSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' + +function loadModule() { + const outfile = join(mkdtempSync(join(tmpdir(), 'modly-model-sources-module-')), 'model-sources.cjs') + const require = createRequire(import.meta.url) + const result = buildSync({ + entryPoints: [resolve('electron/main/model-sources.ts')], + bundle: true, + platform: 'node', + format: 'cjs', + write: false, + }) + writeFileSync(outfile, result.outputFiles[0].text, 'utf8') + return require(outfile) +} + +const validNode = () => ({ + model_sources: [ + { + id: 'primary', + provider: 'huggingface', + repo_id: 'org/main', + destination: '.', + checks: ['pipeline.json'], + }, + { + id: 'encoder', + provider: 'huggingface', + repo_id: 'org/encoder', + revision: 'refs/pr/1', + destination: 'auxiliary/encoder', + include_prefixes: ['config.json', 'weights/'], + skip_prefixes: ['README.md'], + checks: ['config.json', 'model.safetensors'], + }, + ], +}) + +test('validates the new model_sources contract without reinterpreting legacy fields', () => { + const { normalizeModelSources } = loadModule() + const sources = normalizeModelSources(validNode()) + assert.equal(sources.length, 2) + assert.equal(sources[1].destination, 'auxiliary/encoder') + + assert.equal(normalizeModelSources({ + hf_repo: 'legacy/repo', + download_check: '../generate/model.safetensors', + hf_skip_prefixes: ['weights/**'], + }), undefined) +}) + +test('rejects unsafe destinations, unsupported providers, and non-portable source aliases', () => { + const { normalizeModelSources } = loadModule() + const source = validNode().model_sources[0] + for (const destination of ['../outside', 'aux/CON', 'aux/name.', 'C:/models']) { + assert.throws( + () => normalizeModelSources({ model_sources: [{ ...source, destination }] }), + /destination|unsafe/i, + ) + } + assert.throws( + () => normalizeModelSources({ model_sources: [{ ...source, provider: 'url' }] }), + /provider.*huggingface/i, + ) + assert.throws( + () => normalizeModelSources({ model_sources: [source, { ...source, id: 'PRIMARY' }] }), + /portable-unique/i, + ) + assert.throws( + () => normalizeModelSources({ model_sources: [{ ...source, checks: [] }] }), + /checks.*non-empty/i, + ) +}) + +test('requires every declared check and rejects symlinked extension-root ancestry', (t) => { + const { areModelSourcesDownloaded, normalizeModelSources } = loadModule() + const root = mkdtempSync(join(tmpdir(), 'modly-model-readiness-')) + const models = join(root, 'models') + const modelRoot = join(models, 'pixal3d', 'generate') + const sources = normalizeModelSources(validNode()) + mkdirSync(join(modelRoot, 'auxiliary', 'encoder'), { recursive: true }) + writeFileSync(join(modelRoot, 'pipeline.json'), '{}') + writeFileSync(join(modelRoot, 'auxiliary', 'encoder', 'config.json'), '{}') + + try { + assert.equal(areModelSourcesDownloaded(models, 'pixal3d/generate', sources), false) + writeFileSync(join(modelRoot, 'auxiliary', 'encoder', 'model.safetensors'), 'x') + assert.equal(areModelSourcesDownloaded(models, 'pixal3d/generate', sources), true) + + rmSync(join(models, 'pixal3d'), { recursive: true, force: true }) + const outside = join(root, 'outside') + mkdirSync(join(outside, 'generate'), { recursive: true }) + try { + symlinkSync(outside, join(models, 'pixal3d'), 'dir') + } catch (error) { + t.skip(`Symlinks unavailable: ${error}`) + return + } + assert.equal(areModelSourcesDownloaded(models, 'pixal3d/generate', sources), false) + } finally { + rmSync(root, { recursive: true, force: true }) + } +}) diff --git a/electron/main/model-sources.ts b/electron/main/model-sources.ts new file mode 100644 index 00000000..03b82948 --- /dev/null +++ b/electron/main/model-sources.ts @@ -0,0 +1,190 @@ +import { existsSync, lstatSync, readdirSync } from 'node:fs' +import { isAbsolute, relative, resolve } from 'node:path' + +export interface ModelSource { + id: string + provider: 'huggingface' + repo_id: string + revision?: string + destination: string + include_prefixes?: string[] + skip_prefixes?: string[] + checks: string[] +} + +export interface ModelSourceNode { + model_sources?: unknown +} + +const SAFE_ID = /^[A-Za-z0-9][A-Za-z0-9._-]*$/ +const WINDOWS_DEVICE = /^(?:con|prn|aux|nul|com[1-9]|lpt[1-9])(?:\..*)?$/i +const WINDOWS_UNSAFE = /[<>"|?*\u0000-\u001f]/ + +function portableSegment(value: string, field: string): string { + if ( + !value + || value === '.' + || value === '..' + || value.endsWith('.') + || value.endsWith(' ') + || value.includes(':') + || WINDOWS_UNSAFE.test(value) + || WINDOWS_DEVICE.test(value) + ) { + throw new Error(`${field} contains unsafe path segment "${value}"`) + } + return value +} + +export function safeModelSourceId(value: unknown, field = 'model source id'): string { + if (typeof value !== 'string' || !value || value !== value.trim() || !SAFE_ID.test(value)) { + throw new Error(`${field} must be a safe non-empty identifier`) + } + return portableSegment(value, field) +} + +export function safeModelRelativePath(value: unknown, field: string, allowDot = false): string { + if (typeof value !== 'string' || !value || value !== value.trim()) { + throw new Error(`${field} must be a non-empty relative path`) + } + if (allowDot && value === '.') return value + if (value === '.' || value.startsWith('/') || value.includes('\\') || isAbsolute(value)) { + throw new Error(`${field} must be a safe relative POSIX path`) + } + for (const part of value.split('/')) portableSegment(part, field) + return value +} + +function safePrefix(value: unknown, field: string): string { + if (typeof value !== 'string') return safeModelRelativePath(value, field) + const path = value.endsWith('/') ? value.slice(0, -1) : value + safeModelRelativePath(path, field) + return value +} + +function optionalPrefixes(value: unknown, field: string): string[] | undefined { + if (value === undefined) return undefined + if (!Array.isArray(value)) throw new Error(`${field} must be an array`) + return value.map((entry, index) => safePrefix(entry, `${field}[${index}]`)) +} + +function safeRepoId(value: unknown, field: string): string { + if (typeof value !== 'string' || !value || value !== value.trim() || value.includes('\\')) { + throw new Error(`${field} must be a non-empty Hugging Face repository id`) + } + const parts = value.split('/') + if (parts.length > 2 || parts.some((part) => !SAFE_ID.test(part) || part === '.' || part === '..')) { + throw new Error(`${field} is not a safe Hugging Face repository id`) + } + return value +} + +function safeRevision(value: unknown, field: string): string | undefined { + if (value === undefined) return undefined + if (typeof value !== 'string' || !value || value !== value.trim() || value.startsWith('/') || value.includes('\\') || value.includes('\0')) { + throw new Error(`${field} must be a safe non-empty revision`) + } + if (value.split('/').some((part) => !part || part === '.' || part === '..')) { + throw new Error(`${field} must be a safe revision`) + } + return value +} + +export function normalizeModelSources(node: ModelSourceNode): ModelSource[] | undefined { + if (!Object.prototype.hasOwnProperty.call(node, 'model_sources')) return undefined + if (!Array.isArray(node.model_sources) || node.model_sources.length === 0) { + throw new Error('model_sources must be a non-empty array') + } + + const seen = new Map() + return node.model_sources.map((raw, index) => { + const field = `model_sources[${index}]` + if (typeof raw !== 'object' || raw === null || Array.isArray(raw)) { + throw new Error(`${field} must be an object`) + } + const value = raw as Record + const id = safeModelSourceId(value.id, `${field}.id`) + const alias = id.normalize('NFC').toLowerCase() + const previous = seen.get(alias) + if (previous) throw new Error(`model source ids "${previous}" and "${id}" are not portable-unique`) + seen.set(alias, id) + if (value.provider !== 'huggingface') throw new Error(`${field}.provider must be "huggingface"`) + if (!Array.isArray(value.checks) || value.checks.length === 0) { + throw new Error(`${field}.checks must be a non-empty array`) + } + + const source: ModelSource = { + id, + provider: 'huggingface', + repo_id: safeRepoId(value.repo_id, `${field}.repo_id`), + destination: safeModelRelativePath(value.destination, `${field}.destination`, true), + checks: value.checks.map((check, checkIndex) => ( + safeModelRelativePath(check, `${field}.checks[${checkIndex}]`) + )), + } + const revision = safeRevision(value.revision, `${field}.revision`) + const include = optionalPrefixes(value.include_prefixes, `${field}.include_prefixes`) + const skip = optionalPrefixes(value.skip_prefixes, `${field}.skip_prefixes`) + if (revision !== undefined) source.revision = revision + if (include !== undefined) source.include_prefixes = include + if (skip !== undefined) source.skip_prefixes = skip + return source + }) +} + +function pathHasSymlink(root: string, candidate: string): boolean { + const rootPath = resolve(root) + const rel = relative(rootPath, resolve(candidate)) + if (rel === '..' || rel.startsWith('../') || rel.startsWith('..\\') || isAbsolute(rel)) return true + let current = rootPath + try { + if (existsSync(current) && lstatSync(current).isSymbolicLink()) return true + for (const part of rel.split(/[/\\]/).filter(Boolean)) { + current = resolve(current, part) + if (existsSync(current) && lstatSync(current).isSymbolicLink()) return true + } + } catch { + return true + } + return false +} + +export function resolveModelRoot(modelsDir: string, modelId: string): string { + if (typeof modelId !== 'string') throw new Error('Model id must be a string') + const parts = modelId.split('/') + if (parts.length !== 2) throw new Error('Model id must identify one extension node') + const extensionId = safeModelSourceId(parts[0], 'extension id') + const nodeId = safeModelSourceId(parts[1], 'model node id') + const root = resolve(modelsDir) + const modelRoot = resolve(root, extensionId, nodeId) + if (pathHasSymlink(root, modelRoot)) throw new Error('Model path resolves through a symlink') + return modelRoot +} + +export function areModelSourcesDownloaded(modelsDir: string, modelId: string, sources: ModelSource[]): boolean { + try { + const modelRoot = resolveModelRoot(modelsDir, modelId) + if (!existsSync(modelRoot)) return false + return sources.every((source) => { + const destination = source.destination === '.' + ? modelRoot + : resolve(modelRoot, ...source.destination.split('/')) + if (!existsSync(destination) || pathHasSymlink(modelRoot, destination)) return false + return source.checks.every((check) => { + const candidate = resolve(destination, ...check.split('/')) + return existsSync(candidate) && !pathHasSymlink(modelRoot, candidate) + }) + }) + } catch { + return false + } +} + +export function modelHasLocalData(modelsDir: string, modelId: string): boolean { + try { + const modelRoot = resolveModelRoot(modelsDir, modelId) + return existsSync(modelRoot) && readdirSync(modelRoot).length > 0 + } catch { + return false + } +} diff --git a/electron/preload/electron-api.ts b/electron/preload/electron-api.ts index af211f0d..c929204c 100644 --- a/electron/preload/electron-api.ts +++ b/electron/preload/electron-api.ts @@ -111,9 +111,9 @@ export function createElectronApi(ipcRenderer: IpcRendererLike, webFrame: WebFra model: { export: (args: { outputUrl: string; format: string }) => ipcRenderer.invoke('model:export', args), listDownloaded: () => ipcRenderer.invoke('model:listDownloaded'), - isDownloaded: (modelId: string, downloadCheck?: string) => ipcRenderer.invoke('model:isDownloaded', modelId, downloadCheck), - download: (repoId: string, modelId: string, skipPrefixes?: string[], includePrefixes?: string[]) => - ipcRenderer.invoke('model:download', { repoId, modelId, skipPrefixes, includePrefixes }), + isDownloaded: (modelId: string) => ipcRenderer.invoke('model:isDownloaded', modelId), + hasLocalData: (modelId: string) => ipcRenderer.invoke('model:hasLocalData', modelId), + download: (modelId: string) => ipcRenderer.invoke('model:download', modelId), pauseDownload: (modelId: string) => ipcRenderer.invoke('model:pauseDownload', modelId), cancelDownload: (modelId: string) => ipcRenderer.invoke('model:cancelDownload', modelId), delete: (modelId: string) => ipcRenderer.invoke('model:delete', modelId), diff --git a/src/areas/models/ModelsPage.tsx b/src/areas/models/ModelsPage.tsx index 895d4e8f..66f8e5d4 100644 --- a/src/areas/models/ModelsPage.tsx +++ b/src/areas/models/ModelsPage.tsx @@ -2,11 +2,11 @@ import { useEffect, useMemo, useRef, useState } from 'react' import { createPortal } from 'react-dom' import { useExtensionsStore } from '@shared/stores/extensionsStore' import type { AnyExtension, ModelExtension } from '@shared/types/electron.d' -import { formatModelName } from './utils' +import { deleteModelsThenUninstallExtension, formatModelName } from './utils' import { ExtensionCard } from './components/ExtensionCard' import type { ExtensionNode } from './components/ExtensionCard' import { ExtensionDrawer } from './components/ExtensionDrawer' -import { ICONS } from './components/extensionShared' +import { ICONS, nodeHasManagedWeights } from './components/extensionShared' // ─── Filters & sorts ────────────────────────────────────────────────────────── @@ -49,6 +49,7 @@ export default function ModelsPage(): JSX.Element { // Model weight state (needed for node install status + uninstall cleanup) const [installedVariantIds, setInstalledVariantIds] = useState([]) + const [localDataIds, setLocalDataIds] = useState([]) const [downloading, setDownloading] = useState { @@ -161,11 +168,11 @@ export default function ModelsPage(): JSX.Element { // ── Node install / download controls ────────────────────────────────────── function handleInstallNode(node: ExtensionNode, fullId: string) { - if (!node.hfRepo) return + if (!nodeHasManagedWeights(node)) return setDownloading((prev) => ({ ...prev, [fullId]: { ...(prev[fullId] ?? { percent: 0 }), paused: false, status: 'Starting…' } })) - window.electron.model.download(node.hfRepo!, fullId, node.hfSkipPrefixes, node.hfIncludePrefixes).then((result: { success: boolean; paused?: boolean; cancelled?: boolean }) => { + window.electron.model.download(fullId).then((result) => { if (!result.success && !result.paused && !result.cancelled) { - setGhErr('Download failed') + setGhErr(result.error ?? 'Download failed') setDownloading((prev) => { const n = { ...prev }; delete n[fullId]; return n }) } }) @@ -174,7 +181,7 @@ export default function ModelsPage(): JSX.Element { function handleInstallAll(ext: AnyExtension) { if (ext.type !== 'model') return for (const node of ext.nodes) { - if (!node.hfRepo) continue + if (!nodeHasManagedWeights(node)) continue const fullId = `${ext.id}/${node.id}` if (installedVariantIds.includes(fullId) || downloading[fullId]) continue handleInstallNode(node, fullId) @@ -188,7 +195,9 @@ export default function ModelsPage(): JSX.Element { async function handleCancelDownload(fullId: string) { setDownloading((prev) => { const n = { ...prev }; delete n[fullId]; return n }) - await window.electron.model.cancelDownload(fullId) + const result = await window.electron.model.cancelDownload(fullId) + if (!result.success) setGhErr(result.error ?? 'Could not cancel download') + await refreshInstalledIds(useExtensionsStore.getState().modelExtensions) } async function handleUninstallNode(fullId: string) { @@ -228,8 +237,8 @@ export default function ModelsPage(): JSX.Element { function openUninstallModal(extId: string) { const ext = allExtensions.find((e) => e.id === extId) if (ext?.type === 'model') { - const installedModels = ext.nodes.filter((n) => installedVariantIds.includes(`${extId}/${n.id}`)) - setModelsToDelete(new Set(installedModels.map((n) => `${extId}/${n.id}`))) + const localModels = ext.nodes.filter((n) => localDataIds.includes(`${extId}/${n.id}`)) + setModelsToDelete(new Set(localModels.map((n) => `${extId}/${n.id}`))) } else { setModelsToDelete(new Set()) } @@ -237,10 +246,12 @@ export default function ModelsPage(): JSX.Element { } async function handleUninstallExtension(extId: string) { - for (const modelId of modelsToDelete) { - await window.electron.model.delete(modelId) - } - const result = await uninstallExt(extId) + const result = await deleteModelsThenUninstallExtension( + extId, + modelsToDelete, + (modelId) => window.electron.model.delete(modelId), + uninstallExt, + ) if (!result.success) { // Keep the dialog open so the failure is visible (locked folder, etc.) setUninstallError(result.error ?? 'Could not delete the extension folder.') @@ -617,6 +628,7 @@ export default function ModelsPage(): JSX.Element { { const ext = allExtensions.find((e) => e.id === uninstallTarget) const installedModels = ext?.type === 'model' - ? ext.nodes.filter((n) => installedVariantIds.includes(`${uninstallTarget}/${n.id}`)) + ? ext.nodes.filter((n) => localDataIds.includes(`${uninstallTarget}/${n.id}`)) : [] return createPortal( diff --git a/src/areas/models/components/ExtensionDrawer.tsx b/src/areas/models/components/ExtensionDrawer.tsx index 54adab9c..2f154c71 100644 --- a/src/areas/models/components/ExtensionDrawer.tsx +++ b/src/areas/models/components/ExtensionDrawer.tsx @@ -16,6 +16,7 @@ import { finishExtensionRepair, isExtensionRepairable } from '../utils' interface Props { ext: AnyExtension installedIds: string[] + localDataIds: string[] downloading: DownloadMap loadError?: string disabled?: boolean @@ -31,7 +32,7 @@ interface Props { } export function ExtensionDrawer({ - ext, installedIds, downloading, loadError, disabled, + ext, installedIds, localDataIds, downloading, loadError, disabled, onInstall, onInstallAll, onPauseDownload, onCancelDownload, onUninstallNode, onUninstall, onRepaired, onSynced, onClose, }: Props): JSX.Element { @@ -203,11 +204,11 @@ export function ExtensionDrawer({ onResume={() => onInstall(node, fullId)} onCancel={() => onCancelDownload(fullId)} /> - {state.kind === 'installed' && ( + {localDataIds.includes(fullId) && state.kind !== 'downloading' && (