| """Fine-tune TAEF2 into a better tiny VAE for FLUX 3 Klein (same VAE as FLUX.2). |
| |
| Teacher: the full FLUX 3 Klein / FLUX.2 VAE encoder makes the target latents (in the |
| BN-normalized space the transformer uses). Student: TAEF2's decoder learns to turn those |
| latents back into the real image. Optionally the student encoder is trained to match the |
| teacher latents too. |
| |
| python finetune_tiny_vae.py --data data/div2k --steps 3000 --batch 16 --res 512 |
| """ |
| import argparse, copy, glob, math, os, random, time |
|
|
| import lpips |
| import safetensors.torch as stt |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from PIL import Image |
| from torch.utils.data import DataLoader, Dataset |
| from torchvision.transforms import functional as TF |
| from diffusers import AutoencoderKLFlux2 |
|
|
| from eval_tiny_vae import LatentNorm, load_images, load_tiny, psnr |
|
|
|
|
| class ImageFolder(Dataset): |
| def __init__(self, roots, res): |
| exts = (".png", ".jpg", ".jpeg", ".webp") |
| self.files = [p for root in roots for p in glob.glob(os.path.join(root, "**", "*"), recursive=True) |
| if p.lower().endswith(exts)] |
| assert self.files, f"no images under {roots}" |
| print(f"{len(self.files)} training images", flush=True) |
| self.res = res |
|
|
| def __len__(self): |
| return len(self.files) * 50 |
|
|
| def __getitem__(self, i): |
| try: |
| im = Image.open(self.files[i % len(self.files)]).convert("RGB") |
| except Exception: |
| return self[random.randrange(len(self.files))] |
| |
| short = min(im.size) |
| target = random.randint(self.res, max(self.res, min(short, self.res * 3))) |
| if target != short: |
| im = TF.resize(im, target, antialias=True) |
| w, h = im.size |
| x0, y0 = random.randint(0, w - self.res), random.randint(0, h - self.res) |
| im = im.crop((x0, y0, x0 + self.res, y0 + self.res)) |
| if random.random() < 0.5: |
| im = TF.hflip(im) |
| return TF.to_tensor(im) |
|
|
|
|
| class PatchDisc(nn.Module): |
| """Small PatchGAN discriminator for optional adversarial sharpening.""" |
|
|
| def __init__(self, ch=64): |
| super().__init__() |
| sn = nn.utils.spectral_norm |
| self.net = nn.Sequential( |
| sn(nn.Conv2d(3, ch, 4, 2, 1)), nn.LeakyReLU(0.2), |
| sn(nn.Conv2d(ch, ch * 2, 4, 2, 1)), nn.LeakyReLU(0.2), |
| sn(nn.Conv2d(ch * 2, ch * 4, 4, 2, 1)), nn.LeakyReLU(0.2), |
| sn(nn.Conv2d(ch * 4, 1, 3, 1, 1)), |
| ) |
|
|
| def forward(self, x): |
| return self.net(x * 2 - 1) |
|
|
|
|
| @torch.no_grad() |
| def evaluate(student, vae, ln, lp, imgs, dev): |
| student.eval() |
| ps, ls = [], [] |
| for x in imgs: |
| x = x[None].to(dev) |
| zn = ln.to_norm(vae.encode((x * 2 - 1).to(torch.bfloat16)).latent_dist.mode().float()) |
| with torch.autocast("cuda", torch.bfloat16): |
| d = student.decoder(zn) |
| d = d.float().clamp(0, 1) |
| ps.append(psnr(d, x)); ls.append(lp(d * 2 - 1, x * 2 - 1).item()) |
| student.train() |
| return sum(ps) / len(ps), sum(ls) / len(ls) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--data", nargs="+", required=True) |
| ap.add_argument("--eval_images", default="data/kodak") |
| ap.add_argument("--vae", default="weights/flux3-klein-base/vae") |
| ap.add_argument("--init", default="weights/taef2.safetensors") |
| ap.add_argument("--out", default="runs/ft") |
| ap.add_argument("--steps", type=int, default=3000) |
| ap.add_argument("--batch", type=int, default=16) |
| ap.add_argument("--res", type=int, default=512) |
| ap.add_argument("--lr", type=float, default=1e-4) |
| ap.add_argument("--lpips_w", type=float, default=0.5) |
| ap.add_argument("--gan_w", type=float, default=0.0, help="e.g. 0.05 to add adversarial loss") |
| ap.add_argument("--gan_start", type=int, default=1000) |
| ap.add_argument("--disc_init", default=None, help="resume the discriminator from a previous run's disc.pt") |
| ap.add_argument("--noise_p", type=float, default=0.3, help="prob. of jittering latents (robust to diffusion outputs)") |
| ap.add_argument("--train_encoder", action="store_true") |
| ap.add_argument("--eval_every", type=int, default=500) |
| ap.add_argument("--workers", type=int, default=32) |
| ap.add_argument("--ema", type=float, default=0.999, help="EMA decay for the student weights (0 = off)") |
| args = ap.parse_args() |
|
|
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
| dev = "cuda" |
| os.makedirs(args.out, exist_ok=True) |
|
|
| vae = AutoencoderKLFlux2.from_pretrained(args.vae, torch_dtype=torch.bfloat16).to(dev).eval().requires_grad_(False) |
| ln = LatentNorm(vae) |
| student = load_tiny(args.init, dev).train().requires_grad_(True) |
| student.encoder.requires_grad_(args.train_encoder) |
| lp = lpips.LPIPS(net="vgg", verbose=False).to(dev).eval().requires_grad_(False) |
| lp_eval = lpips.LPIPS(net="alex", verbose=False).to(dev).eval().requires_grad_(False) |
|
|
| params = [p for p in student.parameters() if p.requires_grad] |
| opt = torch.optim.AdamW(params, lr=args.lr, betas=(0.9, 0.99), weight_decay=0.0) |
| sched = torch.optim.lr_scheduler.LambdaLR( |
| opt, lambda s: min(1.0, (s + 1) / 200) * (0.5 * (1 + math.cos(math.pi * s / args.steps)) * 0.9 + 0.1)) |
| disc = d_opt = None |
| if args.gan_w > 0: |
| disc = PatchDisc().to(dev) |
| d_opt = torch.optim.AdamW(disc.parameters(), lr=args.lr, betas=(0.5, 0.99)) |
| if args.disc_init: |
| ck = torch.load(args.disc_init, map_location=dev) |
| disc.load_state_dict(ck["disc"]); d_opt.load_state_dict(ck["opt"]) |
|
|
| ema = None |
| if args.ema > 0: |
| ema = copy.deepcopy(student).eval().requires_grad_(False) |
|
|
| eval_imgs, _ = load_images(args.eval_images) |
| p0, l0 = evaluate(student, vae, ln, lp_eval, eval_imgs, dev) |
| print(f"step 0 eval PSNR {p0:.2f} LPIPS {l0:.4f}", flush=True) |
| best = l0 |
| |
| stt.save_file({k: v.detach().float().contiguous().cpu() for k, v in student.state_dict().items()}, |
| os.path.join(args.out, "best.safetensors")) |
|
|
| dl = DataLoader(ImageFolder(args.data, args.res), batch_size=args.batch, shuffle=True, |
| num_workers=args.workers, pin_memory=True, drop_last=True, persistent_workers=True) |
| it = iter(dl) |
| t0 = time.time() |
| for step in range(1, args.steps + 1): |
| try: |
| x = next(it) |
| except StopIteration: |
| it = iter(dl); x = next(it) |
| x = x.to(dev, non_blocking=True) |
| with torch.no_grad(): |
| zn = ln.to_norm(vae.encode((x * 2 - 1).to(torch.bfloat16)).latent_dist.mode().float()) |
| z_in = zn |
| if args.noise_p > 0: |
| mask = (torch.rand(zn.shape[0], 1, 1, 1, device=dev) < args.noise_p).float() |
| z_in = zn + mask * torch.rand_like(mask) * 0.1 * torch.randn_like(zn) |
|
|
| with torch.autocast("cuda", torch.bfloat16): |
| out = student.decoder(z_in) |
| out = out.float() |
| l1 = F.l1_loss(out, x) |
| lpl = lp(out.clamp(0, 1) * 2 - 1, x * 2 - 1).mean() |
| loss = l1 + args.lpips_w * lpl |
| if args.train_encoder: |
| zt = student.encoder(x).float() |
| loss = loss + F.mse_loss(zt, zn) |
| g_loss = torch.zeros((), device=dev) |
| if disc is not None and step >= args.gan_start: |
| g_loss = -disc(out).mean() |
| loss = loss + args.gan_w * g_loss |
|
|
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(params, 1.0) |
| opt.step(); sched.step() |
|
|
| if disc is not None and step >= args.gan_start - 200: |
| with torch.autocast("cuda", torch.bfloat16): |
| d_loss = F.relu(1 - disc(x)).mean() + F.relu(1 + disc(out.detach())).mean() |
| d_opt.zero_grad(set_to_none=True); d_loss.backward(); d_opt.step() |
|
|
| if step % 50 == 0: |
| ips = step * args.batch / (time.time() - t0) |
| print(f"step {step:6d} loss {loss.item():.4f} l1 {l1.item():.4f} lpips {lpl.item():.4f} " |
| f"gan {g_loss.item():.3f} lr {sched.get_last_lr()[0]:.2e} {ips:.0f} img/s", flush=True) |
| if ema is not None: |
| d = min(args.ema, (1 + step) / (10 + step)) |
| with torch.no_grad(): |
| for pe, ps in zip(ema.parameters(), student.parameters()): |
| pe.lerp_(ps, 1 - d) |
|
|
| if step % args.eval_every == 0 or step == args.steps: |
| sd_of = lambda m: {k: v.detach().float().contiguous().cpu() for k, v in m.state_dict().items()} |
| p, l = evaluate(student, vae, ln, lp_eval, eval_imgs, dev) |
| print(f"step {step} eval PSNR {p:.2f} ({p - p0:+.2f}) LPIPS {l:.4f} ({l - l0:+.4f})", flush=True) |
| cands = [(l, sd_of(student))] |
| stt.save_file(cands[0][1], os.path.join(args.out, "last.safetensors")) |
| if ema is not None: |
| pe_, le_ = evaluate(ema, vae, ln, lp_eval, eval_imgs, dev) |
| print(f"step {step} ema eval PSNR {pe_:.2f} ({pe_ - p0:+.2f}) LPIPS {le_:.4f} ({le_ - l0:+.4f})", flush=True) |
| cands.append((le_, sd_of(ema))) |
| stt.save_file(cands[1][1], os.path.join(args.out, "ema.safetensors")) |
| if disc is not None: |
| torch.save({"disc": disc.state_dict(), "opt": d_opt.state_dict()}, os.path.join(args.out, "disc.pt")) |
| lc, sdc = min(cands, key=lambda c: c[0]) |
| if lc < best: |
| best = lc |
| stt.save_file(sdc, os.path.join(args.out, "best.safetensors")) |
| print(f"step {step} new best LPIPS {lc:.4f}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
