simple_diffusion.py

importmath
fromfunctoolsimportpartial, wraps

importtorch
fromtorchimportsqrt
fromtorchimportnn, einsum
importtorch.nn.functionalasF
fromtorch.specialimportexpm1
fromtorch.ampimportautocast

fromtqdmimporttqdm
fromeinopsimportrearrange, repeat, reduce, pack, unpack
fromeinops.layers.torchimportRearrange

# helpers

defexists(val):
returnvalisnotNone

defidentity(t):
returnt

defis_lambda(f):
returncallable(f) andf.__name__=="<lambda>"

defdefault(val, d):
ifexists(val):
returnval
returnd() ifis_lambda(d) elsed

defcast_tuple(t, l=1):
return ((t,) *l) ifnotisinstance(t, tuple) elset

defappend_dims(t, dims):
shape=t.shape
returnt.reshape(*shape, *((1,) *dims))

defl2norm(t):
returnF.normalize(t, dim=-1)

# u-vit related functions and modules

classUpsample(nn.Module):
def__init__(
self,
dim,
dim_out=None,
factor=2
 ):
super().__init__()
self.factor=factor
self.factor_squared=factor**2

dim_out=default(dim_out, dim)
conv=nn.Conv2d(dim, dim_out*self.factor_squared, 1)

self.net=nn.Sequential(
conv,
nn.SiLU(),
nn.PixelShuffle(factor)
 )

self.init_conv_(conv)

definit_conv_(self, conv):
o, i, h, w=conv.weight.shape
conv_weight=torch.empty(o//self.factor_squared, i, h, w)
nn.init.kaiming_uniform_(conv_weight)
conv_weight=repeat(conv_weight, 'o ... -> (o r) ...', r=self.factor_squared)

conv.weight.data.copy_(conv_weight)
nn.init.zeros_(conv.bias.data)

defforward(self, x):
returnself.net(x)

defDownsample(
dim,
dim_out=None,
factor=2
):
returnnn.Sequential(
Rearrange('b c (h p1) (w p2) -> b (c p1 p2) h w', p1=factor, p2=factor),
nn.Conv2d(dim* (factor**2), default(dim_out, dim), 1)
 )

classRMSNorm(nn.Module):
def__init__(self, dim, scale=True, normalize_dim=2):
super().__init__()
self.g=nn.Parameter(torch.ones(dim)) ifscaleelse1

self.scale=scale
self.normalize_dim=normalize_dim

defforward(self, x):
normalize_dim=self.normalize_dim
scale=append_dims(self.g, x.ndim-self.normalize_dim-1) ifself.scaleelse1
returnF.normalize(x, dim=normalize_dim) *scale* (x.shape[normalize_dim] **0.5)

# sinusoidal positional embeds

classLearnedSinusoidalPosEmb(nn.Module):
def__init__(self, dim):
super().__init__()
assert (dim%2) ==0
half_dim=dim//2
self.weights=nn.Parameter(torch.randn(half_dim))

defforward(self, x):
x=rearrange(x, 'b -> b 1')
freqs=x*rearrange(self.weights, 'd -> 1 d') *2*math.pi
fouriered=torch.cat((freqs.sin(), freqs.cos()), dim=-1)
fouriered=torch.cat((x, fouriered), dim=-1)
returnfouriered

# building block modules

classBlock(nn.Module):
def__init__(self, dim, dim_out):
super().__init__()
self.proj=nn.Conv2d(dim, dim_out, 3, padding=1)
self.norm=RMSNorm(dim_out, normalize_dim=1)
self.act=nn.SiLU()

defforward(self, x, scale_shift=None):
x=self.proj(x)
x=self.norm(x)

ifexists(scale_shift):
scale, shift=scale_shift
x=x* (scale+1) +shift

x=self.act(x)
returnx

classResnetBlock(nn.Module):
def__init__(self, dim, dim_out, *, time_emb_dim=None):
super().__init__()
self.mlp=nn.Sequential(
nn.SiLU(),
nn.Linear(time_emb_dim, dim_out*2)
 ) ifexists(time_emb_dim) elseNone

self.block1=Block(dim, dim_out)
self.block2=Block(dim_out, dim_out)
self.res_conv=nn.Conv2d(dim, dim_out, 1) ifdim!=dim_outelsenn.Identity()

defforward(self, x, time_emb=None):

scale_shift=None
ifexists(self.mlp) andexists(time_emb):
time_emb=self.mlp(time_emb)
time_emb=rearrange(time_emb, 'b c -> b c 1 1')
scale_shift=time_emb.chunk(2, dim=1)

h=self.block1(x, scale_shift=scale_shift)

h=self.block2(h)

returnh+self.res_conv(x)

classLinearAttention(nn.Module):
def__init__(self, dim, heads=4, dim_head=32):
super().__init__()
self.scale=dim_head**-0.5
self.heads=heads
hidden_dim=dim_head*heads

self.norm=RMSNorm(dim, normalize_dim=1)
self.to_qkv=nn.Conv2d(dim, hidden_dim*3, 1, bias=False)

self.to_out=nn.Sequential(
nn.Conv2d(hidden_dim, dim, 1),
RMSNorm(dim, normalize_dim=1)
 )

defforward(self, x):
residual=x

b, c, h, w=x.shape

x=self.norm(x)

qkv=self.to_qkv(x).chunk(3, dim=1)
q, k, v=map(lambdat: rearrange(t, 'b (h c) x y -> b h c (x y)', h=self.heads), qkv)

q=q.softmax(dim=-2)
k=k.softmax(dim=-1)

q=q*self.scale

context=torch.einsum('b h d n, b h e n -> b h d e', k, v)

out=torch.einsum('b h d e, b h d n -> b h e n', context, q)
out=rearrange(out, 'b h c (x y) -> b (h c) x y', h=self.heads, x=h, y=w)

returnself.to_out(out) +residual

classAttention(nn.Module):
def__init__(self, dim, heads=4, dim_head=32, scale=8, dropout=0.):
super().__init__()
self.scale=scale
self.heads=heads
hidden_dim=dim_head*heads

self.norm=RMSNorm(dim)

self.attn_dropout=nn.Dropout(dropout)
self.to_qkv=nn.Linear(dim, hidden_dim*3, bias=False)

self.q_scale=nn.Parameter(torch.ones(dim_head))
self.k_scale=nn.Parameter(torch.ones(dim_head))

self.to_out=nn.Linear(hidden_dim, dim, bias=False)

defforward(self, x):
x=self.norm(x)

qkv=self.to_qkv(x).chunk(3, dim=-1)
q, k, v=map(lambdat: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)

q, k=map(l2norm, (q, k))

q=q*self.q_scale
k=k*self.k_scale

sim=einsum('b h i d, b h j d -> b h i j', q, k) *self.scale

attn=sim.softmax(dim=-1)
attn=self.attn_dropout(attn)

out=einsum('b h i j, b h j d -> b h i d', attn, v)

out=rearrange(out, 'b h n d -> b n (h d)')
returnself.to_out(out)

classFeedForward(nn.Module):
def__init__(
self,
dim,
cond_dim,
mult=4,
dropout=0.
 ):
super().__init__()
self.norm=RMSNorm(dim, scale=False)
dim_hidden=dim*mult

self.to_scale_shift=nn.Sequential(
nn.SiLU(),
nn.Linear(cond_dim, dim_hidden*2),
Rearrange('b d -> b 1 d')
 )

to_scale_shift_linear=self.to_scale_shift[-2]
nn.init.zeros_(to_scale_shift_linear.weight)
nn.init.zeros_(to_scale_shift_linear.bias)

self.proj_in=nn.Sequential(
nn.Linear(dim, dim_hidden, bias=False),
nn.SiLU()
 )

self.proj_out=nn.Sequential(
nn.Dropout(dropout),
nn.Linear(dim_hidden, dim, bias=False)
 )

defforward(self, x, t):
x=self.norm(x)
x=self.proj_in(x)

scale, shift=self.to_scale_shift(t).chunk(2, dim=-1)
x=x* (scale+1) +shift

returnself.proj_out(x)

# vit

classTransformer(nn.Module):
def__init__(
self,
dim,
time_cond_dim,
depth,
dim_head=32,
heads=4,
ff_mult=4,
dropout=0.,
 ):
super().__init__()

self.layers=nn.ModuleList([])
for_inrange(depth):
self.layers.append(nn.ModuleList([
Attention(dim=dim, dim_head=dim_head, heads=heads, dropout=dropout),
FeedForward(dim=dim, mult=ff_mult, cond_dim=time_cond_dim, dropout=dropout)
 ]))

defforward(self, x, t):
forattn, ffinself.layers:
x=attn(x) +x
x=ff(x, t) +x

returnx
# model

classUViT(nn.Module):
def__init__(
self,
dim,
init_dim=None,
out_dim=None,
dim_mults= (1, 2, 4, 8),
downsample_factor=2,
channels=3,
vit_depth=6,
vit_dropout=0.2,
attn_dim_head=32,
attn_heads=4,
ff_mult=4,
learned_sinusoidal_dim=16,
init_img_transform: callable=None,
final_img_itransform: callable=None,
patch_size=1,
dual_patchnorm=False
 ):
super().__init__()

# for initial dwt transform (or whatever transform researcher wants to try here)

ifexists(init_img_transform) andexists(final_img_itransform):
init_shape=torch.Size(1, 1, 32, 32)
mock_tensor=torch.randn(init_shape)
assertfinal_img_itransform(init_img_transform(mock_tensor)).shape==init_shape

self.init_img_transform=default(init_img_transform, identity)
self.final_img_itransform=default(final_img_itransform, identity)

input_channels=channels

init_dim=default(init_dim, dim)
self.init_conv=nn.Conv2d(input_channels, init_dim, 7, padding=3)

# whether to do initial patching, as alternative to dwt

self.unpatchify=identity

input_channels=channels* (patch_size**2)
needs_patch=patch_size>1

ifneeds_patch:
ifnotdual_patchnorm:
self.init_conv=nn.Conv2d(channels, init_dim, patch_size, stride=patch_size)
else:
self.init_conv=nn.Sequential(
Rearrange('b c (h p1) (w p2) -> b h w (c p1 p2)', p1=patch_size, p2=patch_size),
nn.LayerNorm(input_channels),
nn.Linear(input_channels, init_dim),
nn.LayerNorm(init_dim),
Rearrange('b h w c -> b c h w')
 )

self.unpatchify=nn.ConvTranspose2d(input_channels, channels, patch_size, stride=patch_size)

# determine dimensions

dims= [init_dim, *map(lambdam: dim*m, dim_mults)]
in_out=list(zip(dims[:-1], dims[1:]))

# time embeddings

time_dim=dim*4

sinu_pos_emb=LearnedSinusoidalPosEmb(learned_sinusoidal_dim)
fourier_dim=learned_sinusoidal_dim+1

self.time_mlp=nn.Sequential(
sinu_pos_emb,
nn.Linear(fourier_dim, time_dim),
nn.GELU(),
nn.Linear(time_dim, time_dim)
 )

# downsample factors

downsample_factor=cast_tuple(downsample_factor, len(dim_mults))
assertlen(downsample_factor) ==len(dim_mults)

# layers

self.downs=nn.ModuleList([])
self.ups=nn.ModuleList([])
num_resolutions=len(in_out)

forind, ((dim_in, dim_out), factor) inenumerate(zip(in_out, downsample_factor)):
is_last=ind>= (num_resolutions-1)

self.downs.append(nn.ModuleList([
ResnetBlock(dim_in, dim_in, time_emb_dim=time_dim),
ResnetBlock(dim_in, dim_in, time_emb_dim=time_dim),
LinearAttention(dim_in),
Downsample(dim_in, dim_out, factor=factor)
 ]))

mid_dim=dims[-1]

self.vit=Transformer(
dim=mid_dim,
time_cond_dim=time_dim,
depth=vit_depth,
dim_head=attn_dim_head,
heads=attn_heads,
ff_mult=ff_mult,
dropout=vit_dropout
 )

forind, ((dim_in, dim_out), factor) inenumerate(zip(reversed(in_out), reversed(downsample_factor))):
is_last=ind== (len(in_out) -1)

self.ups.append(nn.ModuleList([
Upsample(dim_out, dim_in, factor=factor),
ResnetBlock(dim_in*2, dim_in, time_emb_dim=time_dim),
ResnetBlock(dim_in*2, dim_in, time_emb_dim=time_dim),
LinearAttention(dim_in),
 ]))

default_out_dim=input_channels
self.out_dim=default(out_dim, default_out_dim)

self.final_res_block=ResnetBlock(init_dim*2, init_dim, time_emb_dim=time_dim)
self.final_conv=nn.Conv2d(init_dim, self.out_dim, 1)

defforward(self, x, time):
x=self.init_img_transform(x)

x=self.init_conv(x)
r=x.clone()

t=self.time_mlp(time)

h= []

forblock1, block2, attn, downsampleinself.downs:
x=block1(x, t)
h.append(x)

x=block2(x, t)
x=attn(x)
h.append(x)

x=downsample(x)

x=rearrange(x, 'b c h w -> b h w c')
x, ps=pack([x], 'b * c')

x=self.vit(x, t)

x, =unpack(x, ps, 'b * c')
x=rearrange(x, 'b h w c -> b c h w')

forupsample, block1, block2, attninself.ups:
x=upsample(x)

x=torch.cat((x, h.pop()), dim=1)
x=block1(x, t)

x=torch.cat((x, h.pop()), dim=1)
x=block2(x, t)
x=attn(x)

x=torch.cat((x, r), dim=1)

x=self.final_res_block(x, t)
x=self.final_conv(x)

x=self.unpatchify(x)
returnself.final_img_itransform(x)

# normalization functions

defnormalize_to_neg_one_to_one(img):
returnimg*2-1

defunnormalize_to_zero_to_one(t):
return (t+1) *0.5

# diffusion helpers

defright_pad_dims_to(x, t):
padding_dims=x.ndim-t.ndim
ifpadding_dims<=0:
returnt
returnt.view(*t.shape, *((1,) *padding_dims))

# logsnr schedules and shifting / interpolating decorators
# only cosine for now

deflog(t, eps=1e-20):
returntorch.log(t.clamp(min=eps))

deflogsnr_schedule_cosine(t, logsnr_min=-15, logsnr_max=15):
t_min=math.atan(math.exp(-0.5*logsnr_max))
t_max=math.atan(math.exp(-0.5*logsnr_min))
return-2*log(torch.tan(t_min+t* (t_max-t_min)))

deflogsnr_schedule_shifted(fn, image_d, noise_d):
shift=2*math.log(noise_d/image_d)
@wraps(fn)
definner(*args, **kwargs):
nonlocalshift
returnfn(*args, **kwargs) +shift
returninner

deflogsnr_schedule_interpolated(fn, image_d, noise_d_low, noise_d_high):
logsnr_low_fn=logsnr_schedule_shifted(fn, image_d, noise_d_low)
logsnr_high_fn=logsnr_schedule_shifted(fn, image_d, noise_d_high)

@wraps(fn)
definner(t, *args, **kwargs):
nonlocallogsnr_low_fn
nonlocallogsnr_high_fn
returnt*logsnr_low_fn(t, *args, **kwargs) + (1-t) *logsnr_high_fn(t, *args, **kwargs)

returninner

# main gaussian diffusion class

classGaussianDiffusion(nn.Module):
def__init__(
self,
model: UViT,
*,
image_size,
channels=3,
pred_objective='v',
noise_schedule=logsnr_schedule_cosine,
noise_d=None,
noise_d_low=None,
noise_d_high=None,
num_sample_steps=500,
clip_sample_denoised=True,
min_snr_loss_weight=True,
min_snr_gamma=5
 ):
super().__init__()
assertpred_objectivein {'v', 'eps'}, 'whether to predict v-space (progressive distillation paper) or noise'

self.model=model

# image dimensions

self.channels=channels
self.image_size=image_size

# training objective

self.pred_objective=pred_objective

# noise schedule

assertnotall([*map(exists, (noise_d, noise_d_low, noise_d_high))]), 'you must either set noise_d for shifted schedule, or noise_d_low and noise_d_high for shifted and interpolated schedule'

# determine shifting or interpolated schedules

self.log_snr=noise_schedule

ifexists(noise_d):
self.log_snr=logsnr_schedule_shifted(self.log_snr, image_size, noise_d)

ifexists(noise_d_low) orexists(noise_d_high):
assertexists(noise_d_low) andexists(noise_d_high), 'both noise_d_low and noise_d_high must be set'

self.log_snr=logsnr_schedule_interpolated(self.log_snr, image_size, noise_d_low, noise_d_high)

# sampling

self.num_sample_steps=num_sample_steps
self.clip_sample_denoised=clip_sample_denoised

# loss weight

self.min_snr_loss_weight=min_snr_loss_weight
self.min_snr_gamma=min_snr_gamma

@property
defdevice(self):
returnnext(self.model.parameters()).device

defp_mean_variance(self, x, time, time_next):

log_snr=self.log_snr(time)
log_snr_next=self.log_snr(time_next)
c=-expm1(log_snr-log_snr_next)

squared_alpha, squared_alpha_next=log_snr.sigmoid(), log_snr_next.sigmoid()
squared_sigma, squared_sigma_next= (-log_snr).sigmoid(), (-log_snr_next).sigmoid()

alpha, sigma, alpha_next=map(sqrt, (squared_alpha, squared_sigma, squared_alpha_next))

batch_log_snr=repeat(log_snr, ' -> b', b=x.shape[0])
pred=self.model(x, batch_log_snr)

ifself.pred_objective=='v':
x_start=alpha*x-sigma*pred

elifself.pred_objective=='eps':
x_start= (x-sigma*pred) /alpha

x_start.clamp_(-1., 1.)

model_mean=alpha_next* (x* (1-c) /alpha+c*x_start)

posterior_variance=squared_sigma_next*c

returnmodel_mean, posterior_variance

# sampling related functions

@torch.no_grad()
defp_sample(self, x, time, time_next):
batch, *_, device=*x.shape, x.device

model_mean, model_variance=self.p_mean_variance(x=x, time=time, time_next=time_next)

iftime_next==0:
returnmodel_mean

noise=torch.randn_like(x)
returnmodel_mean+sqrt(model_variance) *noise

@torch.no_grad()
defp_sample_loop(self, shape):
batch=shape[0]

img=torch.randn(shape, device=self.device)
steps=torch.linspace(1., 0., self.num_sample_steps+1, device=self.device)

foriintqdm(range(self.num_sample_steps), desc='sampling loop time step', total=self.num_sample_steps):
times=steps[i]
times_next=steps[i+1]
img=self.p_sample(img, times, times_next)

img.clamp_(-1., 1.)
img=unnormalize_to_zero_to_one(img)
returnimg

@torch.no_grad()
defsample(self, batch_size=16):
returnself.p_sample_loop((batch_size, self.channels, self.image_size, self.image_size))

# training related functions - noise prediction

@autocast('cuda', enabled=False)
defq_sample(self, x_start, times, noise=None):
noise=default(noise, lambda: torch.randn_like(x_start))

log_snr=self.log_snr(times)

log_snr_padded=right_pad_dims_to(x_start, log_snr)
alpha, sigma=sqrt(log_snr_padded.sigmoid()), sqrt((-log_snr_padded).sigmoid())
x_noised=x_start*alpha+noise*sigma

returnx_noised, log_snr

defp_losses(self, x_start, times, noise=None):
noise=default(noise, lambda: torch.randn_like(x_start))

x, log_snr=self.q_sample(x_start=x_start, times=times, noise=noise)
model_out=self.model(x, log_snr)

ifself.pred_objective=='v':
padded_log_snr=right_pad_dims_to(x, log_snr)
alpha, sigma=padded_log_snr.sigmoid().sqrt(), (-padded_log_snr).sigmoid().sqrt()
target=alpha*noise-sigma*x_start

elifself.pred_objective=='eps':
target=noise

loss=F.mse_loss(model_out, target, reduction='none')

loss=reduce(loss, 'b ... -> b', 'mean')

snr=log_snr.exp()

maybe_clip_snr=snr.clone()
ifself.min_snr_loss_weight:
maybe_clip_snr.clamp_(max=self.min_snr_gamma)

ifself.pred_objective=='v':
loss_weight=maybe_clip_snr/ (snr+1)

elifself.pred_objective=='eps':
loss_weight=maybe_clip_snr/snr

return (loss*loss_weight).mean()

defforward(self, img, *args, **kwargs):
b, c, h, w, device, img_size, =*img.shape, img.device, self.image_size
asserth==img_sizeandw==img_size, f'height and width of image must be {img_size}'

img=normalize_to_neg_one_to_one(img)
times=torch.zeros((img.shape[0],), device=self.device).float().uniform_(0, 1)

returnself.p_losses(img, times, *args, **kwargs)