title denoising diffusion probabilistic model aknainhttps twitter comaknain date created 20221130 last modified 20221207 description generating images of flowers with denoising diffusion probabilistic models introduction generative modeling experienced tremendous growth in the last five years models like vaes gans and flowbased models proved to be a great success in generating highquality content especially images diffusion models are a new type of generative model that has proven to be better than previous approaches diffusion models are inspired by nonequilibrium thermodynamics and they learn to generate by denoising learning by denoising consists of two processes each of which is a markov chain these are 1 the forward process in the forward process we slowly add random noise to the data in a series of time steps t1 t2 tn samples at the current time step are drawn from a gaussian distribution where the mean of the distribution is conditioned on the sample at the previous time step and the variance of the distribution follows a fixed schedule at the end of the forward process the samples end up with a pure noise distribution 2 the reverse process during the reverse process we try to undo the added noise at every time step we start with the pure noise distribution the last step of the forward process and try to denoise the samples in the backward direction tn tn1 t1 we implement the denoising diffusion probabilistic modelshttps arxiv orgabs2006 11239 paper or ddpms for short in this code example it was the first paper demonstrating the use of diffusion models for generating highquality images the s proved that a certain parameterization of diffusion models reveals an equivalence with denoising score matching over multiple noise levels during training and with annealed langevin dynamics during sampling that generates the best quality results this paper replicates both the markov chains forward process and reverse process involved in the diffusion process but for images the forward process is fixed and gradually adds gaussian noise to the images according to a fixed variance schedule denoted by beta in the paper this is what the diffusion process looks like in case of images image noise noise image diffusion process gifhttps imgur comyn7tho9 gif the paper describes two algorithms one for training the model and the other for sampling from the trained model training is performed by optimizing the usual variational bound on negative loglikelihood the objective function is further simplified and the network is treated as a noise prediction network once optimized we can sample from the network to generate new images from noise samples here is an overview of both algorithms as presented in the paper ddpmshttps i imgur coms7kh5hz png note ddpm is just one way of implementing a diffusion model also the sampling algorithm in the ddpm replicates the complete markov chain hence it s slow in generating new samples compared to other generative models like gans lots of research efforts have been made to address this issue one such example is denoising diffusion implicit models or ddim for short where the s replaced the markov chain with a nonmarkovian process to sample faster you can find the code example for ddim herehttps keras ioexamplesgenerativeddim implementing a ddpm model is simple we define a model that takes two inputs images and the randomly sampled time steps at each training step we perform the following operations to train our model 1 sample random noise to be added to the inputs 2 apply the forward process to diffuse the inputs with the sampled noise 3 your model takes these noisy samples as inputs and outputs the noise prediction for each time step 4 given true noise and predicted noise we calculate the loss values 5 we then calculate the gradients and update the model weights given that our model knows how to denoise a noisy sample at a given time step we can leverage this idea to generate new samples starting from a pure noise distribution setup requires tensorflow 2 11 for the groupnormalization layer hyperparameters dataset we use the oxford flowers 102https www tensorflow orgdatasetscatalogoxfordflowers102 dataset for generating images of flowers in terms of preprocessing we use center cropping for resizing the images to the desired image size and we rescale the pixel values in the range 1 0 1 0 this is in line with the range of the pixel values that was applied by the s of the ddpms paperhttps arxiv orgabs2006 11239 for augmenting training data we randomly flip the images leftright load the dataset flips an image leftright randomly return tf image randomflipleftrightimg def resizeandrescaleimg size height tf shapeimg0 width tf shapeimg1 cropsize tf minimumheight width img tf image croptoboundingbox img height cropsize 2 width cropsize 2 cropsize cropsize resize img tf castimg dtypetf float32 img tf image resizeimg sizesize antialiastrue rescale the pixel values img img 127 5 1 0 img tf clipbyvalueimg clipmin clipmax return img def trainpreprocessingx img ximage img resizeandrescaleimg sizeimgsize imgsize img augmentimg return img trainds ds maptrainpreprocessing numparallelcallstf data autotune batchbatchsize dropremaindertrue shufflebatchsize 2 prefetchtf data autotune gaussian diffusion utilities class gaussiandiffusion def init self betastart1e4 betaend0 02 timesteps1000 clipmin1 0 clipmax1 0 self betastart betastart self betaend betaend self timesteps timesteps self clipmin clipmin self clipmax clipmax define the linear variance schedule self betas betas np linspace betastart betaend timesteps dtypenp float64 using float64 for better precision self numtimesteps inttimesteps alphas 1 0 betas alphascumprod np cumprodalphas axis0 alphascumprodprev np append1 0 alphascumprod 1 self betas tf constantbetas dtypetf float32 self alphascumprod tf constantalphascumprod dtypetf float32 self alphascumprodprev tf constant alphascumprodprev dtypetf float32 calculations for diffusion qxt xt1 and others self sqrtalphascumprod tf constant np sqrtalphascumprod dtypetf float32 self sqrtoneminusalphascumprod tf constant np sqrt1 0 alphascumprod dtypetf float32 self logoneminusalphascumprod tf constant np log1 0 alphascumprod dtypetf float32 self sqrtrecipalphascumprod tf constant np sqrt1 0 alphascumprod dtypetf float32 self sqrtrecipm1alphascumprod tf constant np sqrt1 0 alphascumprod 1 dtypetf float32 calculations for posterior qxt1 xt x0 posteriorvariance betas 1 0 alphascumprodprev 1 0 alphascumprod self posteriorvariance tf constant posteriorvariance dtypetf float32 log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain self posteriorlogvarianceclipped tf constant np lognp maximumposteriorvariance 1e20 dtypetf float32 self posteriormeancoef1 tf constant betas np sqrtalphascumprodprev 1 0 alphascumprod dtypetf float32 self posteriormeancoef2 tf constant 1 0 alphascumprodprev np sqrtalphas 1 0 alphascumprod dtypetf float32 def extractself a t xshape batchsize xshape0 out tf gathera t return tf reshapeout batchsize 1 1 1 def qmeanvarianceself xstart t xstartshape tf shapexstart mean self extractself sqrtalphascumprod t xstartshape xstart variance self extract1 0 self alphascumprod t xstartshape logvariance self extract self logoneminusalphascumprod t xstartshape return mean variance logvariance def qsampleself xstart t noise xstartshape tf shapexstart return self extractself sqrtalphascumprod t tf shapexstart xstart self extract self sqrtoneminusalphascumprod t xstartshape noise def predictstartfromnoiseself xt t noise xtshape tf shapext return self extractself sqrtrecipalphascumprod t xtshape xt self extractself sqrtrecipm1alphascumprod t xtshape noise def qposteriorself xstart xt t xtshape tf shapext posteriormean self extractself posteriormeancoef1 t xtshape xstart self extractself posteriormeancoef2 t xtshape xt posteriorvariance self extract self posteriorvariance t xtshape posteriorlogvarianceclipped self extract self posteriorlogvarianceclipped t xtshape return posteriormean posteriorvariance posteriorlogvarianceclipped def pmeanvarianceself prednoise x t clipdenoisedtrue xrecon self predictstartfromnoisex tt noiseprednoise if clipdenoised xrecon tf clipbyvaluexrecon self clipmin self clipmax modelmean posteriorvariance posteriorlogvariance self qposteriorxstartxrecon xtx tt return modelmean posteriorvariance posteriorlogvariance def psampleself prednoise x t clipdenoisedtrue modelmean modellogvariance self pmeanvariance prednoise xx tt clipdenoisedclipdenoised noise tf random normalshapex shape dtypex dtype no noise when t 0 nonzeromask tf reshape 1 tf casttf equalt 0 tf float32 tf shapex0 1 1 1 return modelmean nonzeromask tf exp0 5 modellogvariance noise network architecture kernel initializer to use def kernelinitscale scale maxscale 1e10 return keras initializers variancescaling scale modefanavg distributionuniform class attentionblocklayers layer def initself units groups8 kwargs self units units self groups groups super initkwargs self norm layers groupnormalizationgroupsgroups self query layers denseunits kernelinitializerkernelinit1 0 self key layers denseunits kernelinitializerkernelinit1 0 self value layers denseunits kernelinitializerkernelinit1 0 self proj layers denseunits kernelinitializerkernelinit0 0 def callself inputs batchsize tf shapeinputs0 height tf shapeinputs1 width tf shapeinputs2 scale tf castself units tf float32 0 5 inputs self norminputs q self queryinputs k self keyinputs v self valueinputs attnscore tf einsumbhwc bhwcbhwhw q k scale attnscore tf reshape attnscore batchsize height width height width attnscore tf nn softmaxattnscore 1 attnscore tf reshape attnscore batchsize height width height width proj tf einsumbhwhw bhwcbhwc attnscore v proj self projproj return inputs proj class timeembeddinglayers layer def initself dim kwargs super initkwargs self dim dim self halfdim dim 2 self emb math log10000 self halfdim 1 self emb tf exptf rangeself halfdim dtypetf float32 self emb def callself inputs inputs tf castinputs dtypetf float32 emb inputs none self embnone emb tf concattf sinemb tf cosemb axis1 return emb def residualblockwidth groups8 activationfnkeras activations swish def applyinputs x t inputs inputwidth x shape3 if inputwidth width residual x else residual layers conv2d width kernelsize1 kernelinitializerkernelinit1 0 x temb activationfnt temb layers densewidth kernelinitializerkernelinit1 0temb none none x layers groupnormalizationgroupsgroupsx x activationfnx x layers conv2d width kernelsize3 paddingsame kernelinitializerkernelinit1 0 x x layers addx temb x layers groupnormalizationgroupsgroupsx x activationfnx x layers conv2d width kernelsize3 paddingsame kernelinitializerkernelinit0 0 x x layers addx residual return x return apply def downsamplewidth def applyx x layers conv2d width kernelsize3 strides2 paddingsame kernelinitializerkernelinit1 0 x return x return apply def upsamplewidth interpolationnearest def applyx x layers upsampling2dsize2 interpolationinterpolationx x layers conv2d width kernelsize3 paddingsame kernelinitializerkernelinit1 0 x return x return apply def timemlpunits activationfnkeras activations swish def applyinputs temb layers dense units activationactivationfn kernelinitializerkernelinit1 0 inputs temb layers denseunits kernelinitializerkernelinit1 0temb return temb return apply def buildmodel imgsize imgchannels widths hasattention numresblocks2 normgroups8 interpolationnearest activationfnkeras activations swish imageinput layers input shapeimgsize imgsize imgchannels nameimageinput timeinput keras inputshape dtypetf int64 nametimeinput x layers conv2d firstconvchannels kernelsize3 3 paddingsame kernelinitializerkernelinit1 0 imageinput temb timeembeddingdimfirstconvchannels 4timeinput temb timemlpunitsfirstconvchannels 4 activationfnactivationfn temb skips x downblock for i in rangelenwidths for in rangenumresblocks x residualblock widthsi groupsnormgroups activationfnactivationfn x temb if hasattentioni x attentionblockwidthsi groupsnormgroupsx skips appendx if widthsi widths1 x downsamplewidthsix skips appendx middleblock x residualblock widths1 groupsnormgroups activationfnactivationfn x temb x attentionblockwidths1 groupsnormgroupsx x residualblock widths1 groupsnormgroups activationfnactivationfn x temb upblock for i in reversedrangelenwidths for in rangenumresblocks 1 x layers concatenateaxis1x skips pop x residualblock widthsi groupsnormgroups activationfnactivationfn x temb if hasattentioni x attentionblockwidthsi groupsnormgroupsx if i 0 x upsamplewidthsi interpolationinterpolationx end block x layers groupnormalizationgroupsnormgroupsx x activationfnx x layers conv2d 3 3 3 paddingsame kernelinitializerkernelinit0 0 x return keras modelimageinput timeinput x nameunet training class diffusionmodelkeras model def initself network emanetwork timesteps gdfutil ema0 999 super init self network network self emanetwork emanetwork self timesteps timesteps self gdfutil gdfutil self ema ema def trainstepself images 1 get the batch size batchsize tf shapeimages0 2 sample timesteps uniformly t tf random uniform minval0 maxvalself timesteps shapebatchsize dtypetf int64 with tf gradienttape as tape 3 sample random noise to be added to the images in the batch noise tf random normalshapetf shapeimages dtypeimages dtype 4 diffuse the images with noise imagest self gdfutil qsampleimages t noise 5 pass the diffused images and time steps to the network prednoise self networkimagest t trainingtrue 6 calculate the loss loss self lossnoise prednoise 7 get the gradients gradients tape gradientloss self network trainableweights 8 update the weights of the network self optimizer applygradients zipgradients self network trainableweights 9 updates the weight values for the network with ema weights for weight emaweight in zip self network weights self emanetwork weights emaweight assignself ema emaweight 1 self ema weight 10 return loss values return loss loss def generateimagesself numimages16 1 randomly sample noise starting point for reverse process samples tf random normal shapenumimages imgsize imgsize imgchannels dtypetf float32 2 sample from the model iteratively for t in reversedrange0 self timesteps tt tf casttf fillnumimages t dtypetf int64 prednoise self emanetwork predict samples tt verbose0 batchsizenumimages samples self gdfutil psample prednoise samples tt clipdenoisedtrue 3 return generated samples return samples def plotimages self epochnone logsnone numrows2 numcols8 figsize12 5 build the unet model get an instance of the gaussian diffusion utilities get the model compile the model train the model results we trained this model for 800 epochs on a v100 gpu and each epoch took almost 8 seconds to finish we load those weights here and we generate a few samples starting from pure noise shell curl lo https github comaakashkumarnainddpmsreleasesdownloadv3 0 0checkpoints zip unzip qq checkpoints zip load the model weights generate and plot some samples conclusion we successfully implemented and trained a diffusion model exactly in the same fashion as implemented by the s of the ddpms paper you can find the original implementation herehttps github comhojonathanhodiffusion there are a few things that you can try to improve the model 1 increasing the width of each block a bigger model can learn to denoise in fewer epochs though you may have to take care of overfitting 2 we implemented the linear schedule for variance scheduling you can implement other schemes like cosine scheduling and compare the performance references 1 denoising diffusion probabilistic modelshttps arxiv orgabs2006 11239 2 s implementationhttps github comhojonathanhodiffusion 3 a deep dive into ddpmshttps magicwithlatents github iolatentpostsddpmspart3 4 denoising diffusion implicit modelshttps keras ioexamplesgenerativeddim 5 annotated diffusion modelhttps huggingface coblogannotateddiffusion 6 aiaiarthttps www youtube comwatch vxts7m6tsk9it14s introduction generative modeling experienced tremendous growth in the last five years models like vaes gans and flow based models proved to be a great success in generating high quality content especially images diffusion models are a new type of generative model that has proven to be better than previous approaches diffusion models are inspired by non equilibrium thermodynamics and they learn to generate by denoising learning by denoising consists of two processes each of which is a markov chain these are 1 the forward process in the forward process we slowly add random noise to the data in a series of time steps t1 t2 tn samples at the current time step are drawn from a gaussian distribution where the mean of the distribution is conditioned on the sample at the previous time step and the variance of the distribution follows a fixed schedule at the end of the forward process the samples end up with a pure noise distribution 2 the reverse process during the reverse process we try to undo the added noise at every time step we start with the pure noise distribution the last step of the forward process and try to denoise the samples in the backward direction tn tn 1 t1 we implement the denoising diffusion probabilistic models https arxiv org abs 2006 11239 paper or ddpms for short in this code example it was the first paper demonstrating the use of diffusion models for generating high quality images the s proved that a certain parameterization of diffusion models reveals an equivalence with denoising score matching over multiple noise levels during training and with annealed langevin dynamics during sampling that generates the best quality results this paper replicates both the markov chains forward process and reverse process involved in the diffusion process but for images the forward process is fixed and gradually adds gaussian noise to the images according to a fixed variance schedule denoted by beta in the paper this is what the diffusion process looks like in case of images image noise noise image diffusion process gif https imgur com yn7tho9 gif the paper describes two algorithms one for training the model and the other for sampling from the trained model training is performed by optimizing the usual variational bound on negative log likelihood the objective function is further simplified and the network is treated as a noise prediction network once optimized we can sample from the network to generate new images from noise samples here is an overview of both algorithms as presented in the paper ddpms https i imgur com s7kh5hz png note ddpm is just one way of implementing a diffusion model also the sampling algorithm in the ddpm replicates the complete markov chain hence it s slow in generating new samples compared to other generative models like gans lots of research efforts have been made to address this issue one such example is denoising diffusion implicit models or ddim for short where the s replaced the markov chain with a non markovian process to sample faster you can find the code example for ddim here https keras io examples generative ddim implementing a ddpm model is simple we define a model that takes two inputs images and the randomly sampled time steps at each training step we perform the following operations to train our model 1 sample random noise to be added to the inputs 2 apply the forward process to diffuse the inputs with the sampled noise 3 your model takes these noisy samples as inputs and outputs the noise prediction for each time step 4 given true noise and predicted noise we calculate the loss values 5 we then calculate the gradients and update the model weights given that our model knows how to denoise a noisy sample at a given time step we can leverage this idea to generate new samples starting from a pure noise distribution setup requires tensorflow 2 11 for the groupnormalization layer hyperparameters just for the sake of demonstration number of groups used in groupnormalization layer number of residual blocks dataset we use the oxford flowers 102 https www tensorflow org datasets catalog oxford_flowers102 dataset for generating images of flowers in terms of preprocessing we use center cropping for resizing the images to the desired image size and we rescale the pixel values in the range 1 0 1 0 this is in line with the range of the pixel values that was applied by the s of the ddpms paper https arxiv org abs 2006 11239 for augmenting training data we randomly flip the images left right load the dataset flips an image left right randomly resize the image to the desired size first and then rescale the pixel values in the range 1 0 1 0 args img image tensor size desired image size for resizing returns resized and rescaled image tensor resize rescale the pixel values gaussian diffusion utilities we define the forward process and the reverse process as a separate utility most of the code in this utility has been borrowed from the original implementation with some slight modifications gaussian diffusion utility args beta_start start value of the scheduled variance beta_end end value of the scheduled variance timesteps number of time steps in the forward process define the linear variance schedule using float64 for better precision calculations for diffusion q x_t x_ t 1 and others calculations for posterior q x_ t 1 x_t x_0 log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain extract some coefficients at specified timesteps then reshape to batch_size 1 1 1 1 for broadcasting purposes args a tensor to extract from t timestep for which the coefficients are to be extracted x_shape shape of the current batched samples extracts the mean and the variance at current timestep args x_start initial sample before the first diffusion step t current timestep diffuse the data args x_start initial sample before the first diffusion step t current timestep noise gaussian noise to be added at the current timestep returns diffused samples at timestep t compute the mean and variance of the diffusion posterior q x_ t 1 x_t x_0 args x_start stating point sample for the posterior computation x_t sample at timestep t t current timestep returns posterior mean and variance at current timestep sample from the diffuison model args pred_noise noise predicted by the diffusion model x samples at a given timestep for which the noise was predicted t current timestep clip_denoised bool whether to clip the predicted noise within the specified range or not no noise when t 0 network architecture u net originally developed for semantic segmentation is an architecture that is widely used for implementing diffusion models but with some slight modifications 1 the network accepts two inputs image and time step 2 self attention between the convolution blocks once we reach a specific resolution 16x16 in the paper 3 group normalization instead of weight normalization we implement most of the things as used in the original paper we use the swish activation function throughout the network we use the variance scaling kernel initializer the only difference here is the number of groups used for the groupnormalization layer for the flowers dataset we found that a value of groups 8 produces better results compared to the default value of groups 32 dropout is optional and should be used where chances of over fitting is high in the paper the s used dropout only when training on cifar10 kernel initializer to use applies self attention args units number of units in the dense layers groups number of groups to be used for groupnormalization layer downblock middleblock upblock end block training we follow the same setup for training the diffusion model as described in the paper we use adam optimizer with a learning rate of 2e 4 we use ema on model parameters with a decay factor of 0 999 we treat our model as noise prediction network i e at every training step we input a batch of images and corresponding time steps to our unet and the network outputs the noise as predictions the only difference is that we aren t using the kernel inception distance kid or frechet inception distance fid for evaluating the quality of generated samples during training this is because both these metrics are compute heavy and are skipped for the brevity of implementation note we are using mean squared error as the loss function which is aligned with the paper and theoretically makes sense in practice though it is also common to use mean absolute error or huber loss as the loss function 1 get the batch size 2 sample timesteps uniformly 3 sample random noise to be added to the images in the batch 4 diffuse the images with noise 5 pass the diffused images and time steps to the network 6 calculate the loss 7 get the gradients 8 update the weights of the network 9 updates the weight values for the network with ema weights 10 return loss values 1 randomly sample noise starting point for reverse process 2 sample from the model iteratively 3 return generated samples utility to plot images using the diffusion model during training build the unet model initially the weights are the same get an instance of the gaussian diffusion utilities get the model compile the model train the model results we trained this model for 800 epochs on a v100 gpu and each epoch took almost 8 seconds to finish we load those weights here and we generate a few samples starting from pure noise shell curl lo https github com aakashkumarnain ddpms releases download v3 0 0 checkpoints zip unzip qq checkpoints zip load the model weights generate and plot some samples conclusion we successfully implemented and trained a diffusion model exactly in the same fashion as implemented by the s of the ddpms paper you can find the original implementation here https github com hojonathanho diffusion there are a few things that you can try to improve the model 1 increasing the width of each block a bigger model can learn to denoise in fewer epochs though you may have to take care of overfitting 2 we implemented the linear schedule for variance scheduling you can implement other schemes like cosine scheduling and compare the performance references 1 denoising diffusion probabilistic models https arxiv org abs 2006 11239 2 s implementation https github com hojonathanho diffusion 3 a deep dive into ddpms https magic with latents github io latent posts ddpms part3 4 denoising diffusion implicit models https keras io examples generative ddim 5 annotated diffusion model https huggingface co blog annotated diffusion 6 aiaiart https www youtube com watch v xts7m6tsk9i t 14s
import math
import numpy as np
import matplotlib.pyplot as plt


import tensorflow as tf
import keras
from keras import layers
import tensorflow_datasets as tfds



batch_size = 32
num_epochs = 1  
total_timesteps = 1000
norm_groups = 8  
learning_rate = 2e-4

img_size = 64
img_channels = 3
clip_min = -1.0
clip_max = 1.0

first_conv_channels = 64
channel_multiplier = [1, 2, 4, 8]
widths = [first_conv_channels * mult for mult in channel_multiplier]
has_attention = [False, False, True, True]
num_res_blocks = 2  

dataset_name = "oxford_flowers102"
splits = ["train"]






(ds,) = tfds.load(
    dataset_name, split=splits, with_info=False, shuffle_files=True
)


def augment(img):
    
    return tf.image.random_flip_left_right(img)


def resize_and_rescale(img, size):
    

    height = tf.shape(img)[0]
    width = tf.shape(img)[1]
    crop_size = tf.minimum(height, width)

    img = tf.image.crop_to_bounding_box(
        img,
        (height - crop_size) // 2,
        (width - crop_size) // 2,
        crop_size,
        crop_size,
    )

    
    img = tf.cast(img, dtype=tf.float32)
    img = tf.image.resize(img, size=size, antialias=True)

    
    img = img / 127.5 - 1.0
    img = tf.clip_by_value(img, clip_min, clip_max)
    return img


def train_preprocessing(x):
    img = x["image"]
    img = resize_and_rescale(img, size=(img_size, img_size))
    img = augment(img)
    return img


train_ds = (
    ds.map(train_preprocessing, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(batch_size, drop_remainder=True)
    .shuffle(batch_size * 2)
    .prefetch(tf.data.AUTOTUNE)
)





class GaussianDiffusion:
    

    def __init__(
        self,
        beta_start=1e-4,
        beta_end=0.02,
        timesteps=1000,
        clip_min=-1.0,
        clip_max=1.0,
    ):
        self.beta_start = beta_start
        self.beta_end = beta_end
        self.timesteps = timesteps
        self.clip_min = clip_min
        self.clip_max = clip_max

        
        self.betas = betas = np.linspace(
            beta_start,
            beta_end,
            timesteps,
            dtype=np.float64,  
        )
        self.num_timesteps = int(timesteps)

        alphas = 1.0 - betas
        alphas_cumprod = np.cumprod(alphas, axis=0)
        alphas_cumprod_prev = np.append(1.0, alphas_cumprod[:-1])

        self.betas = tf.constant(betas, dtype=tf.float32)
        self.alphas_cumprod = tf.constant(alphas_cumprod, dtype=tf.float32)
        self.alphas_cumprod_prev = tf.constant(
            alphas_cumprod_prev, dtype=tf.float32
        )

        
        self.sqrt_alphas_cumprod = tf.constant(
            np.sqrt(alphas_cumprod), dtype=tf.float32
        )

        self.sqrt_one_minus_alphas_cumprod = tf.constant(
            np.sqrt(1.0 - alphas_cumprod), dtype=tf.float32
        )

        self.log_one_minus_alphas_cumprod = tf.constant(
            np.log(1.0 - alphas_cumprod), dtype=tf.float32
        )

        self.sqrt_recip_alphas_cumprod = tf.constant(
            np.sqrt(1.0 / alphas_cumprod), dtype=tf.float32
        )
        self.sqrt_recipm1_alphas_cumprod = tf.constant(
            np.sqrt(1.0 / alphas_cumprod - 1), dtype=tf.float32
        )

        
        posterior_variance = (
            betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
        )
        self.posterior_variance = tf.constant(
            posterior_variance, dtype=tf.float32
        )

        
        
        self.posterior_log_variance_clipped = tf.constant(
            np.log(np.maximum(posterior_variance, 1e-20)), dtype=tf.float32
        )

        self.posterior_mean_coef1 = tf.constant(
            betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod),
            dtype=tf.float32,
        )

        self.posterior_mean_coef2 = tf.constant(
            (1.0 - alphas_cumprod_prev)
            * np.sqrt(alphas)
            / (1.0 - alphas_cumprod),
            dtype=tf.float32,
        )

    def _extract(self, a, t, x_shape):
        
        batch_size = x_shape[0]
        out = tf.gather(a, t)
        return tf.reshape(out, [batch_size, 1, 1, 1])

    def q_mean_variance(self, x_start, t):
        
        x_start_shape = tf.shape(x_start)
        mean = (
            self._extract(self.sqrt_alphas_cumprod, t, x_start_shape) * x_start
        )
        variance = self._extract(1.0 - self.alphas_cumprod, t, x_start_shape)
        log_variance = self._extract(
            self.log_one_minus_alphas_cumprod, t, x_start_shape
        )
        return mean, variance, log_variance

    def q_sample(self, x_start, t, noise):
        
        x_start_shape = tf.shape(x_start)
        return (
            self._extract(self.sqrt_alphas_cumprod, t, tf.shape(x_start))
            * x_start
            + self._extract(
                self.sqrt_one_minus_alphas_cumprod, t, x_start_shape
            )
            * noise
        )

    def predict_start_from_noise(self, x_t, t, noise):
        x_t_shape = tf.shape(x_t)
        return (
            self._extract(self.sqrt_recip_alphas_cumprod, t, x_t_shape) * x_t
            - self._extract(self.sqrt_recipm1_alphas_cumprod, t, x_t_shape)
            * noise
        )

    def q_posterior(self, x_start, x_t, t):
        

        x_t_shape = tf.shape(x_t)
        posterior_mean = (
            self._extract(self.posterior_mean_coef1, t, x_t_shape) * x_start
            + self._extract(self.posterior_mean_coef2, t, x_t_shape) * x_t
        )
        posterior_variance = self._extract(
            self.posterior_variance, t, x_t_shape
        )
        posterior_log_variance_clipped = self._extract(
            self.posterior_log_variance_clipped, t, x_t_shape
        )
        return (
            posterior_mean,
            posterior_variance,
            posterior_log_variance_clipped,
        )

    def p_mean_variance(self, pred_noise, x, t, clip_denoised=True):
        x_recon = self.predict_start_from_noise(x, t=t, noise=pred_noise)
        if clip_denoised:
            x_recon = tf.clip_by_value(x_recon, self.clip_min, self.clip_max)

        (
            model_mean,
            posterior_variance,
            posterior_log_variance,
        ) = self.q_posterior(x_start=x_recon, x_t=x, t=t)
        return model_mean, posterior_variance, posterior_log_variance

    def p_sample(self, pred_noise, x, t, clip_denoised=True):
        
        model_mean, _, model_log_variance = self.p_mean_variance(
            pred_noise, x=x, t=t, clip_denoised=clip_denoised
        )
        noise = tf.random.normal(shape=x.shape, dtype=x.dtype)
        
        nonzero_mask = tf.reshape(
            1 - tf.cast(tf.equal(t, 0), tf.float32), [tf.shape(x)[0], 1, 1, 1]
        )
        return (
            model_mean + nonzero_mask * tf.exp(0.5 * model_log_variance) * noise
        )






def kernel_init(scale):
    scale = max(scale, 1e-10)
    return keras.initializers.VarianceScaling(
        scale, mode="fan_avg", distribution="uniform"
    )


class AttentionBlock(layers.Layer):
    

    def __init__(self, units, groups=8, **kwargs):
        self.units = units
        self.groups = groups
        super().__init__(**kwargs)

        self.norm = layers.GroupNormalization(groups=groups)
        self.query = layers.Dense(units, kernel_initializer=kernel_init(1.0))
        self.key = layers.Dense(units, kernel_initializer=kernel_init(1.0))
        self.value = layers.Dense(units, kernel_initializer=kernel_init(1.0))
        self.proj = layers.Dense(units, kernel_initializer=kernel_init(0.0))

    def call(self, inputs):
        batch_size = tf.shape(inputs)[0]
        height = tf.shape(inputs)[1]
        width = tf.shape(inputs)[2]
        scale = tf.cast(self.units, tf.float32) ** (-0.5)

        inputs = self.norm(inputs)
        q = self.query(inputs)
        k = self.key(inputs)
        v = self.value(inputs)

        attn_score = tf.einsum("bhwc, bHWc->bhwHW", q, k) * scale
        attn_score = tf.reshape(
            attn_score, [batch_size, height, width, height * width]
        )

        attn_score = tf.nn.softmax(attn_score, -1)
        attn_score = tf.reshape(
            attn_score, [batch_size, height, width, height, width]
        )

        proj = tf.einsum("bhwHW,bHWc->bhwc", attn_score, v)
        proj = self.proj(proj)
        return inputs + proj


class TimeEmbedding(layers.Layer):
    def __init__(self, dim, **kwargs):
        super().__init__(**kwargs)
        self.dim = dim
        self.half_dim = dim // 2
        self.emb = math.log(10000) / (self.half_dim - 1)
        self.emb = tf.exp(tf.range(self.half_dim, dtype=tf.float32) * -self.emb)

    def call(self, inputs):
        inputs = tf.cast(inputs, dtype=tf.float32)
        emb = inputs[:, None] * self.emb[None, :]
        emb = tf.concat([tf.sin(emb), tf.cos(emb)], axis=-1)
        return emb


def ResidualBlock(width, groups=8, activation_fn=keras.activations.swish):
    def apply(inputs):
        x, t = inputs
        input_width = x.shape[3]

        if input_width == width:
            residual = x
        else:
            residual = layers.Conv2D(
                width, kernel_size=1, kernel_initializer=kernel_init(1.0)
            )(x)

        temb = activation_fn(t)
        temb = layers.Dense(width, kernel_initializer=kernel_init(1.0))(temb)[
            :, None, None, :
        ]

        x = layers.GroupNormalization(groups=groups)(x)
        x = activation_fn(x)
        x = layers.Conv2D(
            width,
            kernel_size=3,
            padding="same",
            kernel_initializer=kernel_init(1.0),
        )(x)

        x = layers.Add()([x, temb])
        x = layers.GroupNormalization(groups=groups)(x)
        x = activation_fn(x)

        x = layers.Conv2D(
            width,
            kernel_size=3,
            padding="same",
            kernel_initializer=kernel_init(0.0),
        )(x)
        x = layers.Add()([x, residual])
        return x

    return apply


def DownSample(width):
    def apply(x):
        x = layers.Conv2D(
            width,
            kernel_size=3,
            strides=2,
            padding="same",
            kernel_initializer=kernel_init(1.0),
        )(x)
        return x

    return apply


def UpSample(width, interpolation="nearest"):
    def apply(x):
        x = layers.UpSampling2D(size=2, interpolation=interpolation)(x)
        x = layers.Conv2D(
            width,
            kernel_size=3,
            padding="same",
            kernel_initializer=kernel_init(1.0),
        )(x)
        return x

    return apply


def TimeMLP(units, activation_fn=keras.activations.swish):
    def apply(inputs):
        temb = layers.Dense(
            units, activation=activation_fn, kernel_initializer=kernel_init(1.0)
        )(inputs)
        temb = layers.Dense(units, kernel_initializer=kernel_init(1.0))(temb)
        return temb

    return apply


def build_model(
    img_size,
    img_channels,
    widths,
    has_attention,
    num_res_blocks=2,
    norm_groups=8,
    interpolation="nearest",
    activation_fn=keras.activations.swish,
):
    image_input = layers.Input(
        shape=(img_size, img_size, img_channels), name="image_input"
    )
    time_input = keras.Input(shape=(), dtype=tf.int64, name="time_input")

    x = layers.Conv2D(
        first_conv_channels,
        kernel_size=(3, 3),
        padding="same",
        kernel_initializer=kernel_init(1.0),
    )(image_input)

    temb = TimeEmbedding(dim=first_conv_channels * 4)(time_input)
    temb = TimeMLP(units=first_conv_channels * 4, activation_fn=activation_fn)(
        temb
    )

    skips = [x]

    
    for i in range(len(widths)):
        for _ in range(num_res_blocks):
            x = ResidualBlock(
                widths[i], groups=norm_groups, activation_fn=activation_fn
            )([x, temb])
            if has_attention[i]:
                x = AttentionBlock(widths[i], groups=norm_groups)(x)
            skips.append(x)

        if widths[i] != widths[-1]:
            x = DownSample(widths[i])(x)
            skips.append(x)

    
    x = ResidualBlock(
        widths[-1], groups=norm_groups, activation_fn=activation_fn
    )([x, temb])
    x = AttentionBlock(widths[-1], groups=norm_groups)(x)
    x = ResidualBlock(
        widths[-1], groups=norm_groups, activation_fn=activation_fn
    )([x, temb])

    
    for i in reversed(range(len(widths))):
        for _ in range(num_res_blocks + 1):
            x = layers.Concatenate(axis=-1)([x, skips.pop()])
            x = ResidualBlock(
                widths[i], groups=norm_groups, activation_fn=activation_fn
            )([x, temb])
            if has_attention[i]:
                x = AttentionBlock(widths[i], groups=norm_groups)(x)

        if i != 0:
            x = UpSample(widths[i], interpolation=interpolation)(x)

    
    x = layers.GroupNormalization(groups=norm_groups)(x)
    x = activation_fn(x)
    x = layers.Conv2D(
        3, (3, 3), padding="same", kernel_initializer=kernel_init(0.0)
    )(x)
    return keras.Model([image_input, time_input], x, name="unet")





class DiffusionModel(keras.Model):
    def __init__(self, network, ema_network, timesteps, gdf_util, ema=0.999):
        super().__init__()
        self.network = network
        self.ema_network = ema_network
        self.timesteps = timesteps
        self.gdf_util = gdf_util
        self.ema = ema

    def train_step(self, images):
        
        batch_size = tf.shape(images)[0]

        
        t = tf.random.uniform(
            minval=0, maxval=self.timesteps, shape=(batch_size,), dtype=tf.int64
        )

        with tf.GradientTape() as tape:
            
            noise = tf.random.normal(shape=tf.shape(images), dtype=images.dtype)

            
            images_t = self.gdf_util.q_sample(images, t, noise)

            
            pred_noise = self.network([images_t, t], training=True)

            
            loss = self.loss(noise, pred_noise)

        
        gradients = tape.gradient(loss, self.network.trainable_weights)

        
        self.optimizer.apply_gradients(
            zip(gradients, self.network.trainable_weights)
        )

        
        for weight, ema_weight in zip(
            self.network.weights, self.ema_network.weights
        ):
            ema_weight.assign(self.ema * ema_weight + (1 - self.ema) * weight)

        
        return {"loss": loss}

    def generate_images(self, num_images=16):
        
        samples = tf.random.normal(
            shape=(num_images, img_size, img_size, img_channels),
            dtype=tf.float32,
        )
        
        for t in reversed(range(0, self.timesteps)):
            tt = tf.cast(tf.fill(num_images, t), dtype=tf.int64)
            pred_noise = self.ema_network.predict(
                [samples, tt], verbose=0, batch_size=num_images
            )
            samples = self.gdf_util.p_sample(
                pred_noise, samples, tt, clip_denoised=True
            )
        
        return samples

    def plot_images(
        self, epoch=None, logs=None, num_rows=2, num_cols=8, figsize=(12, 5)
    ):
        
        generated_samples = self.generate_images(num_images=num_rows * num_cols)
        generated_samples = (
            tf.clip_by_value(generated_samples * 127.5 + 127.5, 0.0, 255.0)
            .numpy()
            .astype(np.uint8)
        )

        _, ax = plt.subplots(num_rows, num_cols, figsize=figsize)
        for i, image in enumerate(generated_samples):
            if num_rows == 1:
                ax[i].imshow(image)
                ax[i].axis("off")
            else:
                ax[i // num_cols, i % num_cols].imshow(image)
                ax[i // num_cols, i % num_cols].axis("off")

        plt.tight_layout()
        plt.show()



network = build_model(
    img_size=img_size,
    img_channels=img_channels,
    widths=widths,
    has_attention=has_attention,
    num_res_blocks=num_res_blocks,
    norm_groups=norm_groups,
    activation_fn=keras.activations.swish,
)
ema_network = build_model(
    img_size=img_size,
    img_channels=img_channels,
    widths=widths,
    has_attention=has_attention,
    num_res_blocks=num_res_blocks,
    norm_groups=norm_groups,
    activation_fn=keras.activations.swish,
)
ema_network.set_weights(
    network.get_weights()
)  


gdf_util = GaussianDiffusion(timesteps=total_timesteps)


model = DiffusionModel(
    network=network,
    ema_network=ema_network,
    gdf_util=gdf_util,
    timesteps=total_timesteps,
)


model.compile(
    loss=keras.losses.MeanSquaredError(),
    optimizer=keras.optimizers.Adam(learning_rate=learning_rate),
    jit_compile=False,
)


model.fit(
    train_ds,
    epochs=num_epochs,
    batch_size=batch_size,
    callbacks=[keras.callbacks.LambdaCallback(on_epoch_end=model.plot_images)],
)






model.ema_network.load_weights("checkpoints/diffusion_model_checkpoint")


model.plot_images(num_rows=4, num_cols=8)