from omegaconf.dictconfig import DictConfig
from typing import List, Tuple
from ema_pytorch import EMA
import numpy as np
import torch
from torchtyping import TensorType
import torch.nn as nn
import lightning as L
from utils.random_utils import StackedRandomGenerator
# ------------------------------------------------------------------------------------- #
batch_size, num_samples = None, None
num_feats, num_rawfeats, num_cams = None, None, None
RawTrajectory = TensorType["num_samples", "num_rawfeats", "num_cams"]
# ------------------------------------------------------------------------------------- #
class Diffuser(L.LightningModule):
def __init__(
self,
network: nn.Module,
guidance_weight: float,
ema_kwargs: DictConfig,
sampling_kwargs: DictConfig,
edm2_normalization: bool,
**kwargs,
):
super().__init__()
# Network and EMA
self.net = network
self.ema = EMA(self.net, **ema_kwargs)
self.guidance_weight = guidance_weight
self.edm2_normalization = edm2_normalization
self.sigma_data = network.sigma_data
# Sampling
self.num_steps = sampling_kwargs.num_steps
self.sigma_min = sampling_kwargs.sigma_min
self.sigma_max = sampling_kwargs.sigma_max
self.rho = sampling_kwargs.rho
self.S_churn = sampling_kwargs.S_churn
self.S_noise = sampling_kwargs.S_noise
self.S_min = sampling_kwargs.S_min
self.S_max = (
sampling_kwargs.S_max
if isinstance(sampling_kwargs.S_max, float)
else float("inf")
)
# ---------------------------------------------------------------------------------- #
def on_predict_start(self):
eval_dataset = self.trainer.datamodule.eval_dataset
self.modalities = list(eval_dataset.modality_datasets.keys())
self.get_matrix = self.trainer.datamodule.train_dataset.get_matrix
self.v_get_matrix = self.trainer.datamodule.eval_dataset.get_matrix
def predict_step(self, batch, batch_idx):
ref_samples, mask = batch["traj_feat"], batch["padding_mask"]
if len(self.modalities) > 0:
cond_k = [x for x in batch.keys() if "traj" not in x and "feat" in x]
cond_data = [batch[cond] for cond in cond_k]
conds = {}
for cond in cond_k:
cond_name = cond.replace("_feat", "")
if isinstance(batch[f"{cond_name}_raw"], dict):
for cond_name_, x in batch[f"{cond_name}_raw"].items():
conds[cond_name_] = x
else:
conds[cond_name] = batch[f"{cond_name}_raw"]
batch["conds"] = conds
else:
cond_data = None
# cf edm2 sigma_data normalization / https://arxiv.org/pdf/2312.02696.pdf
if self.edm2_normalization:
ref_samples *= self.sigma_data
_, gen_samples = self.sample(self.ema.ema_model, ref_samples, cond_data, mask)
batch["ref_samples"] = torch.stack([self.v_get_matrix(x) for x in ref_samples])
batch["gen_samples"] = torch.stack([self.get_matrix(x) for x in gen_samples])
return batch
# --------------------------------------------------------------------------------- #
def sample(
self,
net: torch.nn.Module,
traj_samples: RawTrajectory,
cond_samples: TensorType["num_samples", "num_feats"],
mask: TensorType["num_samples", "num_feats"],
external_seeds: List[int] = None,
) -> Tuple[RawTrajectory, RawTrajectory]:
# Pick latents
num_samples = traj_samples.shape[0]
seeds = self.gen_seeds if hasattr(self, "gen_seeds") else range(num_samples)
rnd = StackedRandomGenerator(self.device, seeds)
sz = [num_samples, self.net.num_feats, self.net.num_cams]
latents = rnd.randn_rn(sz, device=self.device)
# Generate trajectories.
generations = self.edm_sampler(
net,
latents,
class_labels=cond_samples,
mask=mask,
randn_like=rnd.randn_like,
guidance_weight=self.guidance_weight,
# ----------------------------------- #
num_steps=self.num_steps,
sigma_min=self.sigma_min,
sigma_max=self.sigma_max,
rho=self.rho,
S_churn=self.S_churn,
S_min=self.S_min,
S_max=self.S_max,
S_noise=self.S_noise,
)
return latents, generations
@staticmethod
def edm_sampler(
net,
latents,
class_labels=None,
mask=None,
guidance_weight=2.0,
randn_like=torch.randn_like,
num_steps=18,
sigma_min=0.002,
sigma_max=80,
rho=7,
S_churn=0,
S_min=0,
S_max=float("inf"),
S_noise=1,
):
# Time step discretization.
step_indices = torch.arange(num_steps, device=latents.device)
t_steps = (
sigma_max ** (1 / rho)
+ step_indices
/ (num_steps - 1)
* (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))
) ** rho
t_steps = torch.cat(
[torch.as_tensor(t_steps), torch.zeros_like(t_steps[:1])]
) # t_N = 0
# Main sampling loop.
bool_mask = ~mask.to(bool)
x_next = latents * t_steps[0]
bs = latents.shape[0]
for i, (t_cur, t_next) in enumerate(
zip(t_steps[:-1], t_steps[1:])
): # 0, ..., N-1
x_cur = x_next
# Increase noise temporarily.
gamma = (
min(S_churn / num_steps, np.sqrt(2) - 1)
if S_min <= t_cur <= S_max
else 0
)
t_hat = torch.as_tensor(t_cur + gamma * t_cur)
x_hat = x_cur + (t_hat**2 - t_cur**2).sqrt() * S_noise * randn_like(x_cur)
# Euler step.
if class_labels is not None:
class_label_knot = [torch.zeros_like(label) for label in class_labels]
x_hat_both = torch.cat([x_hat, x_hat], dim=0)
y_label_both = [
torch.cat([y, y_knot], dim=0)
for y, y_knot in zip(class_labels, class_label_knot)
]
bool_mask_both = torch.cat([bool_mask, bool_mask], dim=0)
t_hat_both = torch.cat([t_hat.expand(bs), t_hat.expand(bs)], dim=0)
cond_denoised, denoised = net(
x_hat_both, t_hat_both, y=y_label_both, mask=bool_mask_both
).chunk(2, dim=0)
denoised = denoised + (cond_denoised - denoised) * guidance_weight
else:
denoised = net(x_hat, t_hat.expand(bs), mask=bool_mask)
d_cur = (x_hat - denoised) / t_hat
x_next = x_hat + (t_next - t_hat) * d_cur
# Apply 2nd order correction.
if i < num_steps - 1:
if class_labels is not None:
class_label_knot = [
torch.zeros_like(label) for label in class_labels
]
x_next_both = torch.cat([x_next, x_next], dim=0)
y_label_both = [
torch.cat([y, y_knot], dim=0)
for y, y_knot in zip(class_labels, class_label_knot)
]
bool_mask_both = torch.cat([bool_mask, bool_mask], dim=0)
t_next_both = torch.cat(
[t_next.expand(bs), t_next.expand(bs)], dim=0
)
cond_denoised, denoised = net(
x_next_both, t_next_both, y=y_label_both, mask=bool_mask_both
).chunk(2, dim=0)
denoised = denoised + (cond_denoised - denoised) * guidance_weight
else:
denoised = net(x_next, t_next.expand(bs), mask=bool_mask)
d_prime = (x_next - denoised) / t_next
x_next = x_hat + (t_next - t_hat) * (0.5 * d_cur + 0.5 * d_prime)
return x_next