[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot]
parent
983f2c14a0
commit
3cfb29dd6b
5 changed files with 151 additions and 72 deletions
|
|
@ -43,8 +43,9 @@ def _lpips(ref_arr, arr):
|
|||
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).cuda().eval()
|
||||
|
||||
def t(x):
|
||||
t = torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
|
||||
|
|
@ -53,7 +54,7 @@ def _lpips(ref_arr, arr):
|
|||
with torch.no_grad():
|
||||
return float(_LPIPS["fn"](t(ref_arr), t(arr)).item())
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f" (lpips unavailable: {type(exc).__name__}: {str(exc)[:80]})", flush=True)
|
||||
print(f" (lpips unavailable: {type(exc).__name__}: {str(exc)[:80]})", flush = True)
|
||||
return None
|
||||
|
||||
|
||||
|
|
@ -62,9 +63,9 @@ def _load_dense():
|
|||
import diffusers
|
||||
|
||||
t = diffusers.ZImageTransformer2DModel.from_pretrained(
|
||||
BASE, subfolder="transformer", torch_dtype=torch.bfloat16
|
||||
BASE, subfolder = "transformer", torch_dtype = torch.bfloat16
|
||||
)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype=torch.bfloat16, transformer=t)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
|
@ -76,10 +77,12 @@ def _load_gguf():
|
|||
|
||||
t = diffusers.ZImageTransformer2DModel.from_single_file(
|
||||
hf_hub_download(REPO, GGUF),
|
||||
quantization_config=diffusers.GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
|
||||
torch_dtype=torch.bfloat16, config=BASE, subfolder="transformer",
|
||||
quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = torch.bfloat16),
|
||||
torch_dtype = torch.bfloat16,
|
||||
config = BASE,
|
||||
subfolder = "transformer",
|
||||
)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype=torch.bfloat16, transformer=t)
|
||||
pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
|
@ -91,6 +94,7 @@ def _quant_config(name):
|
|||
Int8DynamicActivationInt8WeightConfig,
|
||||
Float8DynamicActivationFloat8WeightConfig,
|
||||
)
|
||||
|
||||
if name == "int8wo":
|
||||
return Int8WeightOnlyConfig()
|
||||
if name == "int8dq":
|
||||
|
|
@ -105,7 +109,7 @@ def _quant_config(name):
|
|||
try:
|
||||
import torch
|
||||
return MXDynamicActivationMXWeightConfig(
|
||||
activation_dtype=torch.float8_e4m3fn, weight_dtype=torch.float8_e4m3fn
|
||||
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
|
||||
)
|
||||
except TypeError:
|
||||
return MXDynamicActivationMXWeightConfig()
|
||||
|
|
@ -118,7 +122,7 @@ def _make_filter_fn(min_features):
|
|||
timestep/pooled projections (in_features=256) run at M=1 and crash it -- skip them."""
|
||||
import torch.nn as nn
|
||||
|
||||
def filter_fn(module, fqn=""):
|
||||
def filter_fn(module, fqn = ""):
|
||||
return (
|
||||
isinstance(module, nn.Linear)
|
||||
and getattr(module, "in_features", 0) >= min_features
|
||||
|
|
@ -131,11 +135,12 @@ def _make_filter_fn(min_features):
|
|||
def _apply_quant(pipe, name, log, min_features):
|
||||
import torch.nn as nn
|
||||
from torchao.quantization import quantize_
|
||||
|
||||
cfg = _quant_config(name)
|
||||
total = sum(1 for m in pipe.transformer.modules() if isinstance(m, nn.Linear))
|
||||
filt = _make_filter_fn(min_features)
|
||||
q = sum(1 for n, m in pipe.transformer.named_modules() if filt(m, n))
|
||||
quantize_(pipe.transformer, cfg, filter_fn=filt)
|
||||
quantize_(pipe.transformer, cfg, filter_fn = filt)
|
||||
log(f" quantized transformer with {name} ({q}/{total} linears >= {min_features} feat)")
|
||||
|
||||
|
||||
|
|
@ -156,11 +161,17 @@ def _compile(pipe, log):
|
|||
def _gen(pipe, steps, seed, res):
|
||||
import torch
|
||||
|
||||
g = torch.Generator(device="cuda").manual_seed(seed)
|
||||
g = torch.Generator(device = "cuda").manual_seed(seed)
|
||||
torch.cuda.synchronize()
|
||||
t0 = time.time()
|
||||
img = pipe(prompt=PROMPT, width=res, height=res, num_inference_steps=steps,
|
||||
guidance_scale=0.0, generator=g).images[0]
|
||||
img = pipe(
|
||||
prompt = PROMPT,
|
||||
width = res,
|
||||
height = res,
|
||||
num_inference_steps = steps,
|
||||
guidance_scale = 0.0,
|
||||
generator = g,
|
||||
).images[0]
|
||||
torch.cuda.synchronize()
|
||||
return img, time.time() - t0
|
||||
|
||||
|
|
@ -169,23 +180,37 @@ def _median(xs):
|
|||
return sorted(xs)[len(xs) // 2]
|
||||
|
||||
|
||||
def main(argv=None) -> int:
|
||||
def main(argv = None) -> int:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--steps", type=int, default=8)
|
||||
p.add_argument("--res", type=int, default=1024)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--iters", type=int, default=3)
|
||||
p.add_argument("--min-feat", type=int, default=512,
|
||||
help="only quantize Linear with in&out features >= this (int8 _int_mm needs M>16)")
|
||||
p.add_argument("--configs", default="bf16,bf16_c,gguf_c,int8dq_c,fp8dq_c,nvfp4_c,mxfp8_c,int8wo_c",
|
||||
help="comma list; suffix _c = +compile")
|
||||
p.add_argument("--steps", type = int, default = 8)
|
||||
p.add_argument("--res", type = int, default = 1024)
|
||||
p.add_argument("--seed", type = int, default = 42)
|
||||
p.add_argument("--iters", type = int, default = 3)
|
||||
p.add_argument(
|
||||
"--min-feat",
|
||||
type = int,
|
||||
default = 512,
|
||||
help = "only quantize Linear with in&out features >= this (int8 _int_mm needs M>16)",
|
||||
)
|
||||
p.add_argument(
|
||||
"--configs",
|
||||
default = "bf16,bf16_c,gguf_c,int8dq_c,fp8dq_c,nvfp4_c,mxfp8_c,int8wo_c",
|
||||
help = "comma list; suffix _c = +compile",
|
||||
)
|
||||
args = p.parse_args(argv)
|
||||
steps, res, seed, iters = args.steps, args.res, args.seed, args.iters
|
||||
|
||||
import torch
|
||||
OUT.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def run(tag, *, source, quant=None, compile=False):
|
||||
OUT.mkdir(parents = True, exist_ok = True)
|
||||
|
||||
def run(
|
||||
tag,
|
||||
*,
|
||||
source,
|
||||
quant = None,
|
||||
compile = False,
|
||||
):
|
||||
torch.compiler.reset()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
|
|
@ -210,61 +235,79 @@ def main(argv=None) -> int:
|
|||
torch.cuda.empty_cache()
|
||||
return tag, _median(dts), arr, load_peak, gen_peak
|
||||
|
||||
print_ = lambda s: print(s, flush=True) # noqa: E731
|
||||
print_ = lambda s: print(s, flush = True) # noqa: E731
|
||||
|
||||
# config table: tag -> (source, quant, compile)
|
||||
table = {
|
||||
"bf16": ("dense", None, False),
|
||||
"bf16_c": ("dense", None, True),
|
||||
"gguf_c": ("gguf", None, True),
|
||||
"int8wo_c": ("dense", "int8wo", True),
|
||||
"int8dq_c": ("dense", "int8dq", True),
|
||||
"fp8dq_c": ("dense", "fp8dq", True),
|
||||
"nvfp4_c": ("dense", "nvfp4", True),
|
||||
"mxfp8_c": ("dense", "mxfp8", True),
|
||||
"bf16": ("dense", None, False),
|
||||
"bf16_c": ("dense", None, True),
|
||||
"gguf_c": ("gguf", None, True),
|
||||
"int8wo_c": ("dense", "int8wo", True),
|
||||
"int8dq_c": ("dense", "int8dq", True),
|
||||
"fp8dq_c": ("dense", "fp8dq", True),
|
||||
"nvfp4_c": ("dense", "nvfp4", True),
|
||||
"mxfp8_c": ("dense", "mxfp8", True),
|
||||
}
|
||||
want = [c.strip() for c in args.configs.split(",") if c.strip()]
|
||||
|
||||
print(f"== quant probe (Z-Image-Turbo, {res}px, {steps} steps, seed {seed}) ==", flush=True)
|
||||
print(f"== quant probe (Z-Image-Turbo, {res}px, {steps} steps, seed {seed}) ==", flush = True)
|
||||
ref_arr = None
|
||||
rows = []
|
||||
for tag in want:
|
||||
if tag not in table:
|
||||
print(f" {tag}: unknown config, skipping", flush=True)
|
||||
print(f" {tag}: unknown config, skipping", flush = True)
|
||||
continue
|
||||
source, quant, compile = table[tag]
|
||||
print(f"-- {tag} (source={source} quant={quant} compile={compile}) --", flush=True)
|
||||
print(f"-- {tag} (source={source} quant={quant} compile={compile}) --", flush = True)
|
||||
try:
|
||||
_, med, arr, lp, gp = run(tag, source=source, quant=quant, compile=compile)
|
||||
_, med, arr, lp, gp = run(tag, source = source, quant = quant, compile = compile)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
import traceback
|
||||
print(f" {tag:10s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush=True)
|
||||
|
||||
print(f" {tag:10s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True)
|
||||
traceback.print_exc()
|
||||
rows.append((tag, None, None, None, None, None))
|
||||
continue
|
||||
if ref_arr is None and tag == "bf16":
|
||||
ref_arr = arr
|
||||
psnr = _psnr(ref_arr, arr) if ref_arr is not None else None
|
||||
lpips_v = _lpips(ref_arr, arr) if (ref_arr is not None and tag != "bf16") else (0.0 if tag == "bf16" else None)
|
||||
lpips_v = (
|
||||
_lpips(ref_arr, arr)
|
||||
if (ref_arr is not None and tag != "bf16")
|
||||
else (0.0 if tag == "bf16" else None)
|
||||
)
|
||||
rows.append((tag, med, psnr, lpips_v, lp, gp))
|
||||
ps = f"{psnr:.1f}dB" if psnr is not None else "n/a"
|
||||
lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a"
|
||||
print(f" {tag:10s} {med:.3f}s PSNR={ps:>7s} LPIPS={lps:>6s} loadVRAM={lp:.1f}G genVRAM={gp:.1f}G", flush=True)
|
||||
print(
|
||||
f" {tag:10s} {med:.3f}s PSNR={ps:>7s} LPIPS={lps:>6s} loadVRAM={lp:.1f}G genVRAM={gp:.1f}G",
|
||||
flush = True,
|
||||
)
|
||||
|
||||
base = next((r[1] for r in rows if r[0] == "bf16" and r[1]), None)
|
||||
gguf = next((r[1] for r in rows if r[0] == "gguf_c" and r[1]), None)
|
||||
print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush=True)
|
||||
print(f"{'config':10s} {'sec':>7s} {'vs_bf16':>8s} {'vs_gguf':>8s} {'PSNR':>8s} {'LPIPS':>7s} {'loadG':>6s} {'genG':>6s}", flush=True)
|
||||
print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True)
|
||||
print(
|
||||
f"{'config':10s} {'sec':>7s} {'vs_bf16':>8s} {'vs_gguf':>8s} {'PSNR':>8s} {'LPIPS':>7s} {'loadG':>6s} {'genG':>6s}",
|
||||
flush = True,
|
||||
)
|
||||
for tag, med, psnr, lpips_v, lp, gp in rows:
|
||||
if med is None:
|
||||
print(f"{tag:10s} {'FAILED':>7s}", flush=True)
|
||||
print(f"{tag:10s} {'FAILED':>7s}", flush = True)
|
||||
continue
|
||||
vb = f"{base/med:.2f}x" if base else "-"
|
||||
vg = f"{gguf/med:.2f}x" if gguf else "-"
|
||||
ps = f"{psnr:.1f}" if psnr is not None and psnr != float('inf') else ("inf" if psnr == float('inf') else "n/a")
|
||||
ps = (
|
||||
f"{psnr:.1f}"
|
||||
if psnr is not None and psnr != float("inf")
|
||||
else ("inf" if psnr == float("inf") else "n/a")
|
||||
)
|
||||
lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a"
|
||||
print(f"{tag:10s} {med:>7.3f} {vb:>8s} {vg:>8s} {ps:>8s} {lps:>7s} {lp:>6.1f} {gp:>6.1f}", flush=True)
|
||||
print("QUANT-PROBE-DONE", flush=True)
|
||||
print(
|
||||
f"{tag:10s} {med:>7.3f} {vb:>8s} {vg:>8s} {ps:>8s} {lps:>7s} {lp:>6.1f} {gp:>6.1f}",
|
||||
flush = True,
|
||||
)
|
||||
print("QUANT-PROBE-DONE", flush = True)
|
||||
return 0
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue