City
8d619afdf8
This is probably not how it's supposed to be done but it stops the horrible coil whine at least.
125 lines
4.1 KiB
Python
125 lines
4.1 KiB
Python
import os
|
|
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import argparse
|
|
import random
|
|
from PIL import Image
|
|
from tqdm import tqdm
|
|
from safetensors.torch import save_file, load_file
|
|
|
|
from interposer import Interposer
|
|
from vae import get_vae
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Train latent interposer model")
|
|
parser.add_argument("--steps", type=int, default=500000, help="No. of training steps")
|
|
parser.add_argument('--bs', type=int, default=1, help="Batch size")
|
|
parser.add_argument('--lr', default="1e-8", help="Learning rate")
|
|
parser.add_argument("-n", "--save_every_n", type=int, dest="save", default=50000, help="Save model/sample periodically")
|
|
parser.add_argument('--src', choices=["v1","xl"], required=True, help="Source latent format")
|
|
parser.add_argument('--dst', choices=["v1","xl"], required=True, help="Destination latent format")
|
|
parser.add_argument('--resume', help="Checkpoint to resume from")
|
|
args = parser.parse_args()
|
|
if args.src == args.dst:
|
|
parser.error("--src and --dst can't be the same")
|
|
try:
|
|
float(args.lr)
|
|
except:
|
|
parser.error("--lr must be a valid float eg. 0.001 or 1e-3")
|
|
return args
|
|
|
|
class Latent:
|
|
def __init__(self, md5, lat_src, lat_dst, dev):
|
|
if lat_src == "v1": src = os.path.join("latent_v1", f"{md5}.npy")
|
|
if lat_src == "xl": src = os.path.join("latent_xl", f"{md5}.npy")
|
|
|
|
if lat_dst == "v1": dst = os.path.join("latent_v1", f"{md5}.npy")
|
|
if lat_dst == "xl": dst = os.path.join("latent_xl", f"{md5}.npy")
|
|
|
|
self.src = torch.from_numpy(np.load(src)).to(dev)
|
|
self.dst = torch.from_numpy(np.load(dst)).to(dev)
|
|
|
|
def load_latents(src, dst, dev):
|
|
print("Loading latents from disk")
|
|
latents = []
|
|
for i in tqdm(os.listdir("images")):
|
|
md5 = os.path.splitext(i)[0]
|
|
latents.append(Latent(md5, src, dst, dev))
|
|
return latents
|
|
|
|
vae = None
|
|
def sample_decode(latent, filename, version):
|
|
global vae
|
|
if not vae:
|
|
vae = get_vae(version, fp16=True)
|
|
vae.to("cuda")
|
|
|
|
latent = latent.half().to("cuda")
|
|
out = vae.decode(latent).sample
|
|
out = out.cpu().detach().numpy()
|
|
out = np.squeeze(out, 0)
|
|
out = out.transpose((1, 2, 0))
|
|
out = np.clip(out, -1.0, 1.0)
|
|
out = (out+1)/2 * 255
|
|
out = out.astype(np.uint8)
|
|
out = Image.fromarray(out)
|
|
out.save(filename)
|
|
|
|
if __name__ == "__main__":
|
|
args = parse_args()
|
|
target_dev = "cuda"
|
|
latent_src = args.src
|
|
latent_dst = args.dst
|
|
|
|
latents = load_latents(latent_src, latent_dst, target_dev)
|
|
|
|
if not os.path.isdir("models"): os.mkdir("models")
|
|
log = open(f"models/{latent_src}-to-{latent_dst}_interposer.csv", "w")
|
|
|
|
if os.path.isfile(f"test_{latent_src}.npy"):
|
|
sample_latent = torch.from_numpy(np.load(f"test_{latent_src}.npy")).to(target_dev)
|
|
else:
|
|
sample_latent = random.choice(latents).src
|
|
|
|
model = Interposer()
|
|
if args.resume:
|
|
model.load_state_dict(load_file(args.resume))
|
|
model.to(target_dev)
|
|
|
|
criterion = torch.nn.MSELoss(size_average=False)
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=float(args.lr)/args.bs)
|
|
|
|
for t in tqdm(range(int(args.steps/args.bs)), unit_scale=args.bs):
|
|
step = t*args.bs
|
|
# input batch
|
|
lts = [random.choice(latents) for _ in range(args.bs)]
|
|
src = torch.cat([x.src for x in lts],0)
|
|
dst = torch.cat([x.dst for x in lts],0)
|
|
|
|
y_pred = model(src) # forward
|
|
loss = criterion(y_pred, dst) # loss
|
|
|
|
# backward
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
# print loss
|
|
if step%1000 == 0:
|
|
tqdm.write(f"{step} - {loss.data.item()/args.bs:.2f}")
|
|
log.write(f"{step},{loss.data.item()/args.bs:.2f}\n")
|
|
log.flush()
|
|
|
|
# sample/save
|
|
if step%args.save == 0:
|
|
out = model(sample_latent)
|
|
output_name = f"./models/{latent_src}-to-{latent_dst}_interposer_e{step/1000}k"
|
|
sample_decode(out, f"{output_name}.png", latent_dst)
|
|
save_file(model.state_dict(), f"{output_name}.safetensors")
|
|
# save final output
|
|
output_name = f"./models/{latent_src}-to-{latent_dst}_interposer_e{args.steps/1000}k"
|
|
sample_decode(out, f"{output_name}.png", "v1")
|
|
save_file(model.state_dict(), f"{output_name}.safetensors")
|
|
log.close()
|