Address the round of Codex review findings on the merged diffusion phases
Memory planning and dense-quant path: size a local diffusers base's resident companions from its on-disk VAE and text-encoder weights instead of folding them to zero, feed the distilled variant hint into the runtime headroom estimate so turbo and schnell models are not over-reserved, place group-offload companions resident before attaching the transformer hooks so a failed placement falls back to whole-module offload instead of crashing, and bail out of the dense transformer download before it starts when the requested quant scheme is unsupported so the load falls back to GGUF cleanly. sd.cpp stack: scrub the native path lease secret from sd-cli child env, redact native load-progress errors, forward the resolved accelerator when auto-installing a forced-native binary, release stale diffusion GPU ownership on CPU-native loads, and remove the sd.cpp install tree on uninstall. Prequant and scripts: reject prequant artifacts missing base_model_id when a base is requested, expanduser before checkpoint existence checks, record and validate the int8 exclusion filter and fp8 fast-accum in checkpoint metadata, make verify_prequant_backend allowlist its local checkpoint and fail on missing or bad LPIPS and on load-peak regressions, average only finite PSNR values in diffusion_quality, and reset the process-wide attention backend between perf probe variants. API and UI: normalize attention_backend casing before Literal validation, close hidden popovers when leaving the Images page, and clear the stale quant label when loading a direct local GGUF file.
This commit is contained in:
parent
22f49b5ac2
commit
a9e5a80654
21 changed files with 495 additions and 32 deletions
|
|
@ -37,6 +37,14 @@ OUT = Path(os.environ.get("PREQUANT_OUT_DIR", str(_RESEARCH / "prequant_verify_i
|
|||
logging.basicConfig(level = logging.INFO, format = "%(message)s")
|
||||
LOGGER = logging.getLogger("verify_prequant")
|
||||
|
||||
# Prequant and runtime produce the SAME quantized weights, so their images must be
|
||||
# near-identical; anything above this LPIPS means the prequant path diverged and the run
|
||||
# fails. The prequant load peak must also sit clearly below the dense runtime peak (the
|
||||
# whole point of the path); require at least this fractional headroom.
|
||||
LPIPS_MAX = 0.02
|
||||
PREQUANT_PEAK_MAX_FRACTION = 0.75
|
||||
_RUNTIME_PEAK_FILE = OUT / "runtime_peak.txt"
|
||||
|
||||
|
||||
def _target(dtype):
|
||||
import types
|
||||
|
|
@ -81,7 +89,11 @@ def run(mode, steps, seed, res):
|
|||
sys.path.insert(0, str(BACKEND))
|
||||
import torch
|
||||
import diffusers
|
||||
from core.inference.diffusion_prequant import PrequantSource, load_prequantized_transformer
|
||||
from core.inference.diffusion_prequant import (
|
||||
ALLOW_LOCAL_PREQUANT_PATH_ENV,
|
||||
PrequantSource,
|
||||
load_prequantized_transformer,
|
||||
)
|
||||
from core.inference.diffusion_transformer_quant import quantize_transformer
|
||||
|
||||
OUT.mkdir(parents = True, exist_ok = True)
|
||||
|
|
@ -90,6 +102,14 @@ def run(mode, steps, seed, res):
|
|||
torch.cuda.empty_cache()
|
||||
|
||||
if mode == "prequant":
|
||||
# A local checkpoint is refused unless its directory is allowlisted (unpickling an
|
||||
# arbitrary file is unsafe). This verifier's CKPT is operator-supplied and trusted,
|
||||
# so allowlist its directory here or the load returns None and measures nothing.
|
||||
ckpt_dir = os.path.dirname(os.path.realpath(CKPT))
|
||||
existing = os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV, "")
|
||||
os.environ[ALLOW_LOCAL_PREQUANT_PATH_ENV] = (
|
||||
ckpt_dir if not existing else existing + os.pathsep + ckpt_dir
|
||||
)
|
||||
source = PrequantSource(kind = "path", location = CKPT, filename = None)
|
||||
transformer = load_prequantized_transformer(
|
||||
transformer_cls,
|
||||
|
|
@ -122,17 +142,44 @@ def run(mode, steps, seed, res):
|
|||
scheme = quantize_transformer(pipe, _target(torch.bfloat16), mode = "fp8", logger = LOGGER)
|
||||
load_peak = torch.cuda.max_memory_allocated() / 1e9
|
||||
print(f"[runtime] engaged={scheme} load_gpu_peak={load_peak:.1f} GB", flush = True)
|
||||
# Persist the dense reference peak so a later prequant run can enforce its VRAM win.
|
||||
_RUNTIME_PEAK_FILE.write_text(f"{load_peak:.6f}")
|
||||
|
||||
img, dt = _gen(pipe, steps, seed, res) # warmup
|
||||
img, dt = _gen(pipe, steps, seed, res)
|
||||
img.save(OUT / f"{mode}.png")
|
||||
print(f"[{mode}] gen={dt:.3f}s saved {mode}.png", flush = True)
|
||||
|
||||
if mode != "prequant":
|
||||
return 0
|
||||
|
||||
# Enforce the two invariants this verifier exists to check, so a broken prequant
|
||||
# checkpoint fails loudly instead of passing just because generation completed.
|
||||
ref_path = OUT / "runtime.png"
|
||||
if mode == "prequant" and ref_path.exists():
|
||||
from PIL import Image
|
||||
lp = _lpips(np.array(Image.open(ref_path).convert("RGB")), np.array(img))
|
||||
print(f"[prequant] LPIPS_vs_runtime={lp}", flush = True)
|
||||
if not ref_path.exists():
|
||||
print("FAIL: runtime reference image missing; run --mode runtime first", flush = True)
|
||||
return 1
|
||||
from PIL import Image
|
||||
lp = _lpips(np.array(Image.open(ref_path).convert("RGB")), np.array(img))
|
||||
print(f"[prequant] LPIPS_vs_runtime={lp}", flush = True)
|
||||
if lp is None:
|
||||
print("FAIL: LPIPS could not be computed (install lpips)", flush = True)
|
||||
return 1
|
||||
if lp > LPIPS_MAX:
|
||||
print(f"FAIL: LPIPS {lp:.4f} > {LPIPS_MAX} (prequant diverged from runtime)", flush = True)
|
||||
return 1
|
||||
if _RUNTIME_PEAK_FILE.exists():
|
||||
try:
|
||||
runtime_peak = float(_RUNTIME_PEAK_FILE.read_text().strip())
|
||||
except ValueError:
|
||||
runtime_peak = 0.0
|
||||
if runtime_peak > 0.0 and load_peak > runtime_peak * PREQUANT_PEAK_MAX_FRACTION:
|
||||
print(
|
||||
f"FAIL: prequant load peak {load_peak:.1f} GB not below "
|
||||
f"{PREQUANT_PEAK_MAX_FRACTION:.0%} of dense {runtime_peak:.1f} GB",
|
||||
flush = True,
|
||||
)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue