124 lines
3.6 KiB
Python
124 lines
3.6 KiB
Python
#
|
|
# This file just has all the random saving/logging/eval related code
|
|
#
|
|
import os
|
|
import torch
|
|
from tqdm import tqdm
|
|
from diffusers import AutoencoderKL
|
|
from safetensors.torch import save_file
|
|
from torchvision.utils import save_image
|
|
|
|
LOSS_MEMORY = 500
|
|
LOG_EVERY_N = 500
|
|
SAVE_FOLDER = "models"
|
|
|
|
class ModelWrapper:
|
|
def __init__(self, name, specs, model, optimizer, criterion, scheduler, device="cpu", evals=[None,None], stdout=True):
|
|
self.name = name
|
|
self.specs = specs
|
|
self.losses = []
|
|
|
|
self.model = model
|
|
self.optimizer = optimizer
|
|
self.criterion = criterion
|
|
self.scheduler = scheduler
|
|
|
|
self.device = device
|
|
self.vae = self.get_vae(self.specs[1], fp16=True)
|
|
self.eval_src = evals[0]
|
|
self.eval_dst = evals[1]
|
|
|
|
os.makedirs(SAVE_FOLDER, exist_ok=True)
|
|
self.csvlog = open(f"{SAVE_FOLDER}/{self.name}.csv", "w")
|
|
self.stdout = stdout
|
|
|
|
def log_step(self, loss, step=None):
|
|
self.losses.append(loss)
|
|
step = step if step else len(self.losses)
|
|
if step % LOG_EVERY_N == 0:
|
|
self.log_main(step)
|
|
|
|
def log_main(self, step=None):
|
|
lr = float(self.scheduler.get_last_lr()[0])
|
|
avg = sum(self.losses[-LOSS_MEMORY:])/LOSS_MEMORY
|
|
evl = self.eval_model()[0]
|
|
if self.stdout:
|
|
tqdm.write(f"{str(step):<10} {avg:.4e}|{evl:.4e} @ {lr:.4e}")
|
|
if self.csvlog:
|
|
self.csvlog.write(f"{step},{avg},{evl},{lr}\n")
|
|
self.csvlog.flush()
|
|
|
|
def eval_model(self):
|
|
with torch.no_grad():
|
|
pred = self.model(self.eval_src.to(self.device))
|
|
loss = self.criterion(pred, self.eval_dst.to(self.device))
|
|
return loss, pred
|
|
|
|
def save_model(self, step=None, epoch=None):
|
|
step = step if step else len(self.losses)
|
|
if epoch is None and step >= 10**6:
|
|
epoch = f"_e{round(step/10**6,2)}M"
|
|
elif epoch is None:
|
|
epoch = f"_e{round(step/10**3)}K"
|
|
output_name = f"./{SAVE_FOLDER}/{self.name}{epoch}"
|
|
if self.vae:
|
|
out = self.eval_model()[1]
|
|
img = self.vae_decode(out).detach()
|
|
save_image(img, f"{output_name}.png")
|
|
torch.cuda.empty_cache()
|
|
save_file(self.model.state_dict(), f"{output_name}.safetensors")
|
|
torch.save(self.optimizer.state_dict(), f"{output_name}.optim.pth")
|
|
|
|
def close(self):
|
|
del self.vae
|
|
self.csvlog.close()
|
|
|
|
def vae_decode(self, latent):
|
|
latent = latent.to(torch.float16).to("cuda")
|
|
out = self.vae.decode(latent).sample
|
|
out = out.float().to(latent.device)
|
|
out = torch.clamp(out, min=-1.0, max=1.0)
|
|
return ((out + 1.0) / 2.0)
|
|
|
|
def get_vae(self, version, file_path=None, fp16=False):
|
|
"""Load VAE from file or default hf repo. fp16 only works from hf"""
|
|
vae = None
|
|
dtype = torch.float16 if fp16 else torch.float32
|
|
if version == "v1" and file_path:
|
|
vae = AutoencoderKL.from_single_file(
|
|
file_path,
|
|
image_size=512,
|
|
)
|
|
elif version == "v1":
|
|
vae = AutoencoderKL.from_pretrained(
|
|
"runwayml/stable-diffusion-v1-5",
|
|
subfolder="vae",
|
|
torch_dtype=dtype,
|
|
)
|
|
elif version == "xl" and file_path:
|
|
vae = AutoencoderKL.from_single_file(
|
|
file_path,
|
|
image_size=1024
|
|
)
|
|
elif version == "xl" and fp16:
|
|
vae = AutoencoderKL.from_pretrained(
|
|
"madebyollin/sdxl-vae-fp16-fix",
|
|
torch_dtype=torch.float16,
|
|
)
|
|
elif version == "xl":
|
|
vae = AutoencoderKL.from_pretrained(
|
|
"stabilityai/stable-diffusion-xl-base-1.0",
|
|
subfolder="vae"
|
|
)
|
|
else:
|
|
raise NotImplementedError(f"Unknown VAE version '{version}'")
|
|
|
|
# save VRAM
|
|
vae.to(dtype).to("cuda")
|
|
vae.decoder.eval()
|
|
vae.set_use_memory_efficient_attention_xformers(True)
|
|
vae.enable_xformers_memory_efficient_attention()
|
|
vae.enable_gradient_checkpointing()
|
|
del vae.encoder
|
|
return vae
|