| """Compare TAEF2 (or a fine-tuned copy) against the full FLUX 3 Klein VAE. |
| |
| Reports reconstruction PSNR / LPIPS on a folder of images, latent agreement, |
| decode speed, parameter counts, and saves a side-by-side grid. |
| |
| python eval_tiny_vae.py --images data/kodak --tiny weights/taef2.safetensors |
| """ |
| import argparse, glob, os, time |
|
|
| import lpips |
| import numpy as np |
| import safetensors.torch as stt |
| import torch |
| import torch.nn.functional as F |
| from PIL import Image |
| from diffusers import AutoencoderKLFlux2 |
|
|
| from taesd import TAESD |
|
|
|
|
| def convert_diffusers_sd_to_taesd(sd): |
| |
| out = {} |
| for k, v in sd.items(): |
| if not k.startswith(("encoder.layers.", "decoder.layers.")): |
| out[k] = v |
| continue |
| encdec, _layers, index, *suffix = k.split(".") |
| offset = 1 if encdec == "decoder" else 0 |
| out[".".join([encdec, str(int(index) + offset), *suffix])] = v |
| return out |
|
|
|
|
| def load_tiny(path, device, dtype=torch.float32): |
| m = TAESD(encoder_path=None, decoder_path=None, latent_channels=32, arch_variant="flux_2") |
| sd = stt.load_file(path) |
| if any(k.startswith(("encoder.layers.", "decoder.layers.")) for k in sd): |
| sd = convert_diffusers_sd_to_taesd(sd) |
| m.load_state_dict(sd) |
| return m.to(device, dtype).eval().requires_grad_(False) |
|
|
|
|
| def patchify(z): |
| b, c, h, w = z.shape |
| return z.view(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4).reshape(b, c * 4, h // 2, w // 2) |
|
|
|
|
| def unpatchify(z): |
| b, c, h, w = z.shape |
| return z.reshape(b, c // 4, 2, 2, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // 4, h * 2, w * 2) |
|
|
|
|
| class LatentNorm: |
| """Raw VAE latents <-> the BN-normalized space the FLUX transformer (and TAEF2) work in.""" |
|
|
| def __init__(self, vae): |
| self.mean = vae.bn.running_mean.view(1, -1, 1, 1).float() |
| self.std = torch.sqrt(vae.bn.running_var.view(1, -1, 1, 1).float() + vae.config.batch_norm_eps) |
|
|
| def to_norm(self, z): |
| return unpatchify((patchify(z) - self.mean.to(z)) / self.std.to(z)) |
|
|
| def to_raw(self, z): |
| return unpatchify(patchify(z) * self.std.to(z) + self.mean.to(z)) |
|
|
|
|
| def load_images(folder, mult=16): |
| xs, names = [], [] |
| for p in sorted(glob.glob(os.path.join(folder, "*"))): |
| if not p.lower().endswith((".png", ".jpg", ".jpeg", ".webp")): |
| continue |
| im = Image.open(p).convert("RGB") |
| w, h = im.size |
| im = im.crop((0, 0, w - w % mult, h - h % mult)) |
| xs.append(torch.from_numpy(np.asarray(im)).permute(2, 0, 1).float() / 255.0) |
| names.append(os.path.basename(p)) |
| return xs, names |
|
|
|
|
| def psnr(a, b): |
| return (10 * torch.log10(1.0 / F.mse_loss(a, b))).item() |
|
|
|
|
| def timed(fn, x, n=10): |
| fn(x); torch.cuda.synchronize() |
| t = time.perf_counter() |
| for _ in range(n): |
| fn(x) |
| torch.cuda.synchronize() |
| return (time.perf_counter() - t) / n * 1000 |
|
|
|
|
| @torch.no_grad() |
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--images", default="data/kodak") |
| ap.add_argument("--vae", default="weights/flux3-klein-base/vae") |
| ap.add_argument("--tiny", nargs="+", default=["weights/taef2.safetensors"]) |
| ap.add_argument("--grid", default="out/compare.jpg") |
| ap.add_argument("--latent_space", choices=["norm", "raw"], default="norm", |
| help="space the tiny AE reads/writes: BN-normalized (TAEF2) or raw VAE latents") |
| args = ap.parse_args() |
|
|
| dev = "cuda" |
| vae = AutoencoderKLFlux2.from_pretrained(args.vae, torch_dtype=torch.float32).to(dev).eval() |
| ln = LatentNorm(vae) |
| to_tiny = ln.to_norm if args.latent_space == "norm" else (lambda z: z) |
| from_tiny = ln.to_raw if args.latent_space == "norm" else (lambda z: z) |
| tinies = {os.path.basename(p): load_tiny(p, dev) for p in args.tiny} |
| lp = lpips.LPIPS(net="alex", verbose=False).to(dev) |
|
|
| nparam = lambda m: sum(p.numel() for p in m.parameters()) / 1e6 |
| print(f"full VAE : {nparam(vae):7.2f}M params (enc {nparam(vae.encoder):.2f}M / dec {nparam(vae.decoder):.2f}M)") |
| for n, t in tinies.items(): |
| print(f"{n:11s}: {nparam(t):7.2f}M params (enc {nparam(t.encoder):.2f}M / dec {nparam(t.decoder):.2f}M)") |
|
|
| xs, names = load_images(args.images) |
| res = {k: [] for k in ["full"] + [f"{n}|dec" for n in tinies] + [f"{n}|enc" for n in tinies] + [f"{n}|both" for n in tinies]} |
| lat_stats = {n: [] for n in tinies} |
| grid_rows = [] |
| for i, x in enumerate(xs): |
| x = x[None].to(dev) |
| z = vae.encode(x * 2 - 1).latent_dist.mode() |
| full = vae.decode(z).sample.add(1).div(2).clamp(0, 1) |
| res["full"].append((psnr(full, x), lp(full * 2 - 1, x * 2 - 1).item())) |
| row = [x, full] |
| for n, t in tinies.items(): |
| zn = to_tiny(z) |
| d = t.decoder(zn).clamp(0, 1) |
| res[f"{n}|dec"].append((psnr(d, x), lp(d * 2 - 1, x * 2 - 1).item())) |
| zt = t.encoder(x) |
| e = vae.decode(from_tiny(zt)).sample.add(1).div(2).clamp(0, 1) |
| res[f"{n}|enc"].append((psnr(e, x), lp(e * 2 - 1, x * 2 - 1).item())) |
| b = t.decoder(zt).clamp(0, 1) |
| res[f"{n}|both"].append((psnr(b, x), lp(b * 2 - 1, x * 2 - 1).item())) |
| cos = F.cosine_similarity(zt.flatten(), zn.flatten(), dim=0).item() |
| lat_stats[n].append((F.mse_loss(zt, zn).item(), cos, zn.std().item(), zt.std().item())) |
| row.append(d) |
| if i < 4: |
| grid_rows.append(torch.cat([r[..., :512, :512] for r in row], -1)) |
|
|
| print("\nreconstruction vs original image (mean over", len(xs), "images)") |
| print(f"{'path':34s} {'PSNR dB':>8s} {'LPIPS':>7s}") |
| for k, v in res.items(): |
| v = np.array(v) |
| label = {"full": "full VAE enc -> full VAE dec"}.get(k, k.replace("|dec", ": full enc -> TINY dec").replace("|enc", ": TINY enc -> full dec").replace("|both", ": TINY enc -> TINY dec")) |
| print(f"{label:34s} {v[:,0].mean():8.2f} {v[:,1].mean():7.4f}") |
| for n, s in lat_stats.items(): |
| s = np.array(s) |
| print(f"\n{n} encoder vs full encoder latents: MSE {s[:,0].mean():.4f} cosine {s[:,1].mean():.4f} std full {s[:,2].mean():.3f} / tiny {s[:,3].mean():.3f}") |
|
|
| print("\ndecode latency (bf16, batch 1)") |
| vae_bf = vae.to(torch.bfloat16) |
| for res_px in (1024, 2048): |
| z = torch.randn(1, 32, res_px // 8, res_px // 8, device=dev, dtype=torch.bfloat16) |
| torch.cuda.reset_peak_memory_stats() |
| tf = timed(lambda z: vae_bf.decode(z).sample, z, 5) |
| mf = torch.cuda.max_memory_allocated() / 2**30 |
| line = f" {res_px}px full VAE {tf:8.1f} ms ({mf:.1f} GB peak)" |
| for n, t in tinies.items(): |
| tb = t.to(torch.bfloat16) |
| torch.cuda.reset_peak_memory_stats() |
| tt = timed(tb.decoder, z, 10) |
| mt = torch.cuda.max_memory_allocated() / 2**30 |
| line += f" | {n} {tt:6.1f} ms ({mt:.1f} GB) -> {tf/tt:.1f}x faster" |
| t.float() |
| print(line) |
|
|
| os.makedirs(os.path.dirname(args.grid) or ".", exist_ok=True) |
| g = torch.cat(grid_rows, -2)[0].permute(1, 2, 0).float().cpu().numpy() |
| Image.fromarray((g * 255).round().astype(np.uint8)).save(args.grid, quality=92) |
| print("\ngrid columns: original | full VAE |", " | ".join(tinies), "->", args.grid) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
