# ---------------------------------------------------------------
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# This work is licensed under the NVIDIA Source Code License
# for LSGM. To view a copy of this license, see the LICENSE file.
# ---------------------------------------------------------------
from pdb import set_trace as st
from abc import ABC, abstractmethod
import numpy as np
import torch
import gc
from .continuous_distributions import log_p_standard_normal, log_p_var_normal
from .continuous_diffusion_utils import trace_df_dx_hutchinson, sample_gaussian_like, sample_rademacher_like, get_mixed_prediction
from torchdiffeq import odeint
from torch.cuda.amp import autocast
from timeit import default_timer as timer
from guided_diffusion import dist_util, logger
def make_diffusion(args):
""" simple diffusion factory function to return diffusion instances. Only use this to create continuous diffusions """
if args.sde_sde_type == 'geometric_sde':
return DiffusionGeometric(args)
elif args.sde_sde_type == 'vpsde':
return DiffusionVPSDE(args)
elif args.sde_sde_type == 'sub_vpsde':
return DiffusionSubVPSDE(args)
elif args.sde_sde_type == 'vesde':
return DiffusionVESDE(args)
else:
raise ValueError("Unrecognized sde type: {}".format(args.sde_sde_type))
class DiffusionBase(ABC):
"""
Abstract base class for all diffusion implementations.
"""
def __init__(self, args):
super().__init__()
self.args = args
self.sigma2_0 = args.sde_sigma2_0
self.sde_type = args.sde_sde_type
@abstractmethod
def f(self, t):
""" returns the drift coefficient at time t: f(t) """
pass
@abstractmethod
def g2(self, t):
""" returns the squared diffusion coefficient at time t: g^2(t) """
pass
@abstractmethod
def var(self, t):
""" returns variance at time t, \sigma_t^2"""
pass
@abstractmethod
def e2int_f(self, t):
""" returns e^{\int_0^t f(s) ds} which corresponds to the coefficient of mean at time t. """
pass
@abstractmethod
def inv_var(self, var):
""" inverse of the variance function at input variance var. """
pass
@abstractmethod
def mixing_component(self, x_noisy, var_t, t, enabled):
""" returns mixing component which is the optimal denoising model assuming that q(z_0) is N(0, 1) """
pass
def sample_q(self, x_init, noise, var_t, m_t):
""" returns a sample from diffusion process at time t """
return m_t * x_init + torch.sqrt(var_t) * noise
def log_snr(self, m_t, var_t):
return torch.log((torch.square(m_t) / var_t))
def _predict_x0_from_eps(self, z, eps, logsnr):
"""eps = (z - alpha * x0) / sigma
"""
return torch.sqrt(1 + torch.exp(-logsnr)) * (
z - eps * torch.rsqrt(1 + torch.exp(logsnr)))
def _predict_eps_from_x0(self, z, x0, logsnr):
"""x = (z - sigma * eps) / alpha
"""
return torch.sqrt(1 + torch.exp(logsnr)) * (
z - x0 * torch.rsqrt(1 + torch.exp(-logsnr)))
def _predict_eps_from_z_and_v(self, v_t, var_t, z, m_t):
# TODO, use logsnr here?
return torch.sqrt(var_t) * z + m_t * v_t
def _predict_x0_from_z_and_v(self, v_t, var_t, z, m_t):
return torch.sqrt(var_t) * v_t + m_t * z
def cross_entropy_const(self, ode_eps):
""" returns cross entropy factor with variance according to ode integration cutoff ode_eps """
# _, c, h, w = x_init.shape
return 0.5 * (1.0 + torch.log(2.0 * np.pi * self.var(
t=torch.tensor(ode_eps, device=dist_util.dev()))))
def compute_ode_nll(self, dae, eps, ode_eps, ode_solver_tol,
enable_autocast, no_autograd, num_samples, report_std):
""" calculates NLL based on ODE framework, assuming integration cutoff ode_eps """
# ODE solver starts consuming the CPU memory without this on large models
# https://github.com/scipy/scipy/issues/10070
gc.collect()
dae.eval()
def ode_func(t, state):
""" the ode function (including log probability integration for NLL calculation) """
global nfe_counter
nfe_counter = nfe_counter + 1
x = state[0].detach()
x.requires_grad_(True)
noise = sample_gaussian_like(
x) # could also use rademacher noise (sample_rademacher_like)
with torch.set_grad_enabled(True):
with autocast(enabled=enable_autocast):
variance = self.var(t=t)
mixing_component = self.mixing_component(
x_noisy=x,
var_t=variance,
t=t,
enabled=dae.mixed_prediction)
pred_params = dae(x=x, t=t)
params = get_mixed_prediction(dae.mixed_prediction,
pred_params,
dae.mixing_logit,
mixing_component)
dx_dt = self.f(t=t) * x + 0.5 * self.g2(
t=t) * params / torch.sqrt(variance)
with autocast(enabled=False):
dlogp_x_dt = -trace_df_dx_hutchinson(
dx_dt, x, noise, no_autograd).view(x.shape[0], 1)
return (dx_dt, dlogp_x_dt)
# NFE counter
global nfe_counter
nll_all, nfe_all = [], []
for i in range(num_samples):
# integrated log probability
logp_diff_t0 = torch.zeros(eps.shape[0], 1, device=dist_util.dev())
nfe_counter = 0
# solve the ODE
x_t, logp_diff_t = odeint(
ode_func,
(eps, logp_diff_t0),
torch.tensor([ode_eps, 1.0], device=dist_util.dev()),
atol=ode_solver_tol,
rtol=ode_solver_tol,
method="scipy_solver",
options={"solver": 'RK45'},
)
# last output values
x_t0, logp_diff_t0 = x_t[-1], logp_diff_t[-1]
# prior
if self.sde_type == 'vesde':
logp_prior = torch.sum(log_p_var_normal(x_t0,
var=self.sigma2_max),
dim=[1, 2, 3])
else:
logp_prior = torch.sum(log_p_standard_normal(x_t0),
dim=[1, 2, 3])
log_likelihood = logp_prior - logp_diff_t0.view(-1)
nll_all.append(-log_likelihood)
nfe_all.append(nfe_counter)
nfe_mean = np.mean(nfe_all)
nll_all = torch.stack(nll_all, dim=1)
nll_mean = torch.mean(nll_all, dim=1)
if num_samples > 1 and report_std:
nll_stddev = torch.std(nll_all, dim=1)
nll_stddev_batch = torch.mean(nll_stddev)
nll_stderror_batch = nll_stddev_batch / np.sqrt(num_samples)
else:
nll_stddev_batch = None
nll_stderror_batch = None
return nll_mean, nfe_mean, nll_stddev_batch, nll_stderror_batch
def sample_model_ode(self,
dae,
num_samples,
shape,
ode_eps,
ode_solver_tol,
enable_autocast,
temp,
noise=None):
""" generates samples using the ODE framework, assuming integration cutoff ode_eps """
# ODE solver starts consuming the CPU memory without this on large models
# https://github.com/scipy/scipy/issues/10070
gc.collect()
dae.eval()
def ode_func(t, x):
""" the ode function (sampling only, no NLL stuff) """
global nfe_counter
nfe_counter = nfe_counter + 1
with autocast(enabled=enable_autocast):
variance = self.var(t=t)
mixing_component = self.mixing_component(
x_noisy=x,
var_t=variance,
t=t,
enabled=dae.mixed_prediction)
pred_params = dae(x=x, t=t)
params = get_mixed_prediction(dae.mixed_prediction,
pred_params, dae.mixing_logit,
mixing_component)
dx_dt = self.f(t=t) * x + 0.5 * self.g2(
t=t) * params / torch.sqrt(variance)
return dx_dt
# the initial noise
if noise is None:
noise = torch.randn(size=[num_samples] + shape,
device=dist_util.dev())
if self.sde_type == 'vesde':
noise_init = temp * noise * np.sqrt(self.sigma2_max)
else:
noise_init = temp * noise
# NFE counter
global nfe_counter
nfe_counter = 0
# solve the ODE
start = timer()
samples_out = odeint(
ode_func,
noise_init,
torch.tensor([1.0, ode_eps], device=dist_util.dev()),
atol=ode_solver_tol,
rtol=ode_solver_tol,
method="scipy_solver",
options={"solver": 'RK45'},
)
end = timer()
ode_solve_time = end - start
return samples_out[-1], nfe_counter, ode_solve_time
# def iw_quantities(self, size, time_eps, iw_sample_mode, iw_subvp_like_vp_sde):
def iw_quantities(self, iw_sample_mode, size=None):
args = self.args
time_eps, iw_subvp_like_vp_sde = args.sde_time_eps, args.iw_subvp_like_vp_sde
if size is None:
size = args.batch_size
if self.sde_type in ['geometric_sde', 'vpsde']:
return self._iw_quantities_vpsdelike(size, time_eps,
iw_sample_mode)
elif self.sde_type in ['sub_vpsde']:
return self._iw_quantities_subvpsdelike(size, time_eps,
iw_sample_mode,
iw_subvp_like_vp_sde)
elif self.sde_type in ['vesde']:
return self._iw_quantities_vesde(size, time_eps, iw_sample_mode)
else:
raise NotImplementedError
def _iw_quantities_vpsdelike(self, size, time_eps, iw_sample_mode):
"""
For all SDEs where the underlying SDE is of the form dz = -0.5 * beta(t) * z * dt + sqrt{beta(t)} * dw, like
for the VPSDE.
"""
rho = torch.rand(size=[size], device=dist_util.dev())
# In the following, obj_weight_t corresponds to the weight in front of the l2 loss for the given iw_sample_mode.
# obj_weight_t_ll corresponds to the weight that converts the weighting scheme in iw_sample_mode to likelihood
# weighting.
if iw_sample_mode == 'll_uniform':
# uniform t sampling - likelihood obj. for both q and p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = obj_weight_t_ll = g2_t / (2.0 * var_t)
elif iw_sample_mode == 'll_iw': # ! q-obj
# importance sampling for likelihood obj. - likelihood obj. for both q and p
ones = torch.ones_like(rho, device=dist_util.dev())
sigma2_1, sigma2_eps = self.var(ones), self.var(time_eps * ones)
log_sigma2_1, log_sigma2_eps = torch.log(sigma2_1), torch.log(
sigma2_eps)
var_t = torch.exp(rho * log_sigma2_1 +
(1 - rho) * log_sigma2_eps) # sigma square
t = self.inv_var(var_t)
m_t, g2_t = self.e2int_f(t), self.g2(t) # m_t is alpha_bar
obj_weight_t = obj_weight_t_ll = 0.5 * (
log_sigma2_1 - log_sigma2_eps) / (1.0 - var_t)
elif iw_sample_mode == 'drop_all_uniform':
# uniform t sampling - likelihood obj. for q, all-prefactors-dropped obj. for p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = torch.ones(1, device=dist_util.dev())
obj_weight_t_ll = g2_t / (2.0 * var_t)
elif iw_sample_mode == 'drop_all_iw':
# importance sampling for all-pref.-dropped obj. - likelihood obj. for q, all-pref.-dropped obj. for p
assert self.sde_type == 'vpsde', 'Importance sampling for fully unweighted objective is currently only ' \
'implemented for the regular VPSDE.'
t = torch.sqrt(1.0 / self.delta_beta_half) * torch.erfinv(
rho * self.const_norm_2 + self.const_erf) - self.beta_frac
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = self.const_norm / (1.0 - var_t)
obj_weight_t_ll = obj_weight_t * g2_t / (2.0 * var_t)
elif iw_sample_mode == 'drop_sigma2t_iw': # ! default mode for p
# importance sampling for inv_sigma2_t-dropped obj. - likelihood obj. for q, inv_sigma2_t-dropped obj. for p
ones = torch.ones_like(rho, device=dist_util.dev())
sigma2_1, sigma2_eps = self.var(ones), self.var(time_eps * ones)
var_t = rho * sigma2_1 + (1 - rho) * sigma2_eps # ! sigma square
t = self.inv_var(var_t)
m_t, g2_t = self.e2int_f(t), self.g2(t) # ! m_t: alpha_bar sqrt
obj_weight_t = 0.5 * (sigma2_1 - sigma2_eps) / (1.0 - var_t)
obj_weight_t_ll = obj_weight_t / var_t
elif iw_sample_mode == 'drop_sigma2t_uniform':
# uniform sampling for inv_sigma2_t-dropped obj. - likelihood obj. for q, inv_sigma2_t-dropped obj. for p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = g2_t / 2.0
obj_weight_t_ll = g2_t / (2.0 * var_t)
elif iw_sample_mode == 'rescale_iw':
# importance sampling for 1/(1-sigma2_t) resc. obj. - likelihood obj. for q, 1/(1-sigma2_t) resc. obj. for p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = 0.5 / (1.0 - var_t)
obj_weight_t_ll = g2_t / (2.0 * var_t)
else:
raise ValueError(
"Unrecognized importance sampling type: {}".format(
iw_sample_mode))
return t, var_t.view(-1, 1, 1, 1), m_t.view(-1, 1, 1, 1), obj_weight_t.view(-1, 1, 1, 1), \
obj_weight_t_ll.view(-1, 1, 1, 1), g2_t.view(-1, 1, 1, 1)
def _iw_quantities_subvpsdelike(self, size, time_eps, iw_sample_mode,
iw_subvp_like_vp_sde):
"""
For all SDEs where the underlying SDE is of the form
dz = -0.5 * beta(t) * z * dt + sqrt{beta(t) * (1 - exp[-2 * betaintegral])} * dw, like for the Sub-VPSDE.
When iw_subvp_like_vp_sde is True, then we define the importance sampling distributions based on an analogous
VPSDE, while stile using the Sub-VPSDE. The motivation is that deriving the correct importance sampling
distributions for the Sub-VPSDE itself is hard, but the importance sampling distributions from analogous VPSDEs
probably already significantly reduce the variance also for the Sub-VPSDE.
"""
rho = torch.rand(size=[size], device=dist_util.dev())
# In the following, obj_weight_t corresponds to the weight in front of the l2 loss for the given iw_sample_mode.
# obj_weight_t_ll corresponds to the weight that converts the weighting scheme in iw_sample_mode to likelihood
# weighting.
if iw_sample_mode == 'll_uniform':
# uniform t sampling - likelihood obj. for both q and p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = obj_weight_t_ll = g2_t / (2.0 * var_t)
elif iw_sample_mode == 'll_iw':
if iw_subvp_like_vp_sde:
# importance sampling for vpsde likelihood obj. - sub-vpsde likelihood obj. for both q and p
ones = torch.ones_like(rho, device=dist_util.dev())
sigma2_1, sigma2_eps = self.var_vpsde(ones), self.var_vpsde(
time_eps * ones)
log_sigma2_1, log_sigma2_eps = torch.log(sigma2_1), torch.log(
sigma2_eps)
var_t_vpsde = torch.exp(rho * log_sigma2_1 +
(1 - rho) * log_sigma2_eps)
t = self.inv_var_vpsde(var_t_vpsde)
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = obj_weight_t_ll = g2_t / (2.0 * var_t) * \
(log_sigma2_1 - log_sigma2_eps) * var_t_vpsde / (1 - var_t_vpsde) / self.beta(t)
else:
raise NotImplementedError
elif iw_sample_mode == 'drop_all_uniform':
# uniform t sampling - likelihood obj. for q, all-prefactors-dropped obj. for p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = torch.ones(1, device=dist_util.dev())
obj_weight_t_ll = g2_t / (2.0 * var_t)
elif iw_sample_mode == 'drop_all_iw':
if iw_subvp_like_vp_sde:
# importance sampling for all-pref.-dropped obj. - likelihood obj. for q, all-pref.-dropped obj. for p
assert self.sde_type == 'sub_vpsde', 'Importance sampling for fully unweighted objective is ' \
'currently only implemented for the Sub-VPSDE.'
t = torch.sqrt(1.0 / self.delta_beta_half) * torch.erfinv(
rho * self.const_norm_2 + self.const_erf) - self.beta_frac
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = self.const_norm / (1.0 - self.var_vpsde(t))
obj_weight_t_ll = obj_weight_t * g2_t / (2.0 * var_t)
else:
raise NotImplementedError
elif iw_sample_mode == 'drop_sigma2t_iw':
if iw_subvp_like_vp_sde:
# importance sampling for inv_sigma2_t-dropped obj. - likelihood obj. for q, inv_sigma2_t-dropped obj. for p
ones = torch.ones_like(rho, device=dist_util.dev())
sigma2_1, sigma2_eps = self.var_vpsde(ones), self.var_vpsde(
time_eps * ones)
var_t_vpsde = rho * sigma2_1 + (1 - rho) * sigma2_eps
t = self.inv_var_vpsde(var_t_vpsde)
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = 0.5 * g2_t / self.beta(t) * (
sigma2_1 - sigma2_eps) / (1.0 - var_t_vpsde)
obj_weight_t_ll = obj_weight_t / var_t
else:
raise NotImplementedError
elif iw_sample_mode == 'drop_sigma2t_uniform':
# uniform sampling for inv_sigma2_t-dropped obj. - likelihood obj. for q, inv_sigma2_t-dropped obj. for p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = g2_t / 2.0
obj_weight_t_ll = g2_t / (2.0 * var_t)
elif iw_sample_mode == 'rescale_iw':
# importance sampling for 1/(1-sigma2_t) resc. obj. - likelihood obj. for q, 1/(1-sigma2_t) resc. obj. for p
# Note that we use the sub-vpsde variance to scale the p objective! It's not clear what's optimal here!
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = 0.5 / (1.0 - var_t)
obj_weight_t_ll = g2_t / (2.0 * var_t)
else:
raise ValueError(
"Unrecognized importance sampling type: {}".format(
iw_sample_mode))
return t, var_t.view(-1, 1, 1, 1), m_t.view(-1, 1, 1, 1), obj_weight_t.view(-1, 1, 1, 1), \
obj_weight_t_ll.view(-1, 1, 1, 1), g2_t.view(-1, 1, 1, 1)
def _iw_quantities_vesde(self, size, time_eps, iw_sample_mode):
"""
For the VESDE.
"""
rho = torch.rand(size=[size], device=dist_util.dev())
# In the following, obj_weight_t corresponds to the weight in front of the l2 loss for the given iw_sample_mode.
# obj_weight_t_ll corresponds to the weight that converts the weighting scheme in iw_sample_mode to likelihood
# weighting.
if iw_sample_mode == 'll_uniform':
# uniform t sampling - likelihood obj. for both q and p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = obj_weight_t_ll = g2_t / (2.0 * var_t)
elif iw_sample_mode == 'll_iw':
# importance sampling for likelihood obj. - likelihood obj. for both q and p
ones = torch.ones_like(rho, device=dist_util.dev())
nsigma2_1, nsigma2_eps, sigma2_eps = self.var_N(ones), self.var_N(
time_eps * ones), self.var(time_eps * ones)
log_frac_sigma2_1, log_frac_sigma2_eps = torch.log(
self.sigma2_max / nsigma2_1), torch.log(nsigma2_eps /
sigma2_eps)
var_N_t = (1.0 - self.sigma2_min) / (
1.0 - torch.exp(rho *
(log_frac_sigma2_1 + log_frac_sigma2_eps) -
log_frac_sigma2_eps))
t = self.inv_var_N(var_N_t)
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = obj_weight_t_ll = 0.5 * (
log_frac_sigma2_1 +
log_frac_sigma2_eps) * self.var_N(t) / (1.0 - self.sigma2_min)
elif iw_sample_mode == 'drop_all_uniform':
# uniform t sampling - likelihood obj. for q, all-prefactors-dropped obj. for p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = torch.ones(1, device=dist_util.dev())
obj_weight_t_ll = g2_t / (2.0 * var_t)
elif iw_sample_mode == 'drop_all_iw':
# importance sampling for all-pref.-dropped obj. - likelihood obj. for q, all-pref.-dropped obj. for p
ones = torch.ones_like(rho, device=dist_util.dev())
nsigma2_1, nsigma2_eps, sigma2_eps = self.var_N(ones), self.var_N(
time_eps * ones), self.var(time_eps * ones)
log_frac_sigma2_1, log_frac_sigma2_eps = torch.log(
self.sigma2_max / nsigma2_1), torch.log(nsigma2_eps /
sigma2_eps)
var_N_t = (1.0 - self.sigma2_min) / (
1.0 - torch.exp(rho *
(log_frac_sigma2_1 + log_frac_sigma2_eps) -
log_frac_sigma2_eps))
t = self.inv_var_N(var_N_t)
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t_ll = 0.5 * (log_frac_sigma2_1 +
log_frac_sigma2_eps) * self.var_N(t) / (
1.0 - self.sigma2_min)
obj_weight_t = 2.0 * obj_weight_t_ll / np.log(
self.sigma2_max / self.sigma2_min)
elif iw_sample_mode == 'drop_sigma2t_iw':
# importance sampling for inv_sigma2_t-dropped obj. - likelihood obj. for q, inv_sigma2_t-dropped obj. for p
ones = torch.ones_like(rho, device=dist_util.dev())
nsigma2_1, nsigma2_eps = self.var_N(ones), self.var_N(time_eps *
ones)
var_N_t = torch.exp(rho * torch.log(nsigma2_1) +
(1 - rho) * torch.log(nsigma2_eps))
t = self.inv_var_N(var_N_t)
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = 0.5 * torch.log(
nsigma2_1 / nsigma2_eps) * self.var_N(t)
obj_weight_t_ll = obj_weight_t / var_t
elif iw_sample_mode == 'drop_sigma2t_uniform':
# uniform sampling for inv_sigma2_t-dropped obj. - likelihood obj. for q, inv_sigma2_t-dropped obj. for p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = g2_t / 2.0
obj_weight_t_ll = g2_t / (2.0 * var_t)
elif iw_sample_mode == 'rescale_iw':
# uniform sampling for 1/(1-sigma2_t) resc. obj. - likelihood obj. for q, 1/(1-sigma2_t) resc. obj. for p
t = rho * (1. - time_eps) + time_eps
var_t, m_t, g2_t = self.var(t), self.e2int_f(t), self.g2(t)
obj_weight_t = 0.5 / (1.0 - var_t)
obj_weight_t_ll = g2_t / (2.0 * var_t)
else:
raise ValueError(
"Unrecognized importance sampling type: {}".format(
iw_sample_mode))
return t, var_t.view(-1, 1, 1, 1), m_t.view(-1, 1, 1, 1), obj_weight_t.view(-1, 1, 1, 1), \
obj_weight_t_ll.view(-1, 1, 1, 1), g2_t.view(-1, 1, 1, 1)
class DiffusionGeometric(DiffusionBase):
"""
Diffusion implementation with dz = -0.5 * beta(t) * z * dt + sqrt(beta(t)) * dW SDE and geometric progression of
variance. This is our new diffusion.
"""
def __init__(self, args):
super().__init__(args)
self.sigma2_min = args.sde_sigma2_min
self.sigma2_max = args.sde_sigma2_max
def f(self, t):
return -0.5 * self.g2(t)
def g2(self, t):
sigma2_geom = self.sigma2_min * (
(self.sigma2_max / self.sigma2_min)**t)
log_term = np.log(self.sigma2_max / self.sigma2_min)
return sigma2_geom * log_term / (1.0 - self.sigma2_0 +
self.sigma2_min - sigma2_geom)
def var(self, t):
return self.sigma2_min * ((self.sigma2_max / self.sigma2_min)**
t) - self.sigma2_min + self.sigma2_0
def e2int_f(self, t):
return torch.sqrt(1.0 + self.sigma2_min *
(1.0 - (self.sigma2_max / self.sigma2_min)**t) /
(1.0 - self.sigma2_0))
def inv_var(self, var):
return torch.log(
(var + self.sigma2_min - self.sigma2_0) /
self.sigma2_min) / np.log(self.sigma2_max / self.sigma2_min)
def mixing_component(self, x_noisy, var_t, t, enabled):
if enabled:
return torch.sqrt(var_t) * x_noisy
else:
return None
class DiffusionVPSDE(DiffusionBase):
"""
Diffusion implementation of the VPSDE. This uses the same SDE like DiffusionGeometric but with linear beta(t).
Note that we need to scale beta_start and beta_end by 1000 relative to JH's DDPM values, since our t is in [0,1].
"""
def __init__(self, args):
super().__init__(args)
# self.beta_start = args.sde_beta_start # 0.1
# self.beta_end = args.sde_beta_end # 20
# ! hard coded, in the scale of 1000.
# beta_start = scale * 0.0001
# beta_end = scale * 0.02
self.beta_start = 0.1
self.beta_end = 20
# auxiliary constants
self.time_eps = args.sde_time_eps # 0.01 by default in LSGM. Any influence?
self.delta_beta_half = torch.tensor(0.5 *
(self.beta_end - self.beta_start),
device=dist_util.dev())
self.beta_frac = torch.tensor(self.beta_start /
(self.beta_end - self.beta_start),
device=dist_util.dev())
self.const_aq = (1.0 - self.sigma2_0) * torch.exp(
0.5 * self.beta_frac) * torch.sqrt(
0.25 * np.pi / self.delta_beta_half)
self.const_erf = torch.erf(
torch.sqrt(self.delta_beta_half) *
(self.time_eps + self.beta_frac))
self.const_norm = self.const_aq * (torch.erf(
torch.sqrt(self.delta_beta_half) *
(1.0 + self.beta_frac)) - self.const_erf)
self.const_norm_2 = torch.erf(
torch.sqrt(self.delta_beta_half) *
(1.0 + self.beta_frac)) - self.const_erf
def f(self, t):
return -0.5 * self.g2(t)
def g2(self, t):
return self.beta_start + (self.beta_end - self.beta_start) * t
def var(self, t):
return 1.0 - (1.0 - self.sigma2_0
) * torch.exp(-self.beta_start * t - 0.5 *
(self.beta_end - self.beta_start) * t * t)
def e2int_f(self, t):
return torch.exp(-0.5 * self.beta_start * t - 0.25 *
(self.beta_end - self.beta_start) * t * t)
def inv_var(self, var):
c = torch.log((1 - var) / (1 - self.sigma2_0))
a = self.beta_end - self.beta_start
t = (-self.beta_start +
torch.sqrt(np.square(self.beta_start) - 2 * a * c)) / a
return t
def mixing_component(self, x_noisy, var_t, t, enabled):
if enabled:
return torch.sqrt(var_t) * x_noisy
else:
return None
def mixing_component_x0(self, x_noisy, var_t, t, enabled):
if enabled:
# return torch.sqrt(var_t) * x_noisy
return torch.sqrt(1-var_t) * x_noisy # zt * alpha_t
else:
return None
class DiffusionSubVPSDE(DiffusionBase):
"""
Diffusion implementation of the sub-VPSDE. Note that this uses a different SDE compared to the above two diffusions.
"""
def __init__(self, args):
super().__init__(args)
self.beta_start = args.sde_beta_start
self.beta_end = args.sde_beta_end
# auxiliary constants (assumes regular VPSDE)
self.time_eps = args.sde_time_eps
self.delta_beta_half = torch.tensor(0.5 *
(self.beta_end - self.beta_start),
device=dist_util.dev())
self.beta_frac = torch.tensor(self.beta_start /
(self.beta_end - self.beta_start),
device=dist_util.dev())
self.const_aq = (1.0 - self.sigma2_0) * torch.exp(
0.5 * self.beta_frac) * torch.sqrt(
0.25 * np.pi / self.delta_beta_half)
self.const_erf = torch.erf(
torch.sqrt(self.delta_beta_half) *
(self.time_eps + self.beta_frac))
self.const_norm = self.const_aq * (torch.erf(
torch.sqrt(self.delta_beta_half) *
(1.0 + self.beta_frac)) - self.const_erf)
self.const_norm_2 = torch.erf(
torch.sqrt(self.delta_beta_half) *
(1.0 + self.beta_frac)) - self.const_erf
def f(self, t):
return -0.5 * self.beta(t)
def g2(self, t):
return self.beta(t) * (
1.0 - torch.exp(-2.0 * self.beta_start * t -
(self.beta_end - self.beta_start) * t * t))
def var(self, t):
int_term = torch.exp(-self.beta_start * t - 0.5 *
(self.beta_end - self.beta_start) * t * t)
return torch.square(1.0 - int_term) + self.sigma2_0 * int_term
def e2int_f(self, t):
return torch.exp(-0.5 * self.beta_start * t - 0.25 *
(self.beta_end - self.beta_start) * t * t)
def beta(self, t):
""" auxiliary beta function """
return self.beta_start + (self.beta_end - self.beta_start) * t
def inv_var(self, var):
raise NotImplementedError
def mixing_component(self, x_noisy, var_t, t, enabled):
if enabled:
int_term = torch.exp(-self.beta_start * t - 0.5 *
(self.beta_end - self.beta_start) * t *
t).view(-1, 1, 1, 1)
return torch.sqrt(var_t) * x_noisy / (
torch.square(1.0 - int_term) + int_term)
else:
return None
def var_vpsde(self, t):
return 1.0 - (1.0 - self.sigma2_0
) * torch.exp(-self.beta_start * t - 0.5 *
(self.beta_end - self.beta_start) * t * t)
def inv_var_vpsde(self, var):
c = torch.log((1 - var) / (1 - self.sigma2_0))
a = self.beta_end - self.beta_start
t = (-self.beta_start +
torch.sqrt(np.square(self.beta_start) - 2 * a * c)) / a
return t
class DiffusionVESDE(DiffusionBase):
"""
Diffusion implementation of the VESDE with dz = sqrt(beta(t)) * dW
"""
def __init__(self, args):
super().__init__(args)
self.sigma2_min = args.sde_sigma2_min
self.sigma2_max = args.sde_sigma2_max
assert self.sigma2_min == self.sigma2_0, "VESDE was proposed implicitly assuming sigma2_min = sigma2_0!"
def f(self, t):
return torch.zeros_like(t, device=dist_util.dev())
def g2(self, t):
return self.sigma2_min * np.log(self.sigma2_max / self.sigma2_min) * (
(self.sigma2_max / self.sigma2_min)**t)
def var(self, t):
return self.sigma2_min * ((self.sigma2_max / self.sigma2_min)**
t) - self.sigma2_min + self.sigma2_0
def e2int_f(self, t):
return torch.ones_like(t, device=dist_util.dev())
def inv_var(self, var):
return torch.log(
(var + self.sigma2_min - self.sigma2_0) /
self.sigma2_min) / np.log(self.sigma2_max / self.sigma2_min)
def mixing_component(self, x_noisy, var_t, t, enabled):
if enabled:
return torch.sqrt(var_t) * x_noisy / (self.sigma2_min * (
(self.sigma2_max / self.sigma2_min)**t.view(-1, 1, 1, 1)) -
self.sigma2_min + 1.0)
else:
return None
def var_N(self, t):
return 1.0 - self.sigma2_min + self.sigma2_min * (
(self.sigma2_max / self.sigma2_min)**t)
def inv_var_N(self, var):
return torch.log(
(var + self.sigma2_min - 1.0) / self.sigma2_min) / np.log(
self.sigma2_max / self.sigma2_min)