File size: 10,007 Bytes
577a4cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """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 # many random crops per high-res image
def __getitem__(self, i):
try:
im = Image.open(self.files[i % len(self.files)]).convert("RGB")
except Exception: # skip broken downloads
return self[random.randrange(len(self.files))]
# random downscale so the model sees both fine detail and whole-object structure
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
# the starting weights are the best so far until something beats them
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: # warm the discriminator up a bit first
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()
|