Spaces:
Runtime error
Runtime error
# -------------------------------------------------------- | |
# Stable-Diffusion-Torch | |
# Based on Stable-Diffusion (https://github.com/CompVis/stable-diffusion) | |
# Removed Pytorch-lightning by Zigang Geng (zigang@mail.ustc.edu.cn) | |
# -------------------------------------------------------- | |
import torch | |
import torch.nn as nn | |
import torch.nn.functional as F | |
from contextlib import contextmanager | |
from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer | |
from ldm.modules.diffusionmodules.model import Encoder, Decoder | |
from ldm.modules.distributions.distributions import DiagonalGaussianDistribution | |
from ldm.util import instantiate_from_config | |
class VQModel(nn.Module): | |
def __init__(self, | |
ddconfig, | |
lossconfig, | |
n_embed, | |
embed_dim, | |
ckpt_path=None, | |
ignore_keys=[], | |
image_key="image", | |
colorize_nlabels=None, | |
monitor=None, | |
batch_resize_range=None, | |
scheduler_config=None, | |
lr_g_factor=1.0, | |
remap=None, | |
sane_index_shape=False, # tell vector quantizer to return indices as bhw | |
use_ema=False | |
): | |
super().__init__() | |
self.embed_dim = embed_dim | |
self.n_embed = n_embed | |
self.image_key = image_key | |
self.encoder = Encoder(**ddconfig) | |
self.decoder = Decoder(**ddconfig) | |
self.loss = instantiate_from_config(lossconfig) | |
self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25, | |
remap=remap, | |
sane_index_shape=sane_index_shape) | |
self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1) | |
self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) | |
if colorize_nlabels is not None: | |
assert type(colorize_nlabels)==int | |
self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) | |
if monitor is not None: | |
self.monitor = monitor | |
self.batch_resize_range = batch_resize_range | |
if self.batch_resize_range is not None: | |
print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.") | |
self.use_ema = use_ema | |
if self.use_ema: | |
self.model_ema = LitEma(self) | |
print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") | |
if ckpt_path is not None: | |
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) | |
self.scheduler_config = scheduler_config | |
self.lr_g_factor = lr_g_factor | |
def ema_scope(self, context=None): | |
if self.use_ema: | |
self.model_ema.store(self.parameters()) | |
self.model_ema.copy_to(self) | |
if context is not None: | |
print(f"{context}: Switched to EMA weights") | |
try: | |
yield None | |
finally: | |
if self.use_ema: | |
self.model_ema.restore(self.parameters()) | |
if context is not None: | |
print(f"{context}: Restored training weights") | |
def init_from_ckpt(self, path, ignore_keys=list()): | |
sd = torch.load(path, map_location="cpu")["state_dict"] | |
keys = list(sd.keys()) | |
for k in keys: | |
for ik in ignore_keys: | |
if k.startswith(ik): | |
print("Deleting key {} from state_dict.".format(k)) | |
del sd[k] | |
missing, unexpected = self.load_state_dict(sd, strict=False) | |
print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") | |
if len(missing) > 0: | |
print(f"Missing Keys: {missing}") | |
print(f"Unexpected Keys: {unexpected}") | |
def on_train_batch_end(self, *args, **kwargs): | |
if self.use_ema: | |
self.model_ema(self) | |
def encode(self, x): | |
h = self.encoder(x) | |
h = self.quant_conv(h) | |
quant, emb_loss, info = self.quantize(h) | |
return quant, emb_loss, info | |
def encode_to_prequant(self, x): | |
h = self.encoder(x) | |
h = self.quant_conv(h) | |
return h | |
def decode(self, quant): | |
quant = self.post_quant_conv(quant) | |
dec = self.decoder(quant) | |
return dec | |
def decode_code(self, code_b): | |
quant_b = self.quantize.embed_code(code_b) | |
dec = self.decode(quant_b) | |
return dec | |
def forward(self, input, return_pred_indices=False): | |
quant, diff, (_,_,ind) = self.encode(input) | |
dec = self.decode(quant) | |
if return_pred_indices: | |
return dec, diff, ind | |
return dec, diff | |
def get_input(self, batch, k): | |
x = batch[k] | |
if len(x.shape) == 3: | |
x = x[..., None] | |
x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() | |
if self.batch_resize_range is not None: | |
lower_size = self.batch_resize_range[0] | |
upper_size = self.batch_resize_range[1] | |
if self.global_step <= 4: | |
# do the first few batches with max size to avoid later oom | |
new_resize = upper_size | |
else: | |
new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16)) | |
if new_resize != x.shape[2]: | |
x = F.interpolate(x, size=new_resize, mode="bicubic") | |
x = x.detach() | |
return x | |
def training_step(self, batch, batch_idx, optimizer_idx): | |
# https://github.com/pytorch/pytorch/issues/37142 | |
# try not to fool the heuristics | |
x = self.get_input(batch, self.image_key) | |
xrec, qloss, ind = self(x, return_pred_indices=True) | |
if optimizer_idx == 0: | |
# autoencode | |
aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, | |
last_layer=self.get_last_layer(), split="train", | |
predicted_indices=ind) | |
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) | |
return aeloss | |
if optimizer_idx == 1: | |
# discriminator | |
discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, | |
last_layer=self.get_last_layer(), split="train") | |
self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True) | |
return discloss | |
def validation_step(self, batch, batch_idx): | |
log_dict = self._validation_step(batch, batch_idx) | |
with self.ema_scope(): | |
log_dict_ema = self._validation_step(batch, batch_idx, suffix="_ema") | |
return log_dict | |
def _validation_step(self, batch, batch_idx, suffix=""): | |
x = self.get_input(batch, self.image_key) | |
xrec, qloss, ind = self(x, return_pred_indices=True) | |
aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, | |
self.global_step, | |
last_layer=self.get_last_layer(), | |
split="val"+suffix, | |
predicted_indices=ind | |
) | |
discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, | |
self.global_step, | |
last_layer=self.get_last_layer(), | |
split="val"+suffix, | |
predicted_indices=ind | |
) | |
rec_loss = log_dict_ae[f"val{suffix}/rec_loss"] | |
self.log(f"val{suffix}/rec_loss", rec_loss, | |
prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) | |
self.log(f"val{suffix}/aeloss", aeloss, | |
prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) | |
# if version.parse(pl.__version__) >= version.parse('1.4.0'): | |
# del log_dict_ae[f"val{suffix}/rec_loss"] | |
self.log_dict(log_dict_ae) | |
self.log_dict(log_dict_disc) | |
return self.log_dict | |
def configure_optimizers(self): | |
lr_d = self.learning_rate | |
lr_g = self.lr_g_factor*self.learning_rate | |
print("lr_d", lr_d) | |
print("lr_g", lr_g) | |
opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ | |
list(self.decoder.parameters())+ | |
list(self.quantize.parameters())+ | |
list(self.quant_conv.parameters())+ | |
list(self.post_quant_conv.parameters()), | |
lr=lr_g, betas=(0.5, 0.9)) | |
opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), | |
lr=lr_d, betas=(0.5, 0.9)) | |
if self.scheduler_config is not None: | |
scheduler = instantiate_from_config(self.scheduler_config) | |
print("Setting up LambdaLR scheduler...") | |
scheduler = [ | |
{ | |
'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule), | |
'interval': 'step', | |
'frequency': 1 | |
}, | |
{ | |
'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule), | |
'interval': 'step', | |
'frequency': 1 | |
}, | |
] | |
return [opt_ae, opt_disc], scheduler | |
return [opt_ae, opt_disc], [] | |
def get_last_layer(self): | |
return self.decoder.conv_out.weight | |
def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs): | |
log = dict() | |
x = self.get_input(batch, self.image_key) | |
x = x.to(batch.device) | |
if only_inputs: | |
log["inputs"] = x | |
return log | |
xrec, _ = self(x) | |
if x.shape[1] > 3: | |
# colorize with random projection | |
assert xrec.shape[1] > 3 | |
x = self.to_rgb(x) | |
xrec = self.to_rgb(xrec) | |
log["inputs"] = x | |
log["reconstructions"] = xrec | |
if plot_ema: | |
with self.ema_scope(): | |
xrec_ema, _ = self(x) | |
if x.shape[1] > 3: xrec_ema = self.to_rgb(xrec_ema) | |
log["reconstructions_ema"] = xrec_ema | |
return log | |
def to_rgb(self, x): | |
assert self.image_key == "segmentation" | |
if not hasattr(self, "colorize"): | |
self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) | |
x = F.conv2d(x, weight=self.colorize) | |
x = 2.*(x-x.min())/(x.max()-x.min()) - 1. | |
return x | |
class VQModelInterface(VQModel): | |
def __init__(self, embed_dim, *args, **kwargs): | |
super().__init__(embed_dim=embed_dim, *args, **kwargs) | |
self.embed_dim = embed_dim | |
def encode(self, x): | |
h = self.encoder(x) | |
h = self.quant_conv(h) | |
return h | |
def decode(self, h, force_not_quantize=False): | |
# also go through quantization layer | |
if not force_not_quantize: | |
quant, emb_loss, info = self.quantize(h) | |
else: | |
quant = h | |
quant = self.post_quant_conv(quant) | |
dec = self.decoder(quant) | |
return dec | |
class AutoencoderKL(nn.Module): | |
def __init__(self, | |
ddconfig, | |
lossconfig, | |
embed_dim, | |
ckpt_path=None, | |
ignore_keys=[], | |
image_key="image", | |
colorize_nlabels=None, | |
monitor=None, | |
): | |
super().__init__() | |
self.image_key = image_key | |
self.encoder = Encoder(**ddconfig) | |
self.decoder = Decoder(**ddconfig) | |
self.loss = instantiate_from_config(lossconfig) | |
assert ddconfig["double_z"] | |
self.quant_conv = torch.nn.Conv2d(2*ddconfig["z_channels"], 2*embed_dim, 1) | |
self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) | |
self.embed_dim = embed_dim | |
if colorize_nlabels is not None: | |
assert type(colorize_nlabels)==int | |
self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) | |
if monitor is not None: | |
self.monitor = monitor | |
if ckpt_path is not None: | |
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) | |
def init_from_ckpt(self, path, ignore_keys=list()): | |
sd = torch.load(path, map_location="cpu")["state_dict"] | |
keys = list(sd.keys()) | |
for k in keys: | |
for ik in ignore_keys: | |
if k.startswith(ik): | |
print("Deleting key {} from state_dict.".format(k)) | |
del sd[k] | |
self.load_state_dict(sd, strict=False) | |
print(f"Restored from {path}") | |
def encode(self, x): | |
h = self.encoder(x) | |
moments = self.quant_conv(h) | |
posterior = DiagonalGaussianDistribution(moments) | |
return posterior | |
def decode(self, z): | |
z = self.post_quant_conv(z.type(self.post_quant_conv.weight.dtype)) | |
dec = self.decoder(z) | |
return dec | |
def forward(self, input, sample_posterior=True): | |
posterior = self.encode(input) | |
if sample_posterior: | |
z = posterior.sample() | |
else: | |
z = posterior.mode() | |
dec = self.decode(z) | |
return dec, posterior | |
def get_input(self, batch, k): | |
x = batch[k] | |
if len(x.shape) == 3: | |
x = x[..., None] | |
x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() | |
return x | |
def training_step(self, batch, batch_idx, optimizer_idx): | |
inputs = self.get_input(batch, self.image_key) | |
reconstructions, posterior = self(inputs) | |
if optimizer_idx == 0: | |
# train encoder+decoder+logvar | |
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, | |
last_layer=self.get_last_layer(), split="train") | |
self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) | |
self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False) | |
return aeloss | |
if optimizer_idx == 1: | |
# train the discriminator | |
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step, | |
last_layer=self.get_last_layer(), split="train") | |
self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True) | |
self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False) | |
return discloss | |
def validation_step(self, batch, batch_idx): | |
inputs = self.get_input(batch, self.image_key) | |
reconstructions, posterior = self(inputs) | |
aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step, | |
last_layer=self.get_last_layer(), split="val") | |
discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step, | |
last_layer=self.get_last_layer(), split="val") | |
self.log("val/rec_loss", log_dict_ae["val/rec_loss"]) | |
self.log_dict(log_dict_ae) | |
self.log_dict(log_dict_disc) | |
return self.log_dict | |
def configure_optimizers(self): | |
lr = self.learning_rate | |
opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ | |
list(self.decoder.parameters())+ | |
list(self.quant_conv.parameters())+ | |
list(self.post_quant_conv.parameters()), | |
lr=lr, betas=(0.5, 0.9)) | |
opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), | |
lr=lr, betas=(0.5, 0.9)) | |
return [opt_ae, opt_disc], [] | |
def get_last_layer(self): | |
return self.decoder.conv_out.weight | |
def log_images(self, batch, only_inputs=False, **kwargs): | |
log = dict() | |
x = self.get_input(batch, self.image_key) | |
x = x.to(batch.device) | |
if not only_inputs: | |
xrec, posterior = self(x) | |
if x.shape[1] > 3: | |
# colorize with random projection | |
assert xrec.shape[1] > 3 | |
x = self.to_rgb(x) | |
xrec = self.to_rgb(xrec) | |
log["samples"] = self.decode(torch.randn_like(posterior.sample())) | |
log["reconstructions"] = xrec | |
log["inputs"] = x | |
return log | |
def to_rgb(self, x): | |
assert self.image_key == "segmentation" | |
if not hasattr(self, "colorize"): | |
self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) | |
x = F.conv2d(x, weight=self.colorize) | |
x = 2.*(x-x.min())/(x.max()-x.min()) - 1. | |
return x | |
class IdentityFirstStage(torch.nn.Module): | |
def __init__(self, *args, vq_interface=False, **kwargs): | |
self.vq_interface = vq_interface # TODO: Should be true by default but check to not break older stuff | |
super().__init__() | |
def encode(self, x, *args, **kwargs): | |
return x | |
def decode(self, x, *args, **kwargs): | |
return x | |
def quantize(self, x, *args, **kwargs): | |
if self.vq_interface: | |
return x, None, [None, None, None] | |
return x | |
def forward(self, x, *args, **kwargs): | |
return x | |