Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in

load_prequantized_transformer ends in torch.load(weights_only=False), which executes
arbitrary code from the pickle. The transformer_prequant_path load-request field reached
that unpickle for any local file an authenticated caller named, so a request could trigger
remote code execution. Refuse the source.kind=='path' branch unless the operator sets
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted
and unaffected. Document the requirement on the API field and add gate tests.

studio/backend/core/inference/diffusion_prequant.py

@@ -33,6 +33,26 @@ from typing import Any, Optional
 # on-disk structure changes so an old/foreign artifact is rejected rather than mis-loaded.
 PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1"
 
+# Loading a checkpoint ends in ``torch.load(weights_only=False)``, which executes arbitrary
+# code embedded in the pickle. A hosted family *repo* checkpoint is first-party and trusted,
+# but a ``source.kind == "path"`` can originate from the ``transformer_prequant_path`` field
+# of a load request -- i.e. an authenticated API caller naming an arbitrary local file.
+# Unpickling that is remote code execution, so the local-path branch is refused unless an
+# operator explicitly opts in via this env var. The trusted hosted-repo path is unaffected.
+ALLOW_LOCAL_PREQUANT_PATH_ENV = "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH"
+
+
+def _local_prequant_path_allowed() -> bool:
+    """Whether a request-supplied local pre-quant *path* may be unpickled (operator opt-in)."""
+    import os
+
+    return (os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV) or "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+        "on",
+    }
+
 
 @dataclass(frozen = True)
 class PrequantSource:
@@ -93,6 +113,21 @@ def load_prequantized_transformer(
     ordinary unavailable artifact.
     """
     try:
+        # weights_only=False (required below) executes pickle code, so a caller-supplied
+        # local path is refused unless an operator opted in. The hosted family repo is
+        # first-party and always allowed.
+        if source.kind == "path" and not _local_prequant_path_allowed():
+            _warn(
+                logger,
+                f"{scheme}:path",
+                RuntimeError(
+                    "request-supplied local pre-quant path refused (unpickling an "
+                    f"arbitrary file is unsafe); set {ALLOW_LOCAL_PREQUANT_PATH_ENV}=1 "
+                    "to allow trusted local checkpoints",
+                ),
+            )
+            return None
+
         path = _resolve_checkpoint_path(source, hf_token)
         if path is None:
             return None
@@ -100,9 +135,8 @@ def load_prequantized_transformer(
         import torch
 
         # torchao weight subclasses are not safetensors-serializable, so the checkpoint is
-        # a torch.save pickle. weights_only=False is required to rebuild those subclasses;
-        # only a configured family repo (first-party) or an explicit local path reaches
-        # here, which is the trust signal -- this never loads an arbitrary remote pickle.
+        # a torch.save pickle. weights_only=False is required to rebuild those subclasses.
+        # The local-path branch is gated above; the repo branch is a first-party artifact.
         ckpt = torch.load(path, weights_only = False, map_location = "cpu")
         if not _validate_checkpoint(ckpt, scheme, base, logger):
             return None

studio/backend/models/inference.py

@@ -1744,7 +1744,9 @@ class DiffusionLoadRequest(BaseModel):
         "scripts/build_prequant_checkpoint.py) for the requested transformer_quant "
         "scheme. Loads the already-quantized weights with the dense bf16 never on the "
         "GPU (~half the load VRAM and a smaller download). null uses the family's hosted "
-        "checkpoint if configured, else quantises the dense transformer at load time.",
+        "checkpoint if configured, else quantises the dense transformer at load time. "
+        "Loading a local path unpickles the file (arbitrary code execution), so it is "
+        "ignored unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1.",
     )
 
 

studio/backend/tests/test_diffusion_prequant.py

@@ -143,9 +143,16 @@ def _load(
     scheme = "fp8",
     load_raises = False,
     exists = True,
+    allow_local = True,
 ):
     _FakeTransformer.calls = {}
     _stub_torch_accelerate(monkeypatch, ckpt, load_raises = load_raises)
+    # The local-path branch is opt-in (it unpickles an arbitrary file); these tests
+    # exercise the load mechanics, so enable it unless a test is checking the gate.
+    if allow_local:
+        monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, "1")
+    else:
+        monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False)
     path = tmp_path / "ckpt.pt"
     if exists:
         path.write_bytes(b"x")
@@ -195,3 +202,65 @@ def test_load_scheme_mismatch_is_none(monkeypatch, tmp_path):
 
 def test_load_base_mismatch_is_none(monkeypatch, tmp_path):
     assert _load(monkeypatch, tmp_path, _good_ckpt(base = "other/model")) is None
+
+
+# ── local-path opt-in gate (RCE guard) ───────────────────────────────────────────
+def test_load_local_path_refused_by_default(monkeypatch, tmp_path):
+    # A valid checkpoint at a real file is still refused: torch.load must never run on a
+    # request-supplied path without the operator opt-in.
+    called = {"load": False}
+
+    def _explode(*a, **k):
+        called["load"] = True
+        raise AssertionError("torch.load must not run on a refused local path")
+
+    torch = types.ModuleType("torch")
+    torch.load = _explode
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False)
+
+    path = tmp_path / "ckpt.pt"
+    path.write_bytes(b"x")
+    source = PrequantSource(kind = "path", location = str(path), filename = None)
+    result = load_prequantized_transformer(
+        _FakeTransformer,
+        "Tongyi-MAI/Z-Image-Turbo",
+        source,
+        device = "cuda",
+        dtype = "bfloat16",
+        hf_token = None,
+        scheme = "fp8",
+        logger = None,
+    )
+    assert result is None
+    assert called["load"] is False
+
+
+def test_load_local_path_allowed_with_optin(monkeypatch, tmp_path):
+    assert _load(monkeypatch, tmp_path, _good_ckpt(), allow_local = True) is not None
+
+
+def test_load_repo_source_allowed_without_optin(monkeypatch, tmp_path):
+    # The hosted-repo branch is first-party and trusted: it loads with no opt-in env set.
+    _FakeTransformer.calls = {}
+    _stub_torch_accelerate(monkeypatch, _good_ckpt())
+    monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False)
+
+    downloaded = tmp_path / "transformer_fp8.pt"
+    downloaded.write_bytes(b"x")
+    hub = types.ModuleType("huggingface_hub")
+    hub.hf_hub_download = lambda repo_id, filename, token = None: str(downloaded)
+    monkeypatch.setitem(sys.modules, "huggingface_hub", hub)
+
+    source = PrequantSource(kind = "repo", location = "org/hosted-fp8", filename = "transformer_fp8.pt")
+    result = load_prequantized_transformer(
+        _FakeTransformer,
+        "Tongyi-MAI/Z-Image-Turbo",
+        source,
+        device = "cuda",
+        dtype = "bfloat16",
+        hf_token = None,
+        scheme = "fp8",
+        logger = None,
+    )
+    assert result is not None