Snapshot
This commit is contained in:
City
@@ -1,29 +1,31 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
import argparse
|
||||
from PIL import Image
|
||||
from tqdm import tqdm
|
||||
from safetensors.torch import save_file, load_file
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torch.utils.data import DataLoader
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from interposer import Interposer
|
||||
from vae import get_vae
|
||||
from interposer import InterposerModel as Model
|
||||
from dataset import LatentDataset
|
||||
from utils import ModelWrapper
|
||||
|
||||
torch.backends.cudnn.benchmark = True
|
||||
torch.manual_seed(0)
|
||||
|
||||
TARGET_DEV = "cuda"
|
||||
|
||||
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=4, help="Batch size")
|
||||
parser.add_argument('--lr', default="1e-4", help="Learning rate")
|
||||
parser.add_argument("-n", "--save_every_n", type=int, dest="save", default=50000, help="Save model/sample periodically")
|
||||
parser.add_argument("-s", "--steps", type=int, default=500000, help="No. of training steps")
|
||||
parser.add_argument("-b", "--batch", type=int, default= 1, help="Batch size")
|
||||
parser.add_argument("-n", "--nsave", type=int, default= 50000, help="Save model/sample periodically")
|
||||
parser.add_argument('--rev', default="v4.0-rc1", help="Revision/log ID")
|
||||
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('--lr', default="1e-4", help="Learning rate")
|
||||
parser.add_argument('--lrskip', type=int, default=0, help="Constant lr for first N steps")
|
||||
parser.add_argument('--cosine', action=argparse.BooleanOptionalAction, help="Use cosine scheduler")
|
||||
parser.add_argument('--resume', help="Checkpoint to resume from")
|
||||
parser.add_argument('--cosine', action=argparse.BooleanOptionalAction, help="Use cosine scheduler to taper off LR")
|
||||
args = parser.parse_args()
|
||||
if args.src == args.dst:
|
||||
parser.error("--src and --dst can't be the same")
|
||||
@@ -33,120 +35,28 @@ def parse_args():
|
||||
parser.error("--lr must be a valid float eg. 0.001 or 1e-3")
|
||||
return args
|
||||
|
||||
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)
|
||||
|
||||
def get_eval_data(dataset, src_path, dst_path, target_dev):
|
||||
if os.path.isfile(src_path) and os.path.isfile(dst_path):
|
||||
src = LatentDataset.load_latent(None, src_path)
|
||||
dst = LatentDataset.load_latent(None, dst_path)
|
||||
else:
|
||||
src = dataset[0][0]
|
||||
dst = dataset[0][1]
|
||||
src = src.float().to(target_dev).unsqueeze(0)
|
||||
dst = dst.float().to(target_dev).unsqueeze(0)
|
||||
return(src, dst)
|
||||
|
||||
def eval_model(step, model, criterion, scheduler, src, dst):
|
||||
with torch.no_grad():
|
||||
t_pred = model(src)
|
||||
t_loss = criterion(t_pred, dst)
|
||||
tqdm.write(f"{str(step):<10} {loss.data.item():.4e}|{t_loss.data.item():.4e} @ {float(scheduler.get_last_lr()[0]):.4e}")
|
||||
log.write(f"{step},{loss.data.item()},{t_loss.data.item()},{float(scheduler.get_last_lr()[0])}\n")
|
||||
log.flush()
|
||||
|
||||
def save_model(step, model, optim, lat, src, dst):
|
||||
with torch.no_grad():
|
||||
out = model(lat)
|
||||
output_name = f"./models/{src}-to-{dst}_interposer_e{round(step/1000)}k"
|
||||
sample_decode(out, f"{output_name}.png", dst)
|
||||
save_file(model.state_dict(), f"{output_name}.safetensors")
|
||||
torch.save(optim.state_dict(), f"{output_name}.optim.pth")
|
||||
|
||||
class LatentDataset(Dataset):
|
||||
class Shard:
|
||||
def __init__(self, root, fname, res, src, dst):
|
||||
self.fname = fname
|
||||
self.src_path = f"{root}/{src}_{res}px/{fname}.npy"
|
||||
self.dst_path = f"{root}/{dst}_{res}px/{fname}.npy"
|
||||
|
||||
def __init__(self, res, src, dst, root="latents"):
|
||||
print("Loading latents from disk")
|
||||
self.latents = []
|
||||
for i in tqdm(os.listdir(f"{root}/{src}_{res}px")):
|
||||
fname, ext = os.path.splitext(i)
|
||||
assert ext == ".npy"
|
||||
s = self.Shard(root, fname, res, src, dst)
|
||||
if os.path.isfile(s.src_path) and os.path.isfile(s.dst_path):
|
||||
self.latents.append(s)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.latents)
|
||||
|
||||
def __getitem__(self, index):
|
||||
s = self.latents[index]
|
||||
src = self.load_latent(s.src_path)
|
||||
dst = self.load_latent(s.dst_path)
|
||||
return (src, dst)
|
||||
|
||||
def load_latent(self, path):
|
||||
lat = torch.from_numpy(np.load(path))
|
||||
if lat.shape[0] == 1:
|
||||
lat = torch.squeeze(lat, 0)
|
||||
assert not torch.isnan(torch.sum(lat.float()))
|
||||
return lat
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
target_dev = "cuda"
|
||||
resolution = 768
|
||||
|
||||
dataset = LatentDataset(resolution, args.src, args.dst)
|
||||
dataset = LatentDataset([args.src, args.dst])
|
||||
loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=args.bs,
|
||||
shuffle=True,
|
||||
num_workers=0,
|
||||
# num_workers=4,
|
||||
# persistent_workers=True,
|
||||
batch_size = args.batch,
|
||||
shuffle = True,
|
||||
drop_last = True,
|
||||
pin_memory = False,
|
||||
# num_workers = 0,
|
||||
num_workers = 4,
|
||||
persistent_workers=True,
|
||||
)
|
||||
eval_src, eval_dst = get_eval_data(
|
||||
dataset,
|
||||
f"latents/test_{args.src}_{resolution}px.npy",
|
||||
f"latents/test_{args.dst}_{resolution}px.npy",
|
||||
target_dev,
|
||||
)
|
||||
|
||||
os.makedirs("models", exist_ok=True)
|
||||
log = open(f"models/{args.src}-to-{args.dst}_interposer.csv", "w")
|
||||
|
||||
model = Interposer()
|
||||
|
||||
model = Model() # TODO: handle scale factor/channels for non-sd VAEs
|
||||
criterion = torch.nn.L1Loss()
|
||||
optimizer = torch.optim.AdamW(model.parameters(), lr=float(args.lr))
|
||||
# import bitsandbytes as bnb
|
||||
# optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=float(args.lr))
|
||||
|
||||
scheduler = None
|
||||
if args.cosine:
|
||||
print("Using CosineAnnealingLR")
|
||||
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
|
||||
optimizer, T_max = int(args.steps/args.bs),
|
||||
optimizer, T_max = int(args.steps/args.batch),
|
||||
)
|
||||
else:
|
||||
print("Using LinearLR")
|
||||
@@ -154,23 +64,35 @@ if __name__ == "__main__":
|
||||
optimizer,
|
||||
start_factor = 0.1,
|
||||
end_factor = 1.0,
|
||||
total_iters = int(5000/args.bs),
|
||||
total_iters = int(5000/args.batch),
|
||||
)
|
||||
|
||||
if args.resume:
|
||||
model.load_state_dict(load_file(args.resume))
|
||||
model.to(target_dev)
|
||||
model.to(TARGET_DEV)
|
||||
optimizer.load_state_dict(torch.load(
|
||||
f"{os.path.splitext(args.resume)[0]}.optim.pth"
|
||||
))
|
||||
optimizer.param_groups[0]['lr'] = scheduler.base_lrs[0]
|
||||
else:
|
||||
model.to(target_dev)
|
||||
model.to(TARGET_DEV)
|
||||
|
||||
wrapper = ModelWrapper( # model wrapper for saving/eval/etc
|
||||
name = f"{args.src}-to-{args.dst}_interposer-{args.rev}",
|
||||
specs = [args.src, args.dst],
|
||||
model = model,
|
||||
evals = dataset.get_eval(),
|
||||
device = TARGET_DEV,
|
||||
criterion = criterion,
|
||||
optimizer = optimizer,
|
||||
scheduler = scheduler,
|
||||
)
|
||||
|
||||
progress = tqdm(total=args.steps)
|
||||
while progress.n < args.steps:
|
||||
for src, dst in loader:
|
||||
src = src.to(target_dev)
|
||||
dst = dst.to(target_dev)
|
||||
src = src.to(TARGET_DEV)
|
||||
dst = dst.to(TARGET_DEV)
|
||||
with torch.cuda.amp.autocast():
|
||||
y_pred = model(src) # forward
|
||||
loss = criterion(y_pred, dst) # loss
|
||||
@@ -179,19 +101,15 @@ if __name__ == "__main__":
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
if progress.n >= args.lrskip: scheduler.step()
|
||||
|
||||
# eval/save
|
||||
progress.update(args.bs)
|
||||
if progress.n % (1000 + 1000%args.bs) == 0:
|
||||
eval_model(progress.n, model, criterion, scheduler, eval_src, eval_dst)
|
||||
if progress.n % (args.save + args.save%args.bs) == 0:
|
||||
save_model(progress.n, model, optimizer, eval_src, args.src, args.dst)
|
||||
progress.update(args.batch)
|
||||
wrapper.log_step(loss.data.item(), progress.n)
|
||||
if args.nsave > 0 and progress.n % (args.nsave + args.nsave%args.batch) == 0:
|
||||
wrapper.save_model(step=progress.n)
|
||||
if progress.n >= args.steps:
|
||||
break
|
||||
progress.close()
|
||||
|
||||
# save final output
|
||||
eval_model(progress.n, model, criterion, scheduler, eval_src, eval_dst)
|
||||
save_model(progress.n, model, optimizer, eval_src, args.src, args.dst)
|
||||
log.close()
|
||||
wrapper.save_model(epoch="") # final save
|
||||
wrapper.close()
|
||||
|
||||
Reference in New Issue
Block a user