From daad898d05ad16857d9110fbdfff9aee66a8628f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 28 Jun 2026 06:07:11 +0000 Subject: [PATCH] Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. --- scripts/quant_probe.py | 14 ++++++++------ .../core/inference/diffusion_transformer_quant.py | 4 +++- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/scripts/quant_probe.py b/scripts/quant_probe.py index 6d69e7e8e7..8675fb75bf 100644 --- a/scripts/quant_probe.py +++ b/scripts/quant_probe.py @@ -27,7 +27,7 @@ REPO = "unsloth/Z-Image-Turbo-GGUF" GGUF = "z-image-turbo-Q4_K_M.gguf" BASE = "Tongyi-MAI/Z-Image-Turbo" 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/probe_images") +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "probe_images" def _psnr(a, b): @@ -39,17 +39,19 @@ _LPIPS = {"fn": None} def _lpips(ref_arr, arr): - """Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable.""" + """Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable. + + Runs on CPU so the scorer never holds CUDA memory: each row resets peak VRAM, so a + resident GPU LPIPS module would inflate the reported load/gen VRAM and could even OOM.""" try: import torch import lpips if _LPIPS["fn"] is None: - _LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + _LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval() def t(x): - t = torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0 - return t.cuda() + return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0 with torch.no_grad(): return float(_LPIPS["fn"](t(ref_arr), t(arr)).item()) @@ -111,7 +113,7 @@ def _quant_config(name): return MXDynamicActivationMXWeightConfig( activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn ) - except TypeError: + except (TypeError, AttributeError): return MXDynamicActivationMXWeightConfig() raise ValueError(name) diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 1b211e691d..296c30576e 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -265,7 +265,9 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any: return MXDynamicActivationMXWeightConfig( activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn ) - except TypeError: + except (TypeError, AttributeError): + # TypeError: older torchao without the explicit dtype knobs. + # AttributeError: a torch build without torch.float8_e4m3fn. return MXDynamicActivationMXWeightConfig() raise ValueError(f"unknown transformer quant scheme '{scheme}'")