feat(models): install and remove weight variants per node

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

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.

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:

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

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: