| """Generate images with FLUX 3 Klein base, then decode the same latents with the full VAE |
| and each tiny decoder. Saves per-image comparisons plus PSNR/LPIPS vs the full-VAE decode. |
| |
| python gen_compare.py --tiny weights/taef2.safetensors weights/tiny_klein3_ft.safetensors |
| """ |
| import argparse, json, os, time |
|
|
| import lpips |
| import numpy as np |
| import torch |
| from PIL import Image, ImageDraw |
| from diffusers import ModularPipeline |
|
|
| from eval_tiny_vae import LatentNorm, load_tiny, psnr |
|
|
| |
| CAPTIONS = [ |
| "Close-up portrait photograph of an elderly fisherman with a weathered, deeply lined face and a short white beard, wearing a navy knitted wool cap. Soft overcast daylight from the left, shallow depth of field, a blurred harbor with moored boats in the background. Sharp detail on skin pores, individual beard hairs and the wool texture. Natural muted colors, 85mm lens look.", |
| "A vintage bakery storefront at dusk photographed straight on. A hand-painted gold-leaf sign above the window reads \"MARGOT'S BAKERY\" in large serif letters, and a small chalkboard by the door says \"FRESH CROISSANTS 7AM\". Warm interior light spills through the window onto shelves of bread loaves and pastries. Dark green wooden frame, brass door handle, wet cobblestone street reflecting the light.", |
| "Macro photograph of a bright blue morpho butterfly resting on a dew-covered fern frond. Iridescent wing scales visible in fine detail, tiny water droplets refracting light, the fern's fine leaflets crisp in the focal plane, creamy green bokeh background. Early morning backlight.", |
| "Wide landscape photograph of a mountain lake at sunrise. Jagged snow-capped granite peaks reflected in perfectly still turquoise water, a dense pine forest along the shoreline, scattered pink clouds, a weathered wooden rowboat tied to a small dock in the foreground. Crisp, high-detail, natural colors.", |
| "Flat vector illustration poster with bold geometric shapes: a stylized orange sun over layered purple and teal hills, a white minimalist bicycle in the foreground, thin line details, a cream paper texture. At the bottom, clean sans-serif text reads \"RIDE FURTHER\". Limited palette, sharp edges, high contrast.", |
| "Product photograph of a luxury mechanical wristwatch on a dark slate surface. Skeleton dial showing intricate gears and jewels, polished steel case with brushed bezel, fine engraved numerals, a black alligator-leather strap with visible grain. Studio lighting with soft reflections, very sharp detail.", |
| ] |
|
|
|
|
| SIZES = [(1024, 1024), (768, 1344), (1344, 768), (1216, 832), (1088, 1920), (1024, 1024)] |
|
|
|
|
| def label(img, text): |
| d = ImageDraw.Draw(img) |
| d.rectangle((0, 0, 9 * len(text) + 12, 26), fill=(0, 0, 0)) |
| d.text((6, 6), text, fill=(255, 255, 255)) |
| return img |
|
|
|
|
| def to_pil(x): |
| return Image.fromarray((x[0].permute(1, 2, 0).float().clamp(0, 1).cpu().numpy() * 255).round().astype(np.uint8)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--tiny", nargs="+", default=["weights/taef2.safetensors", "weights/tiny_klein3_ft.safetensors"]) |
| ap.add_argument("--res", type=int, default=1024) |
| ap.add_argument("--out", default="out/gen") |
| ap.add_argument("--mixed_sizes", action="store_true", help="cycle through non-square sizes in SIZES") |
| args = ap.parse_args() |
| os.makedirs(args.out, exist_ok=True) |
| dev = "cuda" |
|
|
| pipe = ModularPipeline.from_pretrained("black-forest-labs/flux-3-image-klein-base") |
| pipe.load_components(dtype=torch.bfloat16) |
| pipe.to(dev) |
| vae = pipe.vae |
|
|
| |
| captured = [] |
| orig_decode = vae.decode |
|
|
| def spy_decode(z, *a, **k): |
| captured.append(z.detach().float().clone()) |
| return orig_decode(z, *a, **k) |
|
|
| vae.decode = spy_decode |
|
|
| ln = LatentNorm(vae) |
| tinies = {os.path.basename(p).replace(".safetensors", ""): load_tiny(p, dev) for p in args.tiny} |
| lp = lpips.LPIPS(net="alex", verbose=False).to(dev) |
| results = [] |
|
|
| for i, cap in enumerate(CAPTIONS): |
| captured.clear() |
| t = time.time() |
| h, w = SIZES[i % len(SIZES)] if args.mixed_sizes else (args.res, args.res) |
| pipe(prompt=cap, height=h, width=w, generator=torch.Generator().manual_seed(i), output="images") |
| gen_s = time.time() - t |
| z = captured[-1] |
| with torch.no_grad(): |
| torch.cuda.synchronize(); t = time.time() |
| full = orig_decode(z.to(torch.bfloat16)).sample.float().add(1).div(2).clamp(0, 1) |
| torch.cuda.synchronize(); full_ms = (time.time() - t) * 1000 |
| row = {"idx": i, "gen_s": round(gen_s, 1), "full_decode_ms": round(full_ms, 1)} |
| tiles = [label(to_pil(full), "full VAE")] |
| for n, m in tinies.items(): |
| torch.cuda.synchronize(); t = time.time() |
| d = m.decoder(ln.to_norm(z)).clamp(0, 1) |
| torch.cuda.synchronize(); ms = (time.time() - t) * 1000 |
| row[n] = {"psnr_vs_full": round(psnr(d, full), 2), |
| "lpips_vs_full": round(lp(d * 2 - 1, full * 2 - 1).item(), 4), |
| "decode_ms": round(ms, 1)} |
| tiles.append(label(to_pil(d), n)) |
| W, H = tiles[0].size |
| sheet = Image.new("RGB", (W * len(tiles), H)) |
| for j, tl in enumerate(tiles): |
| sheet.paste(tl, (j * W, 0)) |
| sheet.save(os.path.join(args.out, f"gen_{i}.jpg"), quality=92) |
| |
| c = W // 2 - 192, H // 2 - 192 |
| crops = [tl.crop((c[0], c[1], c[0] + 384, c[1] + 384)) for tl in tiles] |
| zs = Image.new("RGB", (384 * len(crops), 384)) |
| for j, cr in enumerate(crops): |
| zs.paste(cr, (j * 384, 0)) |
| zs.save(os.path.join(args.out, f"gen_{i}_zoom.jpg"), quality=92) |
| results.append(row) |
| print(json.dumps(row), flush=True) |
|
|
| print("\nmean vs full-VAE decode of the same generated latents:") |
| for n in tinies: |
| p = np.mean([r[n]["psnr_vs_full"] for r in results]); l = np.mean([r[n]["lpips_vs_full"] for r in results]) |
| print(f" {n:20s} PSNR {p:.2f} dB LPIPS {l:.4f}") |
| json.dump(results, open(os.path.join(args.out, "results.json"), "w"), indent=1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
