diff --git a/README.md b/README.md index b162cf23..c7a9d7b7 100644 --- a/README.md +++ b/README.md @@ -148,6 +148,69 @@ supported provider is `huggingface`. Existing nodes that use `hf_repo`, `download_check`, `hf_include_prefixes`, and `hf_skip_prefixes` keep their original behavior. +### Separately installable weight variants + +A model node that publishes the same weights in several variants (quantizations, +precisions…) can declare `weight_variants` next to `hf_repo`. The Extensions page +lists every variant under the node, and each one is downloaded or deleted on its own. + +```json +{ + "id": "generate", + "hf_repo": "org/model-gguf", + "download_check": "pipeline.json", + "hf_include_prefixes": ["pipeline.json", "encoder/", "dit/"], + "params_schema": [ + { + "id": "quant", + "label": "Quantization", + "type": "select", + "default": "Q5_K_M", + "options": [ + { "value": "Q4_K_M", "label": "Q4_K_M" }, + { "value": "Q5_K_M", "label": "Q5_K_M" } + ] + } + ], + "weight_variants": { + "param": "quant", + "default": "Q5_K_M", + "options": [ + { + "id": "Q4_K_M", + "label": "Q4_K_M", + "size_gb": 2.4, + "vram_gb": 6, + "include_prefixes": ["dit/model_Q4_K_M.gguf"], + "checks": ["dit/model_Q4_K_M.gguf"] + }, + { + "id": "Q5_K_M", + "include_prefixes": ["dit/model_Q5_K_M.gguf"], + "checks": ["dit/model_Q5_K_M.gguf"] + } + ] + } +} +``` + +- `param` names the `params_schema` select whose values are the variant ids. That + param must exist on the node (or on the extension, as its fallback), and when it + declares `options` they must cover every variant id. +- `size_gb` (download size) and `vram_gb` (approximate VRAM the variant needs) are + optional positive numbers, shown next to the variant when present. +- Every install downloads the shared files (`hf_include_prefixes`, with every + variant's files excluded automatically) plus one variant: the one asked for, or the + `default` one — the first option when `default` is omitted. Files already complete + on disk are skipped, so adding a second variant only fetches that variant. +- A variant is installed when all of its `checks` exist; the node is installed once + its `download_check` and at least one variant are present. +- Generation fails with an explicit message when the selected variant is not + installed, and the node's selector labels those options `(not installed)`. +- `include_prefixes` and `checks` are safe POSIX paths relative to the node's model + directory. Prefixes of two variants cannot overlap, `download_check` stays outside + every variant, and `weight_variants` cannot be combined with `model_sources`. + --- ## Workflows diff --git a/api/routers/generation.py b/api/routers/generation.py index 8481deb4..b4566036 100644 --- a/api/routers/generation.py +++ b/api/routers/generation.py @@ -185,6 +185,7 @@ def progress_cb(pct: int, step: str = "") -> None: try: loop = asyncio.get_running_loop() + generator_registry.assert_weight_variant_installed(params) # Check if the model needs to be loaded BEFORE calling get_active(), # because get_active() loads the model in a blocking manner. diff --git a/api/services/generator_registry.py b/api/services/generator_registry.py index 348a42cb..453a98d1 100644 --- a/api/services/generator_registry.py +++ b/api/services/generator_registry.py @@ -25,7 +25,12 @@ 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 +from services.model_sources import ( + missing_weight_variant, + model_sources_are_downloaded, + normalize_model_sources, + normalize_weight_variants, +) # ------------------------------------------------------------------ # # Global paths @@ -528,6 +533,9 @@ def _discover_extensions( if nodes: for node in nodes: model_sources = normalize_model_sources(node) + weight_variants = normalize_weight_variants( + node, node.get("params_schema", manifest.get("params_schema", [])) + ) node_manifest = { **manifest, "id": f"{ext_id}/{node['id']}", @@ -544,6 +552,8 @@ def _discover_extensions( } if model_sources is not None: node_manifest["model_sources"] = model_sources + if weight_variants is not None: + node_manifest["weight_variants"] = weight_variants full_id = f"{ext_id}/{node['id']}" result[full_id] = (cls_or_None, node_manifest, ext_dir, legacy_context) if subprocess_mode: @@ -724,6 +734,18 @@ def get_active(self) -> BaseGenerator: gen.load() return gen + def assert_weight_variant_installed(self, params: dict) -> None: + """Refuse generation when the weight variant selected by params is not installed.""" + manifest = self._manifests.get(self._active_id, {}) + option = missing_weight_variant( + MODELS_DIR, self._active_id, manifest.get("weight_variants"), params + ) + if option is not None: + raise RuntimeError( + f'{option["label"]} weights for {self._active_id} are not installed. ' + "Install them from the Extensions page, or select an installed variant." + ) + def get_generator(self, model_id: str) -> BaseGenerator: self._assert_not_quarantined(model_id) if model_id not in self._generators: diff --git a/api/services/model_sources.py b/api/services/model_sources.py index 592d1342..1dee2711 100644 --- a/api/services/model_sources.py +++ b/api/services/model_sources.py @@ -2,6 +2,7 @@ from __future__ import annotations +import math import re import unicodedata from pathlib import Path @@ -270,3 +271,171 @@ def validate_source_file_plan( f'"{previous_source}:{previous_target}" and "{source_id}:{value}"' ) aliases[alias] = (source_id, value) + + +def _prefixes_overlap(left: list[str], right: list[str]) -> bool: + return any( + a.lower().startswith(b.lower()) or b.lower().startswith(a.lower()) + for a in left + for b in right + ) + + +def _assert_param_offers_variants(param: str, params_schema: Any, ids: list[str]) -> None: + """Variant ids must be selectable, so the param they key has to exist and offer them.""" + entry = next( + (p for p in params_schema if isinstance(p, dict) and p.get("id") == param), + None, + ) if isinstance(params_schema, list) else None + if entry is None: + raise ValueError(f'weight_variants.param must name a params_schema entry ("{param}")') + options = entry.get("options") + if not isinstance(options, list): + return + values = { + str(option.get("value")) if isinstance(option, dict) else str(option) + for option in options + } + missing = [variant_id for variant_id in ids if variant_id not in values] + if missing: + raise ValueError( + f'the "{param}" param must offer every weight variant id ' + f'(missing: {", ".join(missing)})' + ) + + +def normalize_weight_variants( + node: dict[str, Any], params_schema: Any = None +) -> dict[str, Any] | None: + """Validate a node's separately installable weight variants (e.g. quantizations).""" + if "weight_variants" not in node: + return None + if "model_sources" in node: + raise ValueError("weight_variants cannot be combined with model_sources") + if not isinstance(node.get("hf_repo"), str) or not node["hf_repo"]: + raise ValueError("weight_variants requires hf_repo on the same node") + raw = node["weight_variants"] + if not isinstance(raw, dict): + raise ValueError("weight_variants must be an object") + param = safe_source_id(raw.get("param"), "weight_variants.param") + raw_options = raw.get("options") + if not isinstance(raw_options, list) or not raw_options: + raise ValueError("weight_variants.options must be a non-empty array") + + aliases: dict[str, str] = {} + options: list[dict[str, Any]] = [] + for index, raw_option in enumerate(raw_options): + field = f"weight_variants.options[{index}]" + if not isinstance(raw_option, dict): + raise ValueError(f"{field} must be an object") + variant_id = safe_source_id(raw_option.get("id"), f"{field}.id") + alias = unicodedata.normalize("NFC", variant_id).casefold() + if alias in aliases: + raise ValueError( + f'weight variant ids "{aliases[alias]}" and "{variant_id}" are not portable-unique' + ) + aliases[alias] = variant_id + label = raw_option.get("label", variant_id) + if not isinstance(label, str) or not label.strip(): + raise ValueError(f"{field}.label must be a non-empty string") + size_gb = raw_option.get("size_gb") + if size_gb is not None and ( + isinstance(size_gb, bool) + or not isinstance(size_gb, (int, float)) + or not math.isfinite(size_gb) + or size_gb <= 0 + ): + raise ValueError(f"{field}.size_gb must be a positive number") + vram_gb = raw_option.get("vram_gb") + if vram_gb is not None and ( + isinstance(vram_gb, bool) + or not isinstance(vram_gb, (int, float)) + or not math.isfinite(vram_gb) + or vram_gb <= 0 + ): + raise ValueError(f"{field}.vram_gb must be a positive number") + include = _prefixes(raw_option.get("include_prefixes"), f"{field}.include_prefixes") + if not include: + raise ValueError(f"{field}.include_prefixes must be a non-empty array") + checks = raw_option.get("checks") + if not isinstance(checks, list) or not checks: + raise ValueError(f"{field}.checks must be a non-empty array") + safe_checks: list[str] = [] + for check_index, check in enumerate(checks): + path = safe_relative_path(check, f"{field}.checks[{check_index}]") + if not any(path.startswith(prefix) for prefix in include): + raise ValueError( + f"{field}.checks[{check_index}] is not covered by its include_prefixes" + ) + safe_checks.append(path) + option: dict[str, Any] = { + "id": variant_id, + "label": label, + "include_prefixes": include, + "checks": safe_checks, + } + if size_gb is not None: + option["size_gb"] = size_gb + if vram_gb is not None: + option["vram_gb"] = vram_gb + options.append(option) + + for index, option in enumerate(options): + for other in options[index + 1:]: + if _prefixes_overlap(option["include_prefixes"], other["include_prefixes"]): + raise ValueError( + f'weight variants "{option["id"]}" and "{other["id"]}" share files' + ) + download_check = node.get("download_check") + if isinstance(download_check, str) and any( + _prefixes_overlap([download_check], option["include_prefixes"]) for option in options + ): + raise ValueError("download_check must name a file outside every weight variant") + default = raw.get("default", options[0]["id"]) + if not any(option["id"] == default for option in options): + raise ValueError("weight_variants.default must name one of its options") + _assert_param_offers_variants( + param, + node.get("params_schema") if params_schema is None else params_schema, + [option["id"] for option in options], + ) + return {"param": param, "default": default, "options": options} + + +def installed_weight_variants( + models_dir: Path, model_id: str, variants: dict[str, Any] +) -> list[str]: + try: + model_root = resolve_model_root(models_dir, model_id) + except ValueError: + return [] + + def _present(check: str) -> bool: + try: + candidate = resolve_download_path(model_root, check) + return candidate.is_file() and candidate.stat().st_size > 0 + except (OSError, ValueError): + return False + + return [ + option["id"] + for option in variants["options"] + if all(_present(check) for check in option["checks"]) + ] + + +def missing_weight_variant( + models_dir: Path, + model_id: str, + variants: dict[str, Any] | None, + params: dict[str, Any], +) -> dict[str, Any] | None: + """The variant selected by params when its files are not installed, else None.""" + if not variants: + return None + selected = params.get(variants["param"]) + selected_id = variants["default"] if selected is None else str(selected) + option = next((o for o in variants["options"] if o["id"] == selected_id), None) + if option is None or option["id"] in installed_weight_variants(models_dir, model_id, variants): + return None + return option diff --git a/api/tests/test_generation_router.py b/api/tests/test_generation_router.py index 20fdda94..18a57253 100644 --- a/api/tests/test_generation_router.py +++ b/api/tests/test_generation_router.py @@ -42,6 +42,9 @@ def active_status(self) -> dict: # Report loaded so _run_generation skips the download/load thread. return {"loaded": True, "name": "fake", "downloaded": True} + def assert_weight_variant_installed(self, params: dict) -> None: + pass + def get_active(self) -> _FakeGenerator: return self._gen @@ -101,6 +104,23 @@ def test_output_lands_under_the_current_workspace(self) -> None: self.assertEqual(job.status, "done") self.assertEqual(job.output_url, "/workspace/MyColl/model.glb") + def test_missing_weight_variant_fails_the_job_before_generation(self) -> None: + class _MissingVariantRegistry(_FakeRegistry): + def assert_weight_variant_installed(self, params: dict) -> None: + raise RuntimeError(f'{params["gguf_quant"]} weights for trellis2/generate are not installed.') + + gen = _FakeGenerator() + generation.generator_registry = _MissingVariantRegistry(gen) + job_id = "job-missing-variant" + generation._jobs[job_id] = JobStatus(job_id=job_id, status="pending", progress=0) + generation._cancel_events[job_id] = threading.Event() + asyncio.run(generation._run_generation(job_id, b"img", {"gguf_quant": "Q6_K"}, "MyColl")) + + job = generation._jobs[job_id] + self.assertEqual(job.status, "error") + self.assertIn("Q6_K weights for trellis2/generate are not installed", job.error) + self.assertIsNone(gen.outputs_dir) + class GenerateFromImageWorkspaceTests(unittest.TestCase): """The request path must survive the relocation too, not just the worker: diff --git a/api/tests/test_generator_registry.py b/api/tests/test_generator_registry.py index ff9d090c..86736b39 100644 --- a/api/tests/test_generator_registry.py +++ b/api/tests/test_generator_registry.py @@ -211,6 +211,61 @@ def test_declared_sources_block_generation_even_when_generator_overrides_readine with self.assertRaisesRegex(RuntimeError, "Model sources are incomplete"): self.registry.get_active() + def test_selected_weight_variant_must_be_installed_before_generation(self) -> None: + def variant(quant: str) -> dict: + return { + "id": quant, + "include_prefixes": [f"dit_{quant}.gguf"], + "checks": [f"dit_{quant}.gguf"], + } + + extension = self._make_extension("quantized") + manifest = { + "id": "quantized", + "name": "quantized", + "type": "model", + "generator_class": "TestGenerator", + "params_schema": [ + {"id": "quant", "type": "select", "options": [{"value": "Q4"}, {"value": "Q5"}]} + ], + "nodes": [{ + "id": "generate", + "hf_repo": "org/model-gguf", + "download_check": "pipeline.json", + "weight_variants": { + "param": "quant", + "default": "Q5", + "options": [variant("Q4"), variant("Q5")], + }, + }], + } + (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 = "quantized/generate" + manifest_variants = self.registry.get_manifest("quantized/generate")["weight_variants"] + self.assertEqual([option["id"] for option in manifest_variants["options"]], ["Q4", "Q5"]) + + model_root = self.models_dir / "quantized" / "generate" + model_root.mkdir(parents=True) + (model_root / "dit_Q5.gguf").write_bytes(b"q5") + self.registry.assert_weight_variant_installed({}) + self.registry.assert_weight_variant_installed({"quant": "Q5"}) + self.registry.assert_weight_variant_installed({"quant": "fp16"}) + with self.assertRaisesRegex(RuntimeError, "Q4 weights for quantized/generate are not installed"): + self.registry.assert_weight_variant_installed({"quant": "Q4"}) + 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_sources.py b/api/tests/test_model_sources.py index cdab245a..6bed7b73 100644 --- a/api/tests/test_model_sources.py +++ b/api/tests/test_model_sources.py @@ -4,8 +4,11 @@ from pathlib import Path from services.model_sources import ( + installed_weight_variants, + missing_weight_variant, model_sources_are_downloaded, normalize_model_sources, + normalize_weight_variants, resolve_model_root, validate_source_file_plan, ) @@ -112,6 +115,71 @@ def test_requires_all_checks_and_rejects_symlinked_extension_ancestry(self) -> N resolve_model_root(models, "pixal3d/generate") self.assertFalse(model_sources_are_downloaded(models, "pixal3d/generate", sources)) + def test_weight_variants_validate_and_report_the_missing_selection(self) -> None: + def variant(quant: str) -> dict: + return { + "id": quant, + "include_prefixes": [f"dit/model_{quant}.gguf"], + "checks": [f"dit/model_{quant}.gguf"], + } + + node = { + "hf_repo": "org/model-gguf", + "download_check": "pipeline.json", + "params_schema": [ + {"id": "quant", "type": "select", "options": [{"value": "Q4"}, {"value": "Q5"}]} + ], + "weight_variants": {"param": "quant", "options": [variant("Q4"), variant("Q5")]}, + } + variants = normalize_weight_variants(node) or {} + self.assertEqual(variants["default"], "Q4") + self.assertIsNone(normalize_weight_variants({"hf_repo": "org/model"})) + + def with_options(options: list, **extra) -> dict: + return {**node, "weight_variants": {"param": "quant", "options": options, **extra}} + + overlapping = {**variant("Q5"), "include_prefixes": ["dit/"]} + for broken, message in ( + ({**node, "hf_repo": ""}, "requires hf_repo"), + ({**node, "model_sources": []}, "cannot be combined"), + ({**node, "download_check": "dit/model_Q4.gguf"}, "download_check"), + (with_options([variant("Q4"), overlapping]), "share files"), + (with_options([variant("Q4")], default="Q8"), "default"), + (with_options([{**variant("Q4"), "checks": ["x.gguf"]}]), "not covered"), + (with_options([{**variant("Q4"), "size_gb": True}]), "size_gb"), + (with_options([{**variant("Q4"), "vram_gb": 0}]), "vram_gb"), + (with_options([{**variant("Q4"), "vram_gb": "6"}]), "vram_gb"), + ({**node, "params_schema": []}, "must name a params_schema entry"), + ( + {**node, "params_schema": [{"id": "quant", "options": [{"value": "Q5"}]}]}, + "must offer every weight variant id", + ), + ): + with self.subTest(message=message), self.assertRaisesRegex(ValueError, message): + normalize_weight_variants(broken) + + # A param that declares no options is left to the node to interpret. + self.assertIsNotNone( + normalize_weight_variants({**node, "params_schema": [{"id": "quant", "type": "string"}]}) + ) + self.assertNotIn("vram_gb", variants["options"][0]) + with_vram = normalize_weight_variants(with_options([{**variant("Q4"), "vram_gb": 6.5}])) or {} + self.assertEqual(with_vram["options"][0]["vram_gb"], 6.5) + + with tempfile.TemporaryDirectory(prefix="modly-weight-variants-") as tmp: + models = Path(tmp) / "models" + dit = models / "trellis" / "generate" / "dit" + dit.mkdir(parents=True) + (dit / "model_Q5.gguf").write_bytes(b"q5") + (dit / "model_Q4.gguf").write_bytes(b"") + self.assertEqual(installed_weight_variants(models, "trellis/generate", variants), ["Q5"]) + missing = missing_weight_variant(models, "trellis/generate", variants, {}) + self.assertEqual((missing or {}).get("id"), "Q4") + for params in ({"quant": "Q5"}, {"quant": "fp16"}): + with self.subTest(params=params): + self.assertIsNone(missing_weight_variant(models, "trellis/generate", variants, params)) + self.assertIsNone(missing_weight_variant(models, "trellis/generate", None, {"quant": "Q4"})) + if __name__ == "__main__": unittest.main() diff --git a/electron/main/extension-install-utils.test.mjs b/electron/main/extension-install-utils.test.mjs index 84139f9a..de49d929 100644 --- a/electron/main/extension-install-utils.test.mjs +++ b/electron/main/extension-install-utils.test.mjs @@ -102,6 +102,32 @@ test('validateInstallManifest rejects malformed or process model_sources', () => }, { hasEntryFile: () => true, hasGeneratorFile: () => false }, 'repository'), /only for model nodes/i) }) +test('validateInstallManifest validates weight variants and keeps them off process nodes', () => { + const mod = loadModule() + const files = { hasEntryFile: () => true, hasGeneratorFile: () => true } + const weightVariants = { + param: 'quant', + options: [{ id: 'Q4', include_prefixes: ['dit/model_Q4.gguf'], checks: ['dit/model_Q4.gguf'] }], + } + const paramsSchema = [{ id: 'quant', type: 'select', options: [{ value: 'Q4' }] }] + assert.doesNotThrow(() => mod.validateInstallManifest({ + id: 'quantized', generator_class: 'Generator', params_schema: paramsSchema, + nodes: [{ id: 'generate', hf_repo: 'org/model', weight_variants: weightVariants }], + }, files, 'repository')) + assert.throws(() => mod.validateInstallManifest({ + id: 'quantized', generator_class: 'Generator', + nodes: [{ id: 'generate', hf_repo: 'org/model', weight_variants: weightVariants }], + }, files, 'repository'), /must name a params_schema entry/) + assert.throws(() => mod.validateInstallManifest({ + id: 'quantized', generator_class: 'Generator', params_schema: paramsSchema, + nodes: [{ id: 'generate', weight_variants: weightVariants }], + }, files, 'repository'), /requires hf_repo/) + assert.throws(() => mod.validateInstallManifest({ + id: 'proc', type: 'process', entry: 'processor.js', + nodes: [{ id: 'run', hf_repo: 'org/model', weight_variants: weightVariants }], + }, files, 'repository'), /weight_variants is supported only for model nodes/) +}) + 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 05b965b0..6fc14e7f 100644 --- a/electron/main/extension-install-utils.ts +++ b/electron/main/extension-install-utils.ts @@ -1,7 +1,9 @@ import { normalizeModelSources, + normalizeWeightVariants, safeModelSourceId, type ModelSourceNode, + type WeightVariantNode, } from './model-sources' export interface InstallManifest { @@ -10,7 +12,8 @@ export interface InstallManifest { entry?: string generator_class?: string model_sources?: unknown - nodes?: Array<{ id?: string; model_sources?: unknown } & ModelSourceNode> + params_schema?: unknown + nodes?: Array<{ id?: string; model_sources?: unknown } & ModelSourceNode & WeightVariantNode> } export interface ValidatedInstallManifest { @@ -50,12 +53,14 @@ export function validateInstallManifest( 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 (node.model_sources === undefined && node.weight_variants === undefined) continue if (isProcess) { - throw new Error('manifest.json: model_sources is supported only for model nodes') + const field = node.model_sources !== undefined ? 'model_sources' : 'weight_variants' + throw new Error(`manifest.json: ${field} is supported only for model nodes`) } safeModelSourceId(node.id, 'model node id') normalizeModelSources(node) + normalizeWeightVariants(node, node.params_schema ?? manifest.params_schema) } if (isProcess) { diff --git a/electron/main/ipc-handlers.ts b/electron/main/ipc-handlers.ts index 60f3cf2b..deb9b3e3 100644 --- a/electron/main/ipc-handlers.ts +++ b/electron/main/ipc-handlers.ts @@ -14,12 +14,16 @@ import { listDownloadedModels, downloadModelFromHF, downloadModelSourcesFromHF, + type DownloadProgress, } from './model-downloader' -import { resolveInstalledModelDownloadPlan } from './model-download-plan' +import { legacyDownloadSteps, resolveInstalledModelDownloadPlan } from './model-download-plan' import { areModelSourcesDownloaded, + installedWeightVariants, + listWeightVariantFiles, modelHasLocalData, normalizeModelSources, + normalizeWeightVariants, removePartialDownloadArtifacts, resolveModelRoot, } from './model-sources' @@ -149,11 +153,29 @@ const renameWithRetry = (from: string, to: string, label: string) => export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGetter): void { type ActiveDownload = { - progress: { percent: number; file?: string; fileIndex?: number; totalFiles?: number } + progress: DownloadProgress & { variantId?: string } done: Promise finish: () => void } const activeDownloads = new Map() + const resolveModelPlan = (modelId: unknown) => resolveInstalledModelDownloadPlan({ + modelId, + userExtensionsDir: getSettings(app.getPath('userData')).extensionsDir, + builtinExtensionsDir: getBuiltinExtensionsDir(), + blockedExtensionIds: activeExtensionInstalls, + }) + const LOCKED_MODEL_FILES_ERROR = 'Model files are still locked after several attempts. Close any programs using the model and try again.' + + // Unload and wait for confirmation so file handles are released before removal. + async function unloadModelBeforeRemoval(modelId: string): Promise { + try { + await axios.post(`${API_BASE_URL}/model/unload/${encodeURIComponent(modelId)}`, {}, { timeout: 10_000 }) + // Give the OS a moment to release file locks (Windows holds handles briefly after close) + await new Promise(resolve => setTimeout(resolve, 1_500)) + } catch { + // Unload failed (model may not be loaded) — still attempt deletion + } + } // 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')) @@ -339,35 +361,45 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe } let modelDir: string try { - await resolveInstalledModelDownloadPlan({ - modelId, - userExtensionsDir: getSettings(app.getPath('userData')).extensionsDir, - builtinExtensionsDir: getBuiltinExtensionsDir(), - blockedExtensionIds: activeExtensionInstalls, - }) + await resolveModelPlan(modelId) 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 { - await axios.post(`${API_BASE_URL}/model/unload/${encodeURIComponent(modelId)}`, {}, { timeout: 10_000 }) - // Give the OS a moment to release file locks (Windows holds handles briefly after close) - await new Promise(resolve => setTimeout(resolve, 1_500)) - } catch { - // Unload failed (model may not be loaded) — still attempt deletion - } + await unloadModelBeforeRemoval(modelId) // Retry removal — Windows may return EBUSY/EPERM if handles linger const removed = await rmWithRetry(modelDir, 'model-delete') if (removed.ok) return { success: true } - return { - success: false, - error: removed.locked - ? 'Model files are still locked after several attempts. Close any programs using the model and try again.' - : String(removed.error), + return { success: false, error: removed.locked ? LOCKED_MODEL_FILES_ERROR : String(removed.error) } + }) + + ipcMain.handle('model:deleteWeightVariant', async (_, modelId: string, variantId: string): Promise<{ success: boolean; error?: string }> => { + if (activeDownloads.has(modelId)) { + return { success: false, error: 'Cannot remove model weights while their download is active' } } + let files: string[] + try { + const plan = await resolveModelPlan(modelId) + const variant = plan.kind === 'legacy' + ? plan.weightVariants?.options.find((option) => option.id === variantId) + : undefined + if (!variant) throw new Error(`Model node "${modelId}" has no weight variant "${String(variantId)}"`) + files = await listWeightVariantFiles(getSettings(app.getPath('userData')).modelsDir, modelId, variant) + } catch (err) { + return { success: false, error: String(err) } + } + if (files.length === 0) return { success: true } + + await unloadModelBeforeRemoval(modelId) + for (const file of files) { + const removed = await rmWithRetry(file, 'model-variant-delete') + if (!removed.ok) { + return { success: false, error: removed.locked ? LOCKED_MODEL_FILES_ERROR : String(removed.error) } + } + } + return { success: true } }) ipcMain.handle('model:showInFolder', (_, modelId: string) => { @@ -403,28 +435,31 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe ipcMain.handle('model:isDownloaded', async (_, modelId: string): Promise => { const modelsDir = getSettings(app.getPath('userData')).modelsDir 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) + const plan = await resolveModelPlan(modelId) + if (plan.kind === 'multi-source') return areModelSourcesDownloaded(modelsDir, modelId, plan.sources) + return isModelDownloaded(modelsDir, modelId, plan.downloadCheck) + && (!plan.weightVariants || installedWeightVariants(modelsDir, modelId, plan.weightVariants).length > 0) } catch { return false } }) + // null means "unknown" (unreadable plan, or a node without variants) — the renderer + // must not read an empty array as "no variant installed". + ipcMain.handle('model:installedWeightVariants', async (_, modelId: string): Promise => { + try { + const plan = await resolveModelPlan(modelId) + return plan.kind === 'legacy' && plan.weightVariants + ? installedWeightVariants(getSettings(app.getPath('userData')).modelsDir, modelId, plan.weightVariants) + : null + } catch { + return null + } + }) + ipcMain.handle('model:hasLocalData', async (_, modelId: string): Promise => { try { - await resolveInstalledModelDownloadPlan({ - modelId, - userExtensionsDir: getSettings(app.getPath('userData')).extensionsDir, - builtinExtensionsDir: getBuiltinExtensionsDir(), - blockedExtensionIds: activeExtensionInstalls, - }) + await resolveModelPlan(modelId) return modelHasLocalData(getSettings(app.getPath('userData')).modelsDir, modelId) } catch { return false @@ -438,45 +473,47 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe ipcMain.handle('model:download', async ( event, modelId: string, + requestedVariantId?: string | null, ) => { if (activeDownloads.has(modelId)) { return { success: false, error: 'Download already in progress' } } + const variantId = requestedVariantId ?? undefined let finish!: () => void const done = new Promise((resolveDone) => { finish = resolveDone }) - const active: ActiveDownload = { progress: { percent: 0 }, done, finish } + const active: ActiveDownload = { progress: { percent: 0, variantId }, done, finish } activeDownloads.set(modelId, active) try { - 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 }) + const plan = await resolveModelPlan(modelId) + const onProgress = (progress: DownloadProgress) => { + active.progress = { ...progress, variantId } + event.sender.send('model:downloadProgress', { modelId, variantId, ...progress }) } if (plan.kind === 'multi-source') { + if (variantId !== undefined) throw new Error(`Model node "${modelId}" does not declare weight variants`) await downloadModelSourcesFromHF(modelId, plan.sources, onProgress) } else { - await downloadModelFromHF( - plan.repoId, - modelId, - onProgress, - plan.skipPrefixes, - plan.includePrefixes, - ) + // Shared files and the default variant are separate passes sharing one 0–100 bar. + const steps = legacyDownloadSteps(plan, variantId) + for (const [index, step] of steps.entries()) { + await downloadModelFromHF( + plan.repoId, + modelId, + (progress) => onProgress({ ...progress, percent: Math.round((index * 100 + progress.percent) / steps.length) }), + step.skipPrefixes, + step.includePrefixes, + ) + } } return { success: true } } catch (err: any) { const message = err?.message ?? String(err) if (message.includes('paused')) { - event.sender.send('model:downloadProgress', { modelId, percent: 0, status: 'paused', paused: true }) + event.sender.send('model:downloadProgress', { modelId, variantId, percent: 0, status: 'paused', paused: true }) return { success: false, paused: true } } if (message.includes('cancelled')) { - event.sender.send('model:downloadProgress', { modelId, percent: 0, status: 'cancelled', cancelled: true }) + event.sender.send('model:downloadProgress', { modelId, variantId, percent: 0, status: 'cancelled', cancelled: true }) return { success: false, cancelled: true } } return { success: false, error: String(err) } @@ -835,6 +872,7 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe hf_skip_prefixes?: string[] hf_include_prefixes?: string[] model_sources?: unknown + weight_variants?: unknown }[] } @@ -857,7 +895,11 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe if (parsed.type === 'process' && n.model_sources !== undefined) { throw new Error('manifest.json: model_sources is supported only for model nodes') } + if (parsed.type === 'process' && n.weight_variants !== undefined) { + throw new Error('manifest.json: weight_variants is supported only for model nodes') + } const modelSources = normalizeModelSources(n) + const weightVariants = normalizeWeightVariants(n, n.params_schema ?? parsed.params_schema) return { id: n.id, name: n.name ?? n.id, @@ -872,6 +914,16 @@ export function setupIpcHandlers(pythonBridge: PythonBridge, getWindow: WindowGe hfSkipPrefixes: n.hf_skip_prefixes, hfIncludePrefixes: n.hf_include_prefixes, hasModelSources: modelSources !== undefined, + weightVariants: weightVariants && { + param: weightVariants.param, + default: weightVariants.default, + options: weightVariants.options.map((option) => ({ + id: option.id, + label: option.label, + sizeGb: option.size_gb, + vramGb: option.vram_gb, + })), + }, } }) diff --git a/electron/main/model-download-plan.test.mjs b/electron/main/model-download-plan.test.mjs index e5d037a0..520aa780 100644 --- a/electron/main/model-download-plan.test.mjs +++ b/electron/main/model-download-plan.test.mjs @@ -90,3 +90,55 @@ test('keeps legacy sibling checks and wildcard filters unchanged', async () => { rmSync(fixture.root, { recursive: true, force: true }) } }) + +test('downloads a weight-variant node as shared files first, then one variant at a time', async () => { + const { legacyDownloadSteps, resolveInstalledModelDownloadPlan } = loadModule() + const fixture = setupExtension({ + id: 'trellis', + type: 'model', + nodes: [{ + id: 'generate', + hf_repo: 'org/model-gguf', + download_check: 'pipeline.json', + hf_include_prefixes: ['pipeline.json', 'dit/'], + hf_skip_prefixes: ['README.md'], + params_schema: [{ id: 'quant', type: 'select', options: [{ value: 'Q4' }, { value: 'Q5' }] }], + weight_variants: { + param: 'quant', + default: 'Q5', + options: ['Q4', 'Q5'].map((quant) => ({ + id: quant, + include_prefixes: [`dit/model_${quant}.gguf`], + checks: [`dit/model_${quant}.gguf`], + })), + }, + }], + }) + try { + const plan = await resolveInstalledModelDownloadPlan({ + modelId: 'trellis/generate', + userExtensionsDir: fixture.user, + builtinExtensionsDir: fixture.builtin, + }) + assert.equal(plan.kind, 'legacy') + assert.deepEqual(legacyDownloadSteps(plan), [ + { includePrefixes: ['pipeline.json', 'dit/'], skipPrefixes: ['README.md', 'dit/model_Q4.gguf', 'dit/model_Q5.gguf'] }, + { includePrefixes: ['dit/model_Q5.gguf'], skipPrefixes: ['README.md'] }, + ]) + // Picking one variant still fetches the shared files first, so a node installed + // variant-first is never left without its pipeline files. + assert.deepEqual(legacyDownloadSteps(plan, 'Q4'), [ + { includePrefixes: ['pipeline.json', 'dit/'], skipPrefixes: ['README.md', 'dit/model_Q4.gguf', 'dit/model_Q5.gguf'] }, + { includePrefixes: ['dit/model_Q4.gguf'], skipPrefixes: ['README.md'] }, + ]) + assert.throws(() => legacyDownloadSteps(plan, 'Q8'), /no weight variant "Q8"/) + + const withoutVariants = { ...plan, weightVariants: undefined } + assert.deepEqual(legacyDownloadSteps(withoutVariants), [ + { includePrefixes: ['pipeline.json', 'dit/'], skipPrefixes: ['README.md'] }, + ]) + assert.throws(() => legacyDownloadSteps(withoutVariants, 'Q4'), /no weight variant/) + } finally { + rmSync(fixture.root, { recursive: true, force: true }) + } +}) diff --git a/electron/main/model-download-plan.ts b/electron/main/model-download-plan.ts index 27d2dba6..e4a4905e 100644 --- a/electron/main/model-download-plan.ts +++ b/electron/main/model-download-plan.ts @@ -8,7 +8,13 @@ import { assertSafeExtensionId, resolveExtensionPathWithinRoot, } from './extension-path-guard' -import { normalizeModelSources, safeModelSourceId, type ModelSource } from './model-sources' +import { + normalizeModelSources, + normalizeWeightVariants, + safeModelSourceId, + type ModelSource, + type WeightVariants, +} from './model-sources' interface InstalledNode { id?: unknown @@ -17,12 +23,15 @@ interface InstalledNode { hf_skip_prefixes?: unknown hf_include_prefixes?: unknown model_sources?: unknown + weight_variants?: unknown + params_schema?: unknown } interface InstalledManifest { id?: unknown type?: unknown model_sources?: unknown + params_schema?: unknown nodes?: unknown } @@ -35,6 +44,7 @@ export type InstalledModelDownloadPlan = { downloadCheck?: string skipPrefixes?: string[] includePrefixes?: string[] + weightVariants?: WeightVariants } | { kind: 'multi-source' modelId: string @@ -86,6 +96,7 @@ function parseManifest(raw: string, extensionId: string, nodeId: string): Instal const node = matches[0] const modelId = `${extensionId}/${nodeId}` + const weightVariants = normalizeWeightVariants(node, node.params_schema ?? manifest.params_schema) const sources = normalizeModelSources(node) if (sources) return { kind: 'multi-source', modelId, extensionId, nodeId, sources } @@ -101,7 +112,43 @@ function parseManifest(raw: string, extensionId: string, nodeId: string): Instal 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, + ...(weightVariants ? { weightVariants } : {}), + } +} + +export interface LegacyDownloadStep { + includePrefixes?: string[] + skipPrefixes?: string[] +} + +/** + * Hugging Face filter passes for one download action. A node that declares weight + * variants always fetches its shared files first (every variant excluded), then the + * requested variant — its default one when no variant id is given. The shared pass + * skips files already complete on disk, so it stays cheap on a resume or a second + * variant. + */ +export function legacyDownloadSteps( + plan: Extract, + variantId?: unknown, +): LegacyDownloadStep[] { + const variants = plan.weightVariants + if (!variants) { + if (variantId !== undefined) { + throw new Error(`Model node "${plan.modelId}" has no weight variant "${String(variantId)}"`) + } + return [{ includePrefixes: plan.includePrefixes, skipPrefixes: plan.skipPrefixes }] } + const selected = variantId === undefined ? variants.default : variantId + const variant = variants.options.find((option) => option.id === selected) + if (!variant) throw new Error(`Model node "${plan.modelId}" has no weight variant "${String(variantId)}"`) + return [ + { + includePrefixes: plan.includePrefixes, + skipPrefixes: [...(plan.skipPrefixes ?? []), ...variants.options.flatMap((option) => option.include_prefixes)], + }, + { includePrefixes: variant.include_prefixes, skipPrefixes: plan.skipPrefixes }, + ] } /** Re-read the installed manifest for every model action; renderer metadata is never trusted. */ diff --git a/electron/main/model-sources.test.mjs b/electron/main/model-sources.test.mjs index da8a54fc..63183ce3 100644 --- a/electron/main/model-sources.test.mjs +++ b/electron/main/model-sources.test.mjs @@ -113,3 +113,87 @@ test('requires every declared check and rejects symlinked extension-root ancestr rmSync(root, { recursive: true, force: true }) } }) + +const quantNode = () => ({ + hf_repo: 'org/model-gguf', + download_check: 'pipeline.json', + params_schema: [{ + id: 'quant', + label: 'Quantization', + type: 'select', + default: 'Q5', + options: [{ value: 'Q4', label: 'Q4' }, { value: 'Q5', label: 'Q5' }], + }], + weight_variants: { + param: 'quant', + default: 'Q5', + options: ['Q4', 'Q5'].map((quant) => ({ + id: quant, + include_prefixes: [`dit/model_${quant}.gguf`], + checks: [`dit/model_${quant}.gguf`], + })), + }, +}) + +test('validates weight variants and rejects ambiguous or unsafe declarations', () => { + const { normalizeWeightVariants } = loadModule() + const variants = normalizeWeightVariants(quantNode()) + assert.equal(variants.param, 'quant') + assert.equal(variants.default, 'Q5') + assert.deepEqual(variants.options.map((option) => option.label), ['Q4', 'Q5']) + assert.equal(normalizeWeightVariants({ hf_repo: 'org/model' }), undefined) + + const node = quantNode() + const [q4, q5] = node.weight_variants.options + const withOptions = (options, extra = {}) => ({ ...node, weight_variants: { ...node.weight_variants, options, ...extra } }) + const cases = [ + [{ ...node, hf_repo: undefined }, /requires hf_repo/], + [{ ...node, model_sources: [] }, /cannot be combined/], + [{ ...node, download_check: 'dit/model_Q4.gguf' }, /download_check/], + [withOptions([q4, q5], { default: 'Q8' }), /default/], + [withOptions([q4, { ...q5, include_prefixes: ['dit/'] }]), /share files/], + [withOptions([q4, { ...q5, id: 'q4' }]), /portable-unique/], + [withOptions([{ ...q4, checks: ['other.gguf'] }]), /not covered/], + [withOptions([{ ...q4, include_prefixes: ['../outside'] }]), /unsafe/], + [withOptions([{ ...q4, size_gb: -1 }]), /size_gb/], + [withOptions([{ ...q4, vram_gb: 0 }]), /vram_gb/], + [withOptions([{ ...q4, vram_gb: '6' }]), /vram_gb/], + [{ ...node, params_schema: undefined }, /must name a params_schema entry/], + [{ ...node, params_schema: [{ id: 'steps', type: 'int' }] }, /must name a params_schema entry/], + [ + { ...node, params_schema: [{ id: 'quant', type: 'select', options: [{ value: 'Q5' }] }] }, + /must offer every weight variant id \(missing: Q4\)/, + ], + ] + for (const [candidate, error] of cases) assert.throws(() => normalizeWeightVariants(candidate), error) + + // A param that declares no options (or an unknown shape) is left to the node. + assert.doesNotThrow(() => normalizeWeightVariants({ ...node, params_schema: [{ id: 'quant', type: 'string' }] })) + assert.equal('vram_gb' in variants.options[0], false) + assert.equal(normalizeWeightVariants(withOptions([{ ...q4, vram_gb: 6.5 }, q5])).options[0].vram_gb, 6.5) +}) + +test('reports installed variants and lists only the files of the variant being removed', async () => { + const { installedWeightVariants, listWeightVariantFiles, normalizeWeightVariants } = loadModule() + const variants = normalizeWeightVariants(quantNode()) + const root = mkdtempSync(join(tmpdir(), 'modly-weight-variants-')) + const models = join(root, 'models') + const nodeRoot = join(models, 'trellis', 'generate') + mkdirSync(join(nodeRoot, 'dit'), { recursive: true }) + writeFileSync(join(nodeRoot, 'pipeline.json'), '{}') + writeFileSync(join(nodeRoot, 'dit', 'model_Q5.gguf'), 'q5') + writeFileSync(join(nodeRoot, 'dit', 'model_Q4.gguf'), '') + writeFileSync(join(nodeRoot, 'dit', 'model_Q4.gguf.part'), 'partial') + + try { + assert.deepEqual(installedWeightVariants(models, 'trellis/generate', variants), ['Q5']) + assert.deepEqual(installedWeightVariants(models, '../escape', variants), []) + const files = await listWeightVariantFiles(models, 'trellis/generate', variants.options[0]) + assert.deepEqual( + files.map((file) => file.slice(nodeRoot.length + 1).replaceAll('\\', '/')).sort(), + ['dit/model_Q4.gguf', 'dit/model_Q4.gguf.part'], + ) + } finally { + rmSync(root, { recursive: true, force: true }) + } +}) diff --git a/electron/main/model-sources.ts b/electron/main/model-sources.ts index dac76930..291c05f3 100644 --- a/electron/main/model-sources.ts +++ b/electron/main/model-sources.ts @@ -1,6 +1,6 @@ import { existsSync, lstatSync, readdirSync, statSync } from 'node:fs' import { readdir, rm } from 'node:fs/promises' -import { isAbsolute, relative, resolve } from 'node:path' +import { isAbsolute, relative, resolve, sep } from 'node:path' export interface ModelSource { id: string @@ -17,6 +17,31 @@ export interface ModelSourceNode { model_sources?: unknown } +/** One separately installable set of files inside a node's model directory (e.g. a quantization). */ +export interface WeightVariant { + id: string + label: string + size_gb?: number + vram_gb?: number + include_prefixes: string[] + checks: string[] +} + +export interface WeightVariants { + /** params_schema id whose value selects the variant at generation time */ + param: string + default: string + options: WeightVariant[] +} + +export interface WeightVariantNode { + weight_variants?: unknown + hf_repo?: unknown + download_check?: unknown + model_sources?: unknown + params_schema?: 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]/ @@ -133,6 +158,114 @@ export function normalizeModelSources(node: ModelSourceNode): ModelSource[] | un }) } +function prefixesOverlap(left: string[], right: string[]): boolean { + return left.some((a) => right.some((b) => { + const lowerA = a.toLowerCase() + const lowerB = b.toLowerCase() + return lowerA.startsWith(lowerB) || lowerB.startsWith(lowerA) + })) +} + +/** Variant ids must be selectable, so the param they key has to exist and offer them. */ +function assertParamOffersVariants(param: string, paramsSchema: unknown, ids: string[]): void { + const entry = Array.isArray(paramsSchema) + ? paramsSchema.find((p) => typeof p === 'object' && p !== null && (p as { id?: unknown }).id === param) + : undefined + if (!entry) throw new Error(`weight_variants.param must name a params_schema entry ("${param}")`) + const options = (entry as { options?: unknown }).options + if (!Array.isArray(options)) return + const values = new Set(options.map((option) => + typeof option === 'object' && option !== null ? String((option as { value?: unknown }).value) : String(option))) + const missing = ids.filter((id) => !values.has(id)) + if (missing.length > 0) { + throw new Error(`the "${param}" param must offer every weight variant id (missing: ${missing.join(', ')})`) + } +} + +export function normalizeWeightVariants( + node: WeightVariantNode, + paramsSchema?: unknown, +): WeightVariants | undefined { + if (!Object.prototype.hasOwnProperty.call(node, 'weight_variants')) return undefined + if (Object.prototype.hasOwnProperty.call(node, 'model_sources')) { + throw new Error('weight_variants cannot be combined with model_sources') + } + if (typeof node.hf_repo !== 'string' || !node.hf_repo) { + throw new Error('weight_variants requires hf_repo on the same node') + } + const raw = node.weight_variants + if (typeof raw !== 'object' || raw === null || Array.isArray(raw)) { + throw new Error('weight_variants must be an object') + } + const value = raw as Record + const param = safeModelSourceId(value.param, 'weight_variants.param') + if (!Array.isArray(value.options) || value.options.length === 0) { + throw new Error('weight_variants.options must be a non-empty array') + } + + const seen = new Map() + const options = value.options.map((rawOption, index): WeightVariant => { + const field = `weight_variants.options[${index}]` + if (typeof rawOption !== 'object' || rawOption === null || Array.isArray(rawOption)) { + throw new Error(`${field} must be an object`) + } + const option = rawOption as Record + const id = safeModelSourceId(option.id, `${field}.id`) + const alias = id.normalize('NFC').toLowerCase() + const previous = seen.get(alias) + if (previous) throw new Error(`weight variant ids "${previous}" and "${id}" are not portable-unique`) + seen.set(alias, id) + const label = option.label ?? id + if (typeof label !== 'string' || !label.trim()) throw new Error(`${field}.label must be a non-empty string`) + const sizeGb = option.size_gb + if (sizeGb !== undefined && (typeof sizeGb !== 'number' || !Number.isFinite(sizeGb) || sizeGb <= 0)) { + throw new Error(`${field}.size_gb must be a positive number`) + } + const vramGb = option.vram_gb + if (vramGb !== undefined && (typeof vramGb !== 'number' || !Number.isFinite(vramGb) || vramGb <= 0)) { + throw new Error(`${field}.vram_gb must be a positive number`) + } + const include = optionalPrefixes(option.include_prefixes, `${field}.include_prefixes`) + if (!include?.length) throw new Error(`${field}.include_prefixes must be a non-empty array`) + if (!Array.isArray(option.checks) || option.checks.length === 0) { + throw new Error(`${field}.checks must be a non-empty array`) + } + const checks = option.checks.map((check, checkIndex) => { + const path = safeModelRelativePath(check, `${field}.checks[${checkIndex}]`) + if (!include.some((prefix) => path.startsWith(prefix))) { + throw new Error(`${field}.checks[${checkIndex}] is not covered by its include_prefixes`) + } + return path + }) + return { + id, + label, + ...(sizeGb === undefined ? {} : { size_gb: sizeGb }), + ...(vramGb === undefined ? {} : { vram_gb: vramGb }), + include_prefixes: include, + checks, + } + }) + + options.forEach((variant, index) => { + for (const other of options.slice(index + 1)) { + if (prefixesOverlap(variant.include_prefixes, other.include_prefixes)) { + throw new Error(`weight variants "${variant.id}" and "${other.id}" share files`) + } + } + }) + const downloadCheck = node.download_check + if (typeof downloadCheck === 'string' && options.some((variant) => prefixesOverlap([downloadCheck], variant.include_prefixes))) { + throw new Error('download_check must name a file outside every weight variant') + } + const defaultId = value.default ?? options[0].id + if (!options.some((variant) => variant.id === defaultId)) { + throw new Error('weight_variants.default must name one of its options') + } + assertParamOffersVariants(param, paramsSchema ?? node.params_schema, options.map((variant) => variant.id)) + return { param, default: defaultId as string, options } +} + function pathHasSymlink(root: string, candidate: string): boolean { const rootPath = resolve(root) const rel = relative(rootPath, resolve(candidate)) @@ -196,6 +329,43 @@ export function modelHasLocalData(modelsDir: string, modelId: string): boolean { } } +function isDownloadedFile(root: string, relativePath: string): boolean { + const candidate = resolve(root, ...relativePath.split('/')) + if (!existsSync(candidate) || pathHasSymlink(root, candidate)) return false + try { + const stat = statSync(candidate) + return stat.isFile() && stat.size > 0 + } catch { + return false + } +} + +export function installedWeightVariants(modelsDir: string, modelId: string, variants: WeightVariants): string[] { + try { + const modelRoot = resolveModelRoot(modelsDir, modelId) + return variants.options + .filter((variant) => variant.checks.every((check) => isDownloadedFile(modelRoot, check))) + .map((variant) => variant.id) + } catch { + return [] + } +} + +/** Files of one variant on disk (in-progress `.part` files included), never through a symlink. */ +export async function listWeightVariantFiles(modelsDir: string, modelId: string, variant: WeightVariant): Promise { + const modelRoot = resolveModelRoot(modelsDir, modelId) + if (!existsSync(modelRoot)) return [] + const entries = await readdir(modelRoot, { recursive: true, withFileTypes: true }) + return entries + .filter((entry) => entry.isFile()) + .map((entry) => resolve(entry.parentPath ?? modelRoot, entry.name)) + .filter((path) => { + const relativePath = relative(modelRoot, path).split(sep).join('/') + return variant.include_prefixes.some((prefix) => relativePath.startsWith(prefix)) + && !pathHasSymlink(modelRoot, path) + }) +} + // Mirrors the backend's cancel cleanup (api/routers/model.py): only the in-progress // `.part` files are removed, so completed sources already on disk survive a cancel. export async function removePartialDownloadArtifacts(modelRoot: string): Promise { diff --git a/electron/preload/electron-api.ts b/electron/preload/electron-api.ts index dae65e69..69c033f4 100644 --- a/electron/preload/electron-api.ts +++ b/electron/preload/electron-api.ts @@ -119,10 +119,14 @@ export function createElectronApi(ipcRenderer: IpcRendererLike, webFrame: WebFra listDownloaded: () => ipcRenderer.invoke('model:listDownloaded'), isDownloaded: (modelId: string) => ipcRenderer.invoke('model:isDownloaded', modelId), hasLocalData: (modelId: string) => ipcRenderer.invoke('model:hasLocalData', modelId), - download: (modelId: string) => ipcRenderer.invoke('model:download', modelId), + download: (modelId: string, variantId?: string) => (variantId === undefined + ? ipcRenderer.invoke('model:download', modelId) + : ipcRenderer.invoke('model:download', modelId, variantId)), pauseDownload: (modelId: string) => ipcRenderer.invoke('model:pauseDownload', modelId), cancelDownload: (modelId: string) => ipcRenderer.invoke('model:cancelDownload', modelId), delete: (modelId: string) => ipcRenderer.invoke('model:delete', modelId), + installedWeightVariants: (modelId: string) => ipcRenderer.invoke('model:installedWeightVariants', modelId), + deleteWeightVariant: (modelId: string, variantId: string) => ipcRenderer.invoke('model:deleteWeightVariant', modelId, variantId), unloadAll: () => ipcRenderer.invoke('model:unloadAll'), showInFolder: (modelId: string) => ipcRenderer.invoke('model:showInFolder', modelId), activeDownloads: (): Promise<{ modelId: string; percent: number; file?: string; fileIndex?: number; totalFiles?: number }[]> => diff --git a/src/areas/generate/components/WorkflowPanel.tsx b/src/areas/generate/components/WorkflowPanel.tsx index 2995ddea..c15034b5 100644 --- a/src/areas/generate/components/WorkflowPanel.tsx +++ b/src/areas/generate/components/WorkflowPanel.tsx @@ -17,6 +17,7 @@ import type { WorkflowExtension } from '@areas/workflows/mockExtensions' import type { Workflow, WFNode, WFEdge, ParamSchema } from '@shared/types/electron.d' import { PICKER_LABELS, openParamPicker, resolvePickerIntent } from '@shared/utils/paramPicker' import { PickerIcon } from '@shared/components/ui' +import { isMissingWeightVariant, withWeightVariantAvailability } from '@shared/utils/weightVariants' import ChatPanel from './ChatPanel' type PanelMode = 'basic' | 'chat' @@ -383,6 +384,8 @@ function WaitParamRow({ nodeId }: { nodeId: string }) { function ExtensionParamRow({ nodeId, ext, nodes, onPatch }: { nodeId: string; ext: WorkflowExtension; nodes: FlowNode[]; onPatch: PatchFn }) { const [expanded, setExpanded] = useState(true) + const installedVariants = useExtensionsStore((s) => s.installedWeightVariants[ext.id]) + const openExtension = useNavStore((s) => s.openExtension) const node = nodes.find((n) => n.id === nodeId) const data = node?.data as { enabled: boolean; params: Record } | undefined const enabled = data?.enabled ?? true @@ -430,8 +433,11 @@ function ExtensionParamRow({ nodeId, ext, nodes, onPatch }: { nodeId: string; ex
- onPatch(nodeId, { params: { ...(data?.params ?? {}), [param.id]: v } })} /> + { + onPatch(nodeId, { params: { ...(data?.params ?? {}), [param.id]: v } }) + if (isMissingWeightVariant(param.id, v, ext.weightVariants, installedVariants)) openExtension(ext.extensionId) + }} />
) diff --git a/src/areas/models/ModelsPage.tsx b/src/areas/models/ModelsPage.tsx index 66f8e5d4..c4eaba0e 100644 --- a/src/areas/models/ModelsPage.tsx +++ b/src/areas/models/ModelsPage.tsx @@ -1,12 +1,13 @@ import { useEffect, useMemo, useRef, useState } from 'react' import { createPortal } from 'react-dom' import { useExtensionsStore } from '@shared/stores/extensionsStore' +import { useNavStore } from '@shared/stores/navStore' import type { AnyExtension, ModelExtension } from '@shared/types/electron.d' import { deleteModelsThenUninstallExtension, formatModelName } from './utils' import { ExtensionCard } from './components/ExtensionCard' import type { ExtensionNode } from './components/ExtensionCard' import { ExtensionDrawer } from './components/ExtensionDrawer' -import { ICONS, nodeHasManagedWeights } from './components/extensionShared' +import { ICONS, nodeHasManagedWeights, type DownloadMap } from './components/extensionShared' // ─── Filters & sorts ────────────────────────────────────────────────────────── @@ -48,19 +49,9 @@ export default function ModelsPage(): JSX.Element { ) // Model weight state (needed for node install status + uninstall cleanup) - const [installedVariantIds, setInstalledVariantIds] = useState([]) + const [installedNodeIds, setInstalledNodeIds] = useState([]) const [localDataIds, setLocalDataIds] = useState([]) - const [downloading, setDownloading] = useState>({}) + const [downloading, setDownloading] = useState({}) // Uninstall modal state const [uninstallTarget, setUninstallTarget] = useState(null) @@ -75,6 +66,15 @@ export default function ModelsPage(): JSX.Element { const [selectedId, setSelectedId] = useState(null) const searchRef = useRef(null) + // Another page asked to open an extension (e.g. a weight variant picked but not installed) + const extensionToOpen = useNavStore((s) => s.extensionToOpen) + const clearExtensionToOpen = useNavStore((s) => s.clearExtensionToOpen) + useEffect(() => { + if (!extensionToOpen) return + setSelectedId(extensionToOpen) + clearExtensionToOpen() + }, [extensionToOpen, clearExtensionToOpen]) + // GitHub extension install form const [showGHForm, setShowGHForm] = useState(false) const [ghUrl, setGhUrl] = useState('') @@ -98,8 +98,9 @@ export default function ModelsPage(): JSX.Element { if (hasLocalData) localIds.push(fullId) } } - setInstalledVariantIds(ids) + setInstalledNodeIds(ids) setLocalDataIds(localIds) + await useExtensionsStore.getState().refreshInstalledWeightVariants() } useEffect(() => { @@ -115,7 +116,7 @@ export default function ModelsPage(): JSX.Element { } refreshInstalledIds(exts) }) - window.electron.model.onProgress(({ modelId: id, percent, file, fileIndex, totalFiles, status, bytesDownloaded, totalBytes, stalledSeconds, paused, cancelled }) => { + window.electron.model.onProgress(({ modelId: id, variantId, percent, file, fileIndex, totalFiles, status, bytesDownloaded, totalBytes, stalledSeconds, paused, cancelled }) => { if (cancelled) { setDownloading((prev) => { const n = { ...prev }; delete n[id]; return n }) return @@ -134,6 +135,7 @@ export default function ModelsPage(): JSX.Element { totalBytes: totalBytes ?? current?.totalBytes, stalledSeconds: stalledSeconds ?? current?.stalledSeconds, paused, + variantId: variantId ?? current?.variantId, }, } }) @@ -167,10 +169,10 @@ export default function ModelsPage(): JSX.Element { // ── Node install / download controls ────────────────────────────────────── - function handleInstallNode(node: ExtensionNode, fullId: string) { + function handleInstallNode(node: ExtensionNode, fullId: string, variantId?: string) { if (!nodeHasManagedWeights(node)) return - setDownloading((prev) => ({ ...prev, [fullId]: { ...(prev[fullId] ?? { percent: 0 }), paused: false, status: 'Starting…' } })) - window.electron.model.download(fullId).then((result) => { + setDownloading((prev) => ({ ...prev, [fullId]: { ...(prev[fullId] ?? { percent: 0 }), variantId, paused: false, status: 'Starting…' } })) + window.electron.model.download(fullId, variantId).then((result) => { if (!result.success && !result.paused && !result.cancelled) { setGhErr(result.error ?? 'Download failed') setDownloading((prev) => { const n = { ...prev }; delete n[fullId]; return n }) @@ -183,7 +185,7 @@ export default function ModelsPage(): JSX.Element { for (const node of ext.nodes) { if (!nodeHasManagedWeights(node)) continue const fullId = `${ext.id}/${node.id}` - if (installedVariantIds.includes(fullId) || downloading[fullId]) continue + if (installedNodeIds.includes(fullId) || downloading[fullId]) continue handleInstallNode(node, fullId) } } @@ -205,6 +207,12 @@ export default function ModelsPage(): JSX.Element { refreshInstalledIds(useExtensionsStore.getState().modelExtensions) } + async function handleDeleteWeightVariant(fullId: string, variantId: string) { + const result = await window.electron.model.deleteWeightVariant(fullId, variantId) + await refreshInstalledIds(useExtensionsStore.getState().modelExtensions) + return result + } + // ── GitHub extension install ─────────────────────────────────────────────── async function handleGHInstall() { @@ -322,7 +330,7 @@ export default function ModelsPage(): JSX.Element { } const cardHandlers = { - installedIds: installedVariantIds, + installedIds: installedNodeIds, downloading, disabled: isBusy, onInstall: handleInstallNode, @@ -627,7 +635,7 @@ export default function ModelsPage(): JSX.Element { {selectedExt && ( openUninstallModal(extId)} onRepaired={() => reloadExtensions()} onSynced={() => reloadExtensions()} diff --git a/src/areas/models/components/ExtensionCard.tsx b/src/areas/models/components/ExtensionCard.tsx index 7b0da394..5ceb729a 100644 --- a/src/areas/models/components/ExtensionCard.tsx +++ b/src/areas/models/components/ExtensionCard.tsx @@ -125,7 +125,12 @@ export function ExtensionCard({ className="flex items-center justify-between gap-2.5 px-2.5 py-1.5 rounded-lg bg-white/[0.02] border border-zinc-800" >
- {node.name} + + {node.name} + {isModel && node.weightVariants && ( + · {node.weightVariants.options.length} variants + )} +
{isModel && ( diff --git a/src/areas/models/components/ExtensionDrawer.tsx b/src/areas/models/components/ExtensionDrawer.tsx index 2f154c71..288f12b1 100644 --- a/src/areas/models/components/ExtensionDrawer.tsx +++ b/src/areas/models/components/ExtensionDrawer.tsx @@ -1,11 +1,13 @@ import { useEffect, useState } from 'react' import type { AnyExtension, ExtensionNode } from '@shared/types/electron.d' +import { useExtensionsStore } from '@shared/stores/extensionsStore' import { useNavStore } from '@shared/stores/navStore' import { DownloadMap, ICONS, IOBadge, NodeInstallControl, + NodeUiState, TypePill, extInstallSummary, formatBytes, @@ -20,11 +22,12 @@ interface Props { downloading: DownloadMap loadError?: string disabled?: boolean - onInstall: (node: ExtensionNode, fullId: string) => void + onInstall: (node: ExtensionNode, fullId: string, variantId?: string) => void onInstallAll: (ext: AnyExtension) => void onPauseDownload: (fullId: string) => void onCancelDownload: (fullId: string) => void onUninstallNode: (fullId: string) => void + onDeleteWeightVariant: (fullId: string, variantId: string) => Promise<{ success: boolean; error?: string }> onUninstall: (extId: string) => void onRepaired: () => void | Promise onSynced: () => void @@ -34,13 +37,15 @@ interface Props { export function ExtensionDrawer({ ext, installedIds, localDataIds, downloading, loadError, disabled, onInstall, onInstallAll, onPauseDownload, onCancelDownload, - onUninstallNode, onUninstall, onRepaired, onSynced, onClose, + onUninstallNode, onDeleteWeightVariant, onUninstall, onRepaired, onSynced, onClose, }: Props): JSX.Element { const navigate = useNavStore((s) => s.navigate) + const installedWeightVariants = useExtensionsStore((s) => s.installedWeightVariants) const [repairing, setRepairing] = useState(false) const [repairError, setRepairError] = useState(null) const [syncing, setSyncing] = useState(false) const [syncError, setSyncError] = useState(null) + const [variantError, setVariantError] = useState(null) const isModel = ext.type === 'model' // Built-ins are corrupted-flagged too (builtin-sync repairs them on restart), @@ -85,7 +90,13 @@ export function ExtensionDrawer({ } } - const error = syncError ?? repairError ?? loadError + async function handleDeleteWeightVariant(fullId: string, variantId: string) { + setVariantError(null) + const result = await onDeleteWeightVariant(fullId, variantId) + if (!result.success) setVariantError(result.error ?? 'Could not remove these weights') + } + + const error = variantError ?? syncError ?? repairError ?? loadError return ( <> @@ -178,12 +189,19 @@ export function ExtensionDrawer({ const fullId = `${ext.id}/${node.id}` const state = getNodeState(ext.id, node, installedIds, downloading) const dl = state.kind === 'downloading' ? state.dl : null + const variants = isModel ? node.weightVariants : undefined + // A download started without a variant (Install all) fetches the default one. + const dlVariantId = dl && variants ? (dl.variantId ?? variants.default) : undefined + const dlVariant = variants?.options.find((option) => option.id === dlVariantId) + const installedLabels = variants?.options + .filter((option) => installedWeightVariants[fullId]?.includes(option.id)) + .map((option) => option.label) ?? [] const sub = state.kind === 'ready' ? 'Available on the node graph' - : state.kind === 'installed' ? 'Installed' + : state.kind === 'installed' ? (installedLabels.length > 0 ? `${installedLabels.join(', ')} installed` : 'Installed') : state.kind === 'available' ? 'Not installed' : dl?.paused ? 'Download paused' - : `Downloading… ${dl?.percent ?? 0}%` + : `Downloading${dlVariant ? ` ${dlVariant.label}` : ''}… ${dl?.percent ?? 0}%` return (
@@ -196,14 +214,16 @@ export function ExtensionDrawer({ {isModel && (
- onInstall(node, fullId)} - onPause={() => onPauseDownload(fullId)} - onResume={() => onInstall(node, fullId)} - onCancel={() => onCancelDownload(fullId)} - /> + {!variants && ( + onInstall(node, fullId)} + onPause={() => onPauseDownload(fullId)} + onResume={() => onInstall(node, fullId)} + onCancel={() => onCancelDownload(fullId)} + /> + )} {localDataIds.includes(fullId) && state.kind !== 'downloading' && (
+ {/* Weight variants — each one installs and deletes on its own */} + {variants && ( +
+ {variants.options.map((option) => { + const installed = installedWeightVariants[fullId]?.includes(option.id) ?? false + const variantDl = dlVariantId === option.id ? dl : null + const variantState: NodeUiState = variantDl + ? { kind: 'downloading', dl: variantDl } + : installed ? { kind: 'installed' } : { kind: 'available' } + return ( +
+
+ {option.label} + {option.sizeGb !== undefined && ( + {option.sizeGb} GB + )} + {option.vramGb !== undefined && ( + ~{option.vramGb} GB VRAM + )} + {option.id === variants.default && ( + default + )} +
+
+ onInstall(node, fullId, option.id)} + onPause={() => onPauseDownload(fullId)} + onResume={() => onInstall(node, fullId, option.id)} + onCancel={() => onCancelDownload(fullId)} + /> + {installed && !dl && ( + + )} +
+
+ ) + })} +
+ )} + {/* Download detail */} {dl && (
diff --git a/src/areas/models/components/extensionShared.tsx b/src/areas/models/components/extensionShared.tsx index 8bf865e3..e5682c55 100644 --- a/src/areas/models/components/extensionShared.tsx +++ b/src/areas/models/components/extensionShared.tsx @@ -12,6 +12,7 @@ export interface DownloadInfo { totalBytes?: number stalledSeconds?: number paused?: boolean + variantId?: string // set when the download targets one weight variant of the node } export type DownloadMap = Record diff --git a/src/areas/workflows/mockExtensions.ts b/src/areas/workflows/mockExtensions.ts index 2bbc8cc6..1eab87bb 100644 --- a/src/areas/workflows/mockExtensions.ts +++ b/src/areas/workflows/mockExtensions.ts @@ -1,6 +1,6 @@ import type { ModelExtension, ProcessExtension } from '@shared/stores/extensionsStore' export type { ParamSchema } from '@shared/types/electron.d' -import type { ParamSchema } from '@shared/types/electron.d' +import type { ParamSchema, WeightVariantsInfo } from '@shared/types/electron.d' export interface WorkflowExtension { id: string // "ext_id/node_id" @@ -17,6 +17,7 @@ export interface WorkflowExtension { params: ParamSchema[] builtin: boolean type: 'model' | 'process' + weightVariants?: WeightVariantsInfo } function applyParamDefaults( @@ -77,6 +78,7 @@ export function buildAllWorkflowExtensions( params: applyParamDefaults(node.paramsSchema as ParamSchema[], node.paramDefaults), builtin: ext.builtin, type: 'model', + weightVariants: node.weightVariants, }) } } diff --git a/src/areas/workflows/nodes/ExtensionNode.tsx b/src/areas/workflows/nodes/ExtensionNode.tsx index abe2f9f6..ec5f83d9 100644 --- a/src/areas/workflows/nodes/ExtensionNode.tsx +++ b/src/areas/workflows/nodes/ExtensionNode.tsx @@ -1,11 +1,13 @@ import { useCallback, useEffect, useRef, useLayoutEffect, useState } from 'react' import { Handle, Position, useReactFlow } from '@xyflow/react' import { useExtensionsStore } from '@shared/stores/extensionsStore' +import { useNavStore } from '@shared/stores/navStore' import { buildAllWorkflowExtensions } from '../mockExtensions' import type { ParamSchema } from '../mockExtensions' import type { WFNodeData } from '@shared/types/electron.d' import { PICKER_LABELS, openParamPicker, resolvePickerIntent } from '@shared/utils/paramPicker' import { PickerIcon } from '@shared/components/ui' +import { isMissingWeightVariant, withWeightVariantAvailability } from '@shared/utils/weightVariants' import { useWorkflowRunStore } from '../workflowRunStore' import BaseNode from './BaseNode' @@ -167,8 +169,10 @@ export default function ExtensionNode({ id, data, selected }: { id: string; data const [handleTops, setHandleTops] = useState([]) const { modelExtensions, processExtensions } = useExtensionsStore() + const installedVariants = useExtensionsStore((s) => (data.extensionId ? s.installedWeightVariants[data.extensionId] : undefined)) const allExtensions = buildAllWorkflowExtensions(modelExtensions, processExtensions) const ext = allExtensions.find((e) => e.id === data.extensionId) + const openExtension = useNavStore((s) => s.openExtension) const inputs = ext?.inputs // defined → multi-input mode const isMulti = inputs && inputs.length > 1 @@ -313,7 +317,15 @@ export default function ExtensionNode({ id, data, selected }: { id: string; data
- patchParam(param.id, v)} resolvedParams={resolvedParams} /> + { + patchParam(param.id, v) + if (ext && isMissingWeightVariant(param.id, v, ext.weightVariants, installedVariants)) openExtension(ext.extensionId) + }} + resolvedParams={resolvedParams} + />
) diff --git a/src/shared/stores/extensionsStore.ts b/src/shared/stores/extensionsStore.ts index 91657a0a..ada611f2 100644 --- a/src/shared/stores/extensionsStore.ts +++ b/src/shared/stores/extensionsStore.ts @@ -24,8 +24,11 @@ interface ExtensionsStore { installProgress: InstallProgress | null installError: string | null loadErrors: Record + /** Installed weight variant ids, keyed by "ext_id/node_id", for nodes that declare variants */ + installedWeightVariants: Record loadExtensions: () => Promise + refreshInstalledWeightVariants: () => Promise installFromGitHub: (url: string) => Promise<{ success: boolean; error?: string }> installFromLocal: () => Promise<{ success: boolean; error?: string; cancelled?: boolean; needsRepair?: boolean }> uninstall: (extensionId: string) => Promise<{ success: boolean; error?: string }> @@ -54,6 +57,7 @@ export const useExtensionsStore = create((set, get) => ({ installProgress: null, installError: null, loadErrors: {}, + installedWeightVariants: {}, // ── Load list ────────────────────────────────────────────────────────────── @@ -66,11 +70,30 @@ export const useExtensionsStore = create((set, get) => ({ ...extensions, loading: false, }) + await get().refreshInstalledWeightVariants() } catch { set({ loading: false }) } }, + async refreshInstalledWeightVariants() { + const entries = await Promise.all( + get().modelExtensions.flatMap((ext) => ext.nodes + .filter((node) => node.weightVariants) + .map(async (node) => { + const fullId = `${ext.id}/${node.id}` + return [fullId, await window.electron.model.installedWeightVariants(fullId)] as const + })), + ) + // A node whose state could not be read stays absent: undefined reads as "unknown", + // which the UI keeps neutral, while [] would claim no variant is installed. + set({ + installedWeightVariants: Object.fromEntries( + entries.filter((entry): entry is readonly [string, string[]] => entry[1] !== null), + ), + }) + }, + // ── Install from GitHub ──────────────────────────────────────────────────── async installFromGitHub(url: string) { diff --git a/src/shared/stores/navStore.ts b/src/shared/stores/navStore.ts index c6854424..62afa2f5 100644 --- a/src/shared/stores/navStore.ts +++ b/src/shared/stores/navStore.ts @@ -4,10 +4,16 @@ export type Page = 'generate' | 'workflows' | 'models' | 'settings' interface NavState { currentPage: Page + extensionToOpen: string | null navigate: (page: Page) => void + openExtension: (extensionId: string) => void + clearExtensionToOpen: () => void } export const useNavStore = create((set) => ({ currentPage: 'generate', - navigate: (page) => set({ currentPage: page }) + extensionToOpen: null, + navigate: (page) => set({ currentPage: page }), + openExtension: (extensionId) => set({ currentPage: 'models', extensionToOpen: extensionId }), + clearExtensionToOpen: () => set({ extensionToOpen: null }), })) diff --git a/src/shared/types/electron.d.ts b/src/shared/types/electron.d.ts index 840c5e76..26f9bf3b 100644 --- a/src/shared/types/electron.d.ts +++ b/src/shared/types/electron.d.ts @@ -25,6 +25,13 @@ export interface ExtensionNode { hfSkipPrefixes?: string[] hfIncludePrefixes?: string[] hasModelSources?: boolean + weightVariants?: WeightVariantsInfo +} + +export interface WeightVariantsInfo { + param: string // params_schema id whose value selects the variant + default: string + options: { id: string; label: string; sizeGb?: number; vramGb?: number }[] } export interface ModelExtension { @@ -207,17 +214,21 @@ declare global { model: { export: (args: { outputUrl: string; format: string }) => Promise<{ success: boolean; error?: string }> listDownloaded: () => Promise<{ id: string; name: string; size_gb: number }[]> - activeDownloads: () => Promise<{ modelId: string; percent: number; file?: string; fileIndex?: number; totalFiles?: number }[]> + activeDownloads: () => Promise<{ modelId: string; variantId?: string; percent: number; file?: string; fileIndex?: number; totalFiles?: number }[]> isDownloaded: (modelId: string) => Promise hasLocalData: (modelId: string) => Promise - download: (modelId: string) => Promise<{ success: boolean; error?: string; paused?: boolean; cancelled?: boolean }> + download: (modelId: string, variantId?: string) => Promise<{ success: boolean; error?: string; paused?: boolean; cancelled?: boolean }> pauseDownload: (modelId: string) => Promise<{ success: boolean; error?: string }> cancelDownload: (modelId: string) => Promise<{ success: boolean; error?: string }> delete: (modelId: string) => Promise<{ success: boolean; error?: string }> + /** Installed variant ids, or null when the node's install state could not be read */ + installedWeightVariants: (modelId: string) => Promise + deleteWeightVariant: (modelId: string, variantId: string) => Promise<{ success: boolean; error?: string }> unloadAll: () => Promise<{ success: boolean; error?: string }> showInFolder: (modelId: string) => Promise onProgress: (cb: (data: { modelId: string + variantId?: string percent: number file?: string fileIndex?: number diff --git a/src/shared/utils/weightVariants.test.mjs b/src/shared/utils/weightVariants.test.mjs new file mode 100644 index 00000000..1659391b --- /dev/null +++ b/src/shared/utils/weightVariants.test.mjs @@ -0,0 +1,63 @@ +import test from 'node:test' +import assert from 'node:assert/strict' +import { buildSync } from 'esbuild' +import { createRequire } from 'node:module' +import { mkdtempSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' + +function loadModule() { + const outfile = join(mkdtempSync(join(tmpdir(), 'modly-weight-variants-test-')), 'weightVariants.cjs') + const require = createRequire(import.meta.url) + const result = buildSync({ + entryPoints: [resolve('src/shared/utils/weightVariants.ts')], + bundle: true, + platform: 'node', + format: 'cjs', + write: false, + }) + writeFileSync(outfile, result.outputFiles[0].text, 'utf8') + return require(outfile) +} + +const { withWeightVariantAvailability, isMissingWeightVariant } = loadModule() + +const variants = { + param: 'gguf_quant', + default: 'Q5_K_M', + options: [{ id: 'Q4_K_M', label: 'Q4_K_M' }, { id: 'Q5_K_M', label: 'Q5_K_M' }], +} + +const quantParam = { + id: 'gguf_quant', + label: 'Quantization', + type: 'select', + default: 'Q5_K_M', + options: [ + { value: 'Q4_K_M', label: 'Q4_K_M' }, + { value: 'Q5_K_M', label: 'Q5_K_M' }, + { value: 'auto', label: 'Auto' }, + ], +} + +test('labels declared variants that are not installed and keeps their values', () => { + const marked = withWeightVariantAvailability(quantParam, variants, ['Q5_K_M']) + assert.deepEqual(marked.options.map((option) => option.label), ['Q4_K_M (not installed)', 'Q5_K_M', 'Auto']) + assert.deepEqual(marked.options.map((option) => option.value), ['Q4_K_M', 'Q5_K_M', 'auto']) +}) + +test('returns the param untouched when availability is unknown or the param selects nothing', () => { + assert.equal(withWeightVariantAvailability(quantParam, variants, undefined), quantParam) + assert.equal(withWeightVariantAvailability(quantParam, undefined, []), quantParam) + const steps = { id: 'steps', label: 'Steps', type: 'int', default: 25 } + assert.equal(withWeightVariantAvailability(steps, variants, []), steps) +}) + +test('flags a selected variant only when availability is known and it is not installed', () => { + assert.equal(isMissingWeightVariant('gguf_quant', 'Q4_K_M', variants, ['Q5_K_M']), true) + assert.equal(isMissingWeightVariant('gguf_quant', 'Q5_K_M', variants, ['Q5_K_M']), false) + assert.equal(isMissingWeightVariant('gguf_quant', 'auto', variants, []), false) + assert.equal(isMissingWeightVariant('steps', 'Q4_K_M', variants, []), false) + assert.equal(isMissingWeightVariant('gguf_quant', 'Q4_K_M', variants, undefined), false) + assert.equal(isMissingWeightVariant('gguf_quant', 'Q4_K_M', undefined, []), false) +}) diff --git a/src/shared/utils/weightVariants.ts b/src/shared/utils/weightVariants.ts new file mode 100644 index 00000000..252fd6d5 --- /dev/null +++ b/src/shared/utils/weightVariants.ts @@ -0,0 +1,35 @@ +import type { ParamSchema, WeightVariantsInfo } from '@shared/types/electron.d' + +/** + * Suffix the options of a node's variant-selecting param that are declared + * weight variants but not installed. Other params, and values that are not + * variant ids, are returned untouched. `installed` undefined means not known yet. + */ +export function withWeightVariantAvailability( + param: ParamSchema, + variants: WeightVariantsInfo | undefined, + installed: string[] | undefined, +): ParamSchema { + if (!variants || !installed || param.id !== variants.param || !param.options) return param + const variantIds = new Set(variants.options.map((option) => option.id)) + return { + ...param, + options: param.options.map((option) => { + const value = String(option.value) + if (!variantIds.has(value) || installed.includes(value)) return option + return { ...option, label: `${option.label ?? value} (not installed)` } + }), + } +} + +/** True when `value` selects a declared weight variant that is known not to be installed. */ +export function isMissingWeightVariant( + paramId: string, + value: unknown, + variants: WeightVariantsInfo | undefined, + installed: string[] | undefined, +): boolean { + if (!variants || !installed || paramId !== variants.param) return false + const id = String(value) + return variants.options.some((option) => option.id === id) && !installed.includes(id) +}