|
|
"""
|
|
|
From https://github.com/CompVis/latent-diffusion/main/ldm/models/diffusion/ddpm.py
|
|
|
Pared down to simplify code.
|
|
|
|
|
|
The original file acknowledges:
|
|
|
https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
|
|
|
https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bbbd2d9902ce/improved_diffusion/gaussian_diffusion.py
|
|
|
https://github.com/CompVis/taming-transformers
|
|
|
"""
|
|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import numpy as np
|
|
|
import pytorch_lightning as pl
|
|
|
from contextlib import contextmanager
|
|
|
from functools import partial
|
|
|
from torchmetrics import MeanSquaredError
|
|
|
|
|
|
from .utils import make_beta_schedule, extract_into_tensor, noise_like, timestep_embedding
|
|
|
from .ema import LitEma
|
|
|
from ..blocks.afno import PatchEmbed3d, PatchExpand3d, AFNOBlock3d
|
|
|
|
|
|
|
|
|
class LatentDiffusion(pl.LightningModule):
|
|
|
def __init__(self,
|
|
|
model,
|
|
|
autoencoder,
|
|
|
context_encoder=None,
|
|
|
timesteps=1000,
|
|
|
beta_schedule="linear",
|
|
|
loss_type="l2",
|
|
|
use_ema=True,
|
|
|
lr=1e-4,
|
|
|
lr_warmup=0,
|
|
|
linear_start=1e-4,
|
|
|
linear_end=2e-2,
|
|
|
cosine_s=8e-3,
|
|
|
parameterization="eps",
|
|
|
):
|
|
|
super().__init__()
|
|
|
self.model = model
|
|
|
self.autoencoder = autoencoder.requires_grad_(False)
|
|
|
self.conditional = (context_encoder is not None)
|
|
|
self.context_encoder = context_encoder
|
|
|
self.lr = lr
|
|
|
self.lr_warmup = lr_warmup
|
|
|
|
|
|
self.val_loss = MeanSquaredError()
|
|
|
|
|
|
assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"'
|
|
|
self.parameterization = parameterization
|
|
|
|
|
|
self.use_ema = use_ema
|
|
|
if self.use_ema:
|
|
|
self.model_ema = LitEma(self.model)
|
|
|
|
|
|
self.register_schedule(
|
|
|
beta_schedule=beta_schedule, timesteps=timesteps,
|
|
|
linear_start=linear_start, linear_end=linear_end,
|
|
|
cosine_s=cosine_s
|
|
|
)
|
|
|
|
|
|
self.loss_type = loss_type
|
|
|
|
|
|
def register_schedule(self, beta_schedule="linear", timesteps=1000,
|
|
|
linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
|
|
|
|
|
betas = make_beta_schedule(
|
|
|
beta_schedule, timesteps,
|
|
|
linear_start=linear_start, linear_end=linear_end,
|
|
|
cosine_s=cosine_s
|
|
|
)
|
|
|
alphas = 1. - betas
|
|
|
alphas_cumprod = np.cumprod(alphas, axis=0)
|
|
|
alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
|
|
|
|
|
|
timesteps, = betas.shape
|
|
|
self.num_timesteps = int(timesteps)
|
|
|
self.linear_start = linear_start
|
|
|
self.linear_end = linear_end
|
|
|
assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
|
|
|
|
|
|
to_torch = partial(torch.tensor, dtype=torch.float32)
|
|
|
|
|
|
self.register_buffer('betas', to_torch(betas))
|
|
|
self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
|
|
|
self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
|
|
|
|
|
|
|
|
|
self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
|
|
|
self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
|
|
|
|
|
|
@contextmanager
|
|
|
def ema_scope(self, context=None):
|
|
|
if self.use_ema:
|
|
|
self.model_ema.store(self.model.parameters())
|
|
|
self.model_ema.copy_to(self.model)
|
|
|
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.model.parameters())
|
|
|
if context is not None:
|
|
|
print(f"{context}: Restored training weights")
|
|
|
|
|
|
def apply_model(self, x_noisy, t, cond=None, return_ids=False):
|
|
|
if self.conditional:
|
|
|
cond = self.context_encoder(cond)
|
|
|
with self.ema_scope():
|
|
|
return self.model(x_noisy, t, context=cond)
|
|
|
|
|
|
def q_sample(self, x_start, t, noise=None):
|
|
|
if noise is None:
|
|
|
noise = torch.randn_like(x_start)
|
|
|
return (
|
|
|
extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
|
|
|
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
|
|
|
)
|
|
|
|
|
|
def get_loss(self, pred, target, mean=True):
|
|
|
if self.loss_type == 'l1':
|
|
|
loss = (target - pred).abs()
|
|
|
if mean:
|
|
|
loss = loss.mean()
|
|
|
elif self.loss_type == 'l2':
|
|
|
if mean:
|
|
|
loss = torch.nn.functional.mse_loss(target, pred)
|
|
|
else:
|
|
|
loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
|
|
|
else:
|
|
|
raise NotImplementedError("unknown loss type '{loss_type}'")
|
|
|
|
|
|
return loss
|
|
|
|
|
|
def p_losses(self, x_start, t, noise=None, context=None):
|
|
|
if noise is None:
|
|
|
noise = torch.randn_like(x_start)
|
|
|
x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
|
|
|
model_out = self.model(x_noisy, t, context=context)
|
|
|
|
|
|
if self.parameterization == "eps":
|
|
|
target = noise
|
|
|
elif self.parameterization == "x0":
|
|
|
target = x_start
|
|
|
else:
|
|
|
raise NotImplementedError(f"Parameterization {self.parameterization} not yet supported")
|
|
|
|
|
|
return self.get_loss(model_out, target, mean=False).mean()
|
|
|
|
|
|
def forward(self, x, *args, **kwargs):
|
|
|
t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
|
|
|
return self.p_losses(x, t, *args, **kwargs)
|
|
|
|
|
|
def shared_step(self, batch):
|
|
|
(x,y) = batch
|
|
|
y = self.autoencoder.encode(y)[0]
|
|
|
context = self.context_encoder(x) if self.conditional else None
|
|
|
return self(y, context=context)
|
|
|
|
|
|
def training_step(self, batch, batch_idx):
|
|
|
loss = self.shared_step(batch)
|
|
|
self.log("train_loss", loss)
|
|
|
return loss
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def validation_step(self, batch, batch_idx):
|
|
|
|
|
|
|
|
|
|
|
|
loss = self.shared_step(batch)
|
|
|
with self.ema_scope():
|
|
|
loss_ema = self.shared_step(batch)
|
|
|
log_params = {"on_step": False, "on_epoch": True, "prog_bar": True}
|
|
|
self.log("val_loss", loss, **log_params)
|
|
|
self.log("val_loss_ema", loss, **log_params)
|
|
|
|
|
|
|
|
|
def test_step(self, batch, batch_idx):
|
|
|
return self.validation_step(batch, batch_idx)
|
|
|
|
|
|
def on_train_batch_end(self, *args, **kwargs):
|
|
|
if self.use_ema:
|
|
|
self.model_ema(self.model)
|
|
|
|
|
|
def configure_optimizers(self):
|
|
|
optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr,
|
|
|
betas=(0.5, 0.9), weight_decay=1e-3)
|
|
|
reduce_lr = torch.optim.lr_scheduler.ReduceLROnPlateau(
|
|
|
optimizer, patience=3, factor=0.25, verbose=True
|
|
|
)
|
|
|
return {
|
|
|
"optimizer": optimizer,
|
|
|
"lr_scheduler": {
|
|
|
"scheduler": reduce_lr,
|
|
|
"monitor": "val_loss_ema",
|
|
|
"frequency": 1,
|
|
|
},
|
|
|
}
|
|
|
|
|
|
def optimizer_step(
|
|
|
self,
|
|
|
epoch,
|
|
|
batch_idx,
|
|
|
optimizer,
|
|
|
optimizer_idx,
|
|
|
|
|
|
**kwargs
|
|
|
):
|
|
|
if self.trainer.global_step < self.lr_warmup:
|
|
|
lr_scale = (self.trainer.global_step+1) / self.lr_warmup
|
|
|
for pg in optimizer.param_groups:
|
|
|
pg['lr'] = lr_scale * self.lr
|
|
|
|
|
|
super().optimizer_step(
|
|
|
epoch, batch_idx, optimizer,
|
|
|
optimizer_idx,
|
|
|
|
|
|
**kwargs
|
|
|
)
|
|
|
|
|
|
|
