feat(models): support extension-scoped shared weight groups

.github/workflows/model-weights.yml

@@ -0,0 +1,40 @@
+name: Model weight regressions
+
+on:
+  pull_request:
+    branches: [dev, main]
+    paths:
+      - 'api/**'
+      - 'electron/**'
+      - 'src/areas/models/**'
+      - 'src/shared/types/electron.d.ts'
+      - 'package*.json'
+      - '.github/workflows/model-weights.yml'
+
+permissions:
+  contents: read
+
+jobs:
+  model-weights:
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, windows-latest]
+        python: ['3.11', '3.12']
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-node@v4
+        with:
+          node-version: '22'
+          cache: npm
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python }}
+      - run: npm ci --ignore-scripts --no-audit --no-fund
+      - run: python -m pip install fastapi httpx
+      - name: Download, deletion, manifests and install queue
+        run: node --test electron/main/model-weight-ipc.test.mjs electron/main/model-sources.test.mjs electron/main/model-download-plan.test.mjs electron/main/model-download-preload.test.mjs electron/main/extension-install-utils.test.mjs src/areas/models/utils.test.mjs
+      - name: Python runtime and download regressions
+        working-directory: api
+        run: python -m unittest tests.test_model_sources tests.test_model_router tests.test_generator_registry tests.test_extension_process tests.test_runner

README.md

@@ -148,6 +148,65 @@ supported provider is `huggingface`. Existing nodes that use `hf_repo`,
 `download_check`, `hf_include_prefixes`, and `hf_skip_prefixes` keep their
 original behavior.
 
+### Shared weights inside one model extension
+
+Multi-node model extensions can declare extension-scoped `weight_groups` and
+reference them from any sibling node. Shared files are downloaded once under
+`<models-dir>/<extension-id>/_shared/<group-id>`, while node-specific
+`model_sources` stay under the node's existing model directory.
+
+```json
+{
+  "id": "pixal3d",
+  "type": "model",
+  "weight_groups": [
+    {
+      "id": "pixal3d-base",
+      "model_sources": [
+        {
+          "id": "base",
+          "provider": "huggingface",
+          "repo_id": "TencentARC/Pixal3D",
+          "revision": "<pinned-revision>",
+          "destination": ".",
+          "checks": ["pipeline.json"]
+        }
+      ]
+    }
+  ],
+  "nodes": [
+    {
+      "id": "generate",
+      "weight_groups": ["pixal3d-base"]
+    },
+    {
+      "id": "worldsculpt",
+      "weight_groups": ["pixal3d-base"],
+      "model_sources": [
+        {
+          "id": "adapter",
+          "provider": "huggingface",
+          "repo_id": "AlayaLab/WorldSculpt",
+          "revision": "<pinned-revision>",
+          "destination": ".",
+          "checks": ["model.safetensors"]
+        }
+      ]
+    }
+  ]
+}
+```
+
+At runtime, `MODEL_DIR` remains the selected node's private directory.
+Subprocess extensions also receive `MODEL_ID`, `MODEL_NODE_ID`, and a JSON
+`SHARED_MODEL_DIRS` map in their environment. Both direct and subprocess generator
+instances receive `MODEL_ID`, `MODEL_NODE_ID`, and the resolved mapping in
+`shared_model_dirs` before `load()`. Direct generators use these instance attributes,
+not process-global environment variables, to distinguish sibling nodes.
+Shared groups are installed through their dependent nodes; the drawer exposes
+shared-group status and explicit removal. Removing private node data never removes a shared group;
+shared-group removal is a separate action that identifies every affected node.
+
 ---
 
 ## Workflows

api/routers/model.py

@@ -10,12 +10,14 @@
 from urllib.request import Request, urlopen
 from fastapi import APIRouter, HTTPException, Request as FastAPIRequest
 from fastapi.responses import StreamingResponse
-from services.generator_registry import generator_registry, MODELS_DIR
+from services.generator_registry import generator_registry
+import services.generator_registry as registry_module
+from services.extension_process import ExtensionProcess
 from services.model_sources import (
     normalize_model_sources,
     resolve_download_path,
-    resolve_model_root,
-    resolve_source_destination,
+    resolve_source_destination_at_root,
+    resolve_weight_storage_root,
     validate_source_file_plan,
 )
 
@@ -109,10 +111,19 @@ async def unload_model(model_id: str):
     """Unloads a model from memory so its files can be safely deleted."""
     try:
         gen = generator_registry.get_generator(model_id)
+    except ValueError as exc:
+        if model_id in generator_registry._generators:
+            raise HTTPException(409, str(exc)) from exc
+        return {"unloaded": True}  # No runtime registered for these files.
+    # unload() on ExtensionProcess deliberately swallows IPC errors; deletion
+    # needs a confirmed process exit so no worker can retain file handles.
+    if isinstance(gen, ExtensionProcess):
+        gen.stop()
+    else:
         gen.unload()
-        return {"unloaded": True}
-    except ValueError:
-        return {"unloaded": True}  # already not loaded, that's fine
+        if gen.is_loaded():
+            raise HTTPException(409, "Model is still loaded; weights were preserved")
+    return {"unloaded": True}
 
 
 @router.post("/hf-download/pause")
@@ -131,7 +142,7 @@ async def cancel_hf_download(model_id: str):
 
 @router.post("/hf-download-sources")
 async def hf_download_sources(request: FastAPIRequest, model_id: str):
-    """Download all Hugging Face sources declared for one model node."""
+    """Download sources into one validated node or extension-shared target."""
     try:
         body = await request.json()
         if not isinstance(body, dict):
@@ -140,10 +151,10 @@ async def hf_download_sources(request: FastAPIRequest, model_id: str):
         if raw_sources is None:
             raise ValueError("sources are required")
         sources = normalize_model_sources({"model_sources": raw_sources})
-        model_root = resolve_model_root(MODELS_DIR, model_id)
+        model_root = resolve_weight_storage_root(registry_module.MODELS_DIR, model_id)
         destinations = {
-            source["id"]: resolve_source_destination(
-                MODELS_DIR, model_id, source["destination"]
+            source["id"]: resolve_source_destination_at_root(
+                model_root, source["destination"]
             )
             for source in sources
         }
@@ -305,7 +316,7 @@ async def hf_download(
     """
     import json as _json
     import os
-    dest_dir  = str(MODELS_DIR / model_id)
+    dest_dir  = str(registry_module.MODELS_DIR / model_id)
     # Prefer skip_prefixes passed directly from the client (authoritative, no registry dep)
     if skip_prefixes:
         try:

api/runner.py

@@ -34,6 +34,15 @@
 # MODEL_DIR is set by ExtensionProcess to match its own model_dir (composite node id path).
 # Falls back to MODELS_DIR/manifest_id for standalone/legacy use.
 _MODEL_DIR_OVERRIDE = os.environ.get("MODEL_DIR", "")
+_MODEL_ID_OVERRIDE = os.environ.get("MODEL_ID", "")
+_MODEL_NODE_ID_OVERRIDE = os.environ.get("MODEL_NODE_ID", "")
+try:
+    _SHARED_MODEL_DIRS = {
+        str(group_id): Path(path)
+        for group_id, path in json.loads(os.environ.get("SHARED_MODEL_DIRS", "{}")).items()
+    }
+except (AttributeError, TypeError, ValueError, json.JSONDecodeError):
+    _SHARED_MODEL_DIRS = {}
 
 # Inject Modly's api/ so generator.py can do:
 #   from services.generators.base import BaseGenerator, ...
@@ -87,8 +96,12 @@ def load_generator(manifest: dict):
     return getattr(mod, manifest["generator_class"])
 
 
-def _select_node(manifest: dict, model_dir_override: str) -> dict:
+def _select_node(
+    manifest: dict, model_dir_override: str, node_id_override: str = ""
+) -> dict:
     nodes = manifest.get("nodes") or []
+    if nodes and node_id_override:
+        return next((n for n in nodes if n.get("id") == node_id_override), nodes[0])
     if nodes and model_dir_override:
         node_id = Path(model_dir_override).name
         return next((n for n in nodes if n.get("id") == node_id), nodes[0])
@@ -157,7 +170,7 @@ def _apply_manifest_metadata(gen, manifest: dict, node: dict) -> None:
 
 def main() -> None:
     manifest = json.loads((EXT_DIR / "manifest.json").read_text(encoding="utf-8"))
-    model_id = manifest["id"]
+    model_id = _MODEL_ID_OVERRIDE or manifest["id"]
 
     try:
         GenClass = load_generator(manifest)
@@ -167,11 +180,9 @@ def main() -> None:
               "traceback": traceback.format_exc()})
         return
 
-    # Support both flat manifest (legacy) and nodes[] format.
-    # Use MODEL_DIR to find the correct node for multi-node extensions:
-    # MODEL_DIR is set by ExtensionProcess to MODELS_DIR/ext_id/node_id,
-    # so its last component matches the node id.
-    node = _select_node(manifest, _MODEL_DIR_OVERRIDE)
+    # Support both flat manifest (legacy) and nodes[] format. The host passes an
+    # explicit node id; MODEL_DIR name inference remains only as a legacy fallback.
+    node = _select_node(manifest, _MODEL_DIR_OVERRIDE, _MODEL_NODE_ID_OVERRIDE)
 
     # Announce readiness and send params_schema so ExtensionProcess
     # can serve it without needing to query the subprocess later.
@@ -184,6 +195,9 @@ def main() -> None:
     # Falls back to MODELS_DIR/manifest_id for legacy / standalone use.
     model_dir = Path(_MODEL_DIR_OVERRIDE) if _MODEL_DIR_OVERRIDE else MODELS_DIR / model_id
     gen = GenClass(model_dir, WORKSPACE_DIR)
+    gen.MODEL_ID = model_id
+    gen.MODEL_NODE_ID = node.get("id", "")
+    gen.shared_model_dirs = dict(_SHARED_MODEL_DIRS)
     _apply_manifest_metadata(gen, manifest, node)
 
     # Active cancel events keyed by request id

api/services/extension_process.py

@@ -45,6 +45,7 @@ def __init__(self, ext_dir: Path, manifest: dict) -> None:
         self.manifest      = manifest
         self.model_dir     = None   # set by registry after init
         self.outputs_dir   = None   # set by registry after init
+        self.shared_model_dirs: dict[str, Path] = {}
 
         self._proc:   Optional[subprocess.Popen] = None
         self._queue:  queue.Queue                = queue.Queue()
@@ -84,11 +85,17 @@ def _build_env(self) -> dict:
             # Setting it inside generator.py is too late, since generator.py
             # itself imports torch before calling select_device().
             env.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
-        # Pass the exact model_dir so runner.py doesn't have to re-derive it
-        # from manifest["id"] (which is the ext_id, not the composite node id).
-        # runner.py extracts the node id from MODEL_DIR's trailing path component.
+        # Keep capability identity separate from storage identity. MODEL_DIR
+        # retains its node-private meaning; shared roots are passed explicitly.
         if self.model_dir is not None:
             env["MODEL_DIR"] = str(self.model_dir)
+        env["MODEL_ID"] = self.MODEL_ID
+        env["MODEL_NODE_ID"] = self.manifest.get(
+            "node_id", self.MODEL_ID.split("/", 1)[-1]
+        )
+        env["SHARED_MODEL_DIRS"] = json.dumps(
+            {group_id: str(path) for group_id, path in self.shared_model_dirs.items()}
+        )
         # Extension venvs are based on python-embed which ships without a CA bundle.
         # Only set SSL_CERT_FILE if not already provided (preserves corporate/custom certs).
         if "SSL_CERT_FILE" not in env: