Studio diffusion (Phase 9) review fixes: prequant safety + validation

- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
  resolves inside an operator-configured ALLOWLIST of directories
  (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
  once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
  any path a load request named (arbitrary code execution). realpath() blocks symlink
  escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
  checkpoint that quantised a different layer set is rejected instead of silently
  loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
  path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
  author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.
This commit is contained in:
Daniel Han 2026-06-29 05:27:19 +00:00
commit 44770ce2d9
5 changed files with 177 additions and 31 deletions

View file

@ -19,17 +19,20 @@ from __future__ import annotations
import argparse
import logging
import os
import sys
import time
from pathlib import Path
import numpy as np
BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend"
_REPO = Path(__file__).resolve().parent.parent
_RESEARCH = _REPO / "outputs" / "quant_research"
BACKEND = _REPO / "studio" / "backend"
BASE = "Tongyi-MAI/Z-Image-Turbo"
CKPT = "/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/prequant_fp8/transformer_fp8.pt"
CKPT = os.environ.get("PREQUANT_CKPT", str(_RESEARCH / "prequant_fp8" / "transformer_fp8.pt"))
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/prequant_verify_images")
OUT = Path(os.environ.get("PREQUANT_OUT_DIR", str(_RESEARCH / "prequant_verify_images")))
logging.basicConfig(level = logging.INFO, format = "%(message)s")
LOGGER = logging.getLogger("verify_prequant")