前言

在扩散模型的学习过程中，up主deep_thoughts的视频给了我很大的帮助，下面是我在学习过程中观看的一些视频，大家也可以按照这个顺序进行学习。DDPM论文链接也在此给大家列出：https://arxiv.org/abs/2006.11239

up主deep_thoughts
————————
54 DDPM 原来讲解和demo代码演示
58 借助IDDPM代码讲解DDPM
64 DDIM原理和代码
————————
62 分数扩散模型一
63 分数扩散模型二
————————
66 Guided Diffusion

注：本文章难度较大，篇幅较长，大家在阅读过程中需要一些耐心！

一、前置数学知识

在学习DDPM原理之前，了解一些数学概率论的知识是必要的，因为在DDPM中涉及到很多概率公式的推导。

1、先验概率和后验概率
先验概率：根据以往经验和分析得到的概率,它往往作为“由因求果”问题中的“因”出现，如 $q(x_t|x_{t-1})$
后验概率：指在得到“结果”的信息后重新修正的概率,是“执果寻因”问题中的“因", 如$p(x_{t-1}|x_t)$

2、条件概率
设A、B为任意两个事件，若P(A)>0,称在已知事件A发生的条件下，事件B发生的概率为条件概率，记为$P(B|A)$
$$P(B|A)=\frac{P(A,B)}{P(A)}$$

3、乘法公式：
$$P(A,B)=P(B|A)P(A)$$

4、乘法公式一般形式：
$$P(A,B,C)=P(C|B,A)P(B,A)=P(C|B,A)P(B|A)P(A)$$

5、贝叶斯公式：
$$P(A|B)=\frac{P(B|A)P(A)}{P(B)}$$

6、多元贝叶斯公式:
$$P(A|B,C)=\frac{P(A,B,C)}{P(B,C)}=\frac{P(B|A,C)P(A,C)}{P(B,C)}=\frac{P(B|A,C)P(A|C)P(C)}{P(B|C)P(C)}=\frac{P(B|A,C)P(A|C))}{P(B|C)}$$

7、正态分布的叠加性：
当有两个独立的正态分布变量$N_1$和$N_2$,它们的均值和方差分别为${\mu }_1,{\mu }_2$和${\sigma }_1^2,{\sigma }_2^2$它们的和为$N=a{N}_1+b{N}_2$的均值和方差可以表示如下：
$$\begin{gathered} E(N)=E(a{N}_1+b{N}_2)=a{\mu }_1+b{\mu }_2 \\ Var(N)=Var(a{N}_1+b{N}_2)=a^2{\sigma }_1^2+b^2{\sigma }_2^2 \end{gathered}$$
相减时：
$$\begin{gathered} E(N)=E(a{N}_1-b{N}_2)=a{\mu }_1-b{\mu }_2 \\ Var(N)=Var(a{N}_1-b{N}_2)=a^2{\sigma }_1^2+b^2{\sigma }_2^2 \end{gathered}$$

8、参数重整化(重参数化）：
若希望从高斯分布$N(\mu ,\sigma )$中采样，可以先从标准分布$N(0,1)$采样出$z$,再得到$\sigma \ast z+\mu$,这样做的好处是将随机性转移到$z$这个常量上，而$\sigma$和$\mu$则当做仿射变换网络的一部分。

9、高斯分布的概率密度函数
$$f(x)=\frac{1}{\sqrt{2\pi }\sigma }{e}^{-\frac{(x-\mu {)}^2}{2{\sigma }^2}}$$

10、高斯分布的KL散度公式
对于两个单一变量的高斯分布p和q而言，它们的KL散度公式：
$$KL(p|q)=log\frac{{\sigma }_2}{{\sigma }_1}+\frac{{\sigma }^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}-\frac{1}{2}$$

11、二次函数配方:$a{x}^2+bx=a(x+\frac{b}{2a}{)}^2+c$

12、随机变量的期望公式 设$X$是随机变量,$Y=g(X)$,则：$$E(Y)=E[g(X)]=\{\begin{array}{l}\sum _{k=1}^{\infty }g(x_k)p_k\\ {\int }_{-\infty }^{\infty }g(x)p(x)dx\end{array}$$

13、KL散度公式：$KL(p(x)|q(x))=E_{x\sim p(x)}[\frac{p(x)}{q(x)}]=\int p(x)\frac{p(x)}{q(x)}dx$

二、VAE模型回顾

$p(x)={\int }_z{p}_{\theta }(x|z)p(z)$
$p(x)=\int q_{\Phi }(z|x)\frac{p_{\theta }(x|z)p(z)}{q_{\Phi }(z|x)}$
$logp(x)=log{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}]$
$logp(x)\ge log{\mathbb{E}}_{z\sim q_{\varphi }(z|x)}[\frac{p_{\theta }(x|z)p(z)}{q_{\varphi }(z|x)}]$

三、DDPM介绍

2020年Berkeley大学的学生提出的DDPM（Denoising Diffusion Probabilistic Models），简称扩散模型，是AIGC的核心算法，在生成图像的真实性和多样性方面均超越了GAN，而且训练过程稳定。

扩散模型包括两个过程：前向扩散过程(前向加噪过程)和反向去噪过程

前向过程和反向过程都是马尔可夫链（有人可能对马尔科夫链比较陌生，这里打个不太恰当的比方，就好比明天的天气如何只以今天的天气作为变量和前天还有大前天都无关，也就是只和离他最近的一个结果有关，全过程大约需要1000步，其中反向过程用来生成数据，它的推导过程可以描述成：

四、前向扩散的过程

前向扩散过程是对原始数据逐渐增加高斯噪声，直至变成标准高斯分布的过程。

1）给定初始数据分布$x_0\sim q(x)$，可以不断地向分布中添加高斯噪声，该噪声的标准差是以固定值${\beta }_t$而确定的，均值是以固定值${\beta }_t$和当前t-1时刻的数据$x_{t-1}$决定的，这个过程是一个马尔可夫过程。

2）随着t不断增大，最终数据分布$x_T$变成了一个各项独立的高斯分布。

$$\begin{array}{rl}q(x_{1:T}|x_0)& =\prod _{t=1}^{T}q(x_t|x_{t-1})\\ q(x_t|x_{t-1})& =\mathcal{N}(x_t;\sqrt{{\alpha }_t}{x}_{t-1},{\beta }_{t}I)\end{array}$$

注：${\alpha }_t=1-{\beta }_t$！！！！这是作者设定

将$x_t$ 进行参数重整化，得到$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{{\beta }_t}{ϵ}_t{ϵ}_t\sim \mathcal{N}(0,I)\end{array}$

3）任意时刻的$q(x_t)$ 推导也可以完全基于$x_0$和${\beta }_t$来计算出来，而不需要迭代

$$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{{\beta }_t}{ϵ}_t\\ & =\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\beta }_{t-1}}{ϵ}_{t-1})+\sqrt{{\beta }_t}{ϵ}_t\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}(\sqrt{{\alpha }_{t-2}}{x}_{t-3}+\sqrt{{\beta }_{t-2}}{ϵ}_{t-2})+\sqrt{{\alpha }_t{\beta }_{t-1}}{ϵ}_{t-1}+\sqrt{{\beta }_t}{ϵ}_t\\ & =\sqrt{({\alpha }_t\dots {\alpha }_1)}{x}_0+\sqrt{({\alpha }_t\dots {\alpha }_2){\beta }_1}{ϵ}_1+\sqrt{({\alpha }_t\dots {\alpha }_3){\beta }_2}{ϵ}_2+\cdots +\sqrt{{\alpha }_t{\beta }_{t-1}}{ϵ}_{t-1}+\sqrt{{\beta }_t}{ϵ}_t\end{array}$$

设：$\overline{{\alpha }_t}={\alpha }_1{\alpha }_2\dots {\alpha }_t$,接下来需要用到正态分布的叠加性

注意这里，两个正态分布$X\sim N({\mu }_1,{\sigma }_1)$和$Y\sim N({\mu }_2,{\sigma }_2)$的叠加后的分布$aX+bY$的均值为$a{\mu }_1+b{\mu }_2$,方差为$a^2{\sigma }_1^2+b^2{\sigma }_2^2$。所以$\sqrt{a_t-a_t{a}_{t-1}}{ϵ}_{t-2}+\sqrt{1-a_t}{ϵ}_{t-1}$可以重参数化成只含一个随机变量$ϵ$构成的$\sqrt{1-a_t{a}_{t-1}}{ϵ}_{t-1}$
因此均值和方差可得：
$$\begin{array}{rl}\mu & =\sqrt{\overline{{\alpha }_t}}\\ {\sigma }^2& =({\alpha }_t\dots {\alpha }_2){\beta }_1+({\alpha }_t\dots {\alpha }_3){\beta }_2+\cdots +{\alpha }_t{\beta }_{t-1}+{\beta }_t\\ & =({\alpha }_t\dots {\alpha }_2)(1-{\alpha }_1)+({\alpha }_t\dots {\alpha }_3)(1-{\alpha }_2)+\cdots +{\alpha }_t(1-{\alpha }_{t-1})+1-{\alpha }_t\\ & =1-\overline{{\alpha }_t}\end{array}$$

从而得到前向过程的最终表达式
$$\begin{array}{rl}{x}_t& =\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-\overline{{\alpha }_t}}ϵϵ\sim \mathcal{N}(0,I)\\ q(x_t|x_0)& =\mathcal{N}(x_t;\sqrt{\overline{{\alpha }_t}}{x}_0,(1-\overline{{\alpha }_t})I)\end{array}$$

这个公式表示任意步骤t的噪声图像$x_t$，都可以通过$x_0$直接加噪得到，后面需要用到。

五、反向去噪过程,神经网络拟合过程

反向去噪过程就是数据生成过程，它首先是从标准高斯分布中采样得到一个噪声样本，再一步步地迭代去噪，最后得到数据分布中的一个样本。

如果知道反向过程的每一步真实的条件分布$q(x_{t-1}|x_t)$ ，那么从一个随机噪声开始，逐步采样就能生成一个真实的样本。但是真实的条件分布利用贝叶斯公式$q(x_{t-1}|x_t)=\frac{q(x_t|x_{t-1})q(x_{t-1})}{q(x_t)}$无法直接求解，原因是其中 $q(x_{t-1}),q(x_t)$未知，因此无法从$x_t$推导到$x_{t-1}$，所以必须通过神经网络$p_{\theta }(x_{t-1}|x_t)$ 来近似。为了简化起见，将反向过程也定义为一个马尔卡夫链，且服从高斯分布，建模如下:
$$\begin{gathered} p_{\theta }(x_{0:T})=p(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t) \\ {p}_{\theta }(x_{t-1}|x_t)=N(x_{t-1};{\mu }_{\theta }(x_t,t),\sum _{\theta }(x_t,t)) \end{gathered}$$

虽然真实条件分布$q(x_{t-1}|x_t)$无法直接求解，但是加上已知条件$x_0$的后验分布$q(x_{t-1}|x_t,x_0)$却可以通过贝叶斯公式求解，再结合前向马尔科夫性质可得：
$$q(x_{t-1}|x_t,x_0)=\frac{q(x_t|x_{t-1},x_0)q(x_{t-1}|x_0)}{q(x_t|x_0)}=\frac{q(x_t|x_{t-1})q(x_{t-1}|x_0)}{q(x_t|x_0)}$$
因此可以得到：
$$\begin{array}{rl}q(x_{t-1}|x_0)& =\sqrt{{\bar{\alpha }}_{t-1}}{x}_0+\sqrt{1-{\bar{\alpha }}_{t-1}}ϵ\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_{t-1}}{x}_0,(1-{\bar{\alpha }}_{t-1})I)\\ q(x_t|x_0)& =\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I)\\ q(x_t|x_{t-1})& ={\sqrt{\alpha }}_t{x}_{t-1}+{\beta }_tϵ\sim \mathcal{N}({\sqrt{\alpha }}_t{x}_{t-1},{\beta }_{t}I)\end{array}$$
接下来需要用到：

所以
$$\begin{array}{rl}q(x_{t-1}|x_t,x_0)& \propto exp(-\frac{1}{2}(\frac{(x_t-\sqrt{{\alpha }_t}{x}_{t-1}{)}^2}{{\beta }_t})+\frac{(x_{t-1}-{\sqrt{\bar{\alpha }}}_{t-1}{x}_0{)}^2}{1-{\bar{\alpha }}_{t-1}}-\frac{(x_t-\sqrt{{\bar{\alpha }}_t}{x}_0{)}^2}{1-{\bar{\alpha }}_t})\\ & =exp(-\frac{1}{2}(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})x_{t-1}^2-(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{2\sqrt{\overline{{\alpha }_t}}}{1-\overline{{\alpha }_t}}{x}_0)x_{t-1}+C(x_t,x_0))\end{array}$$
通过配方就可以得到
$$\begin{gathered} {\tilde{\beta }}_t=1/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t \\ {\tilde{\mu }}_t=(\frac{{\sqrt{\alpha }}_t}{{\beta }_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_t}}{1-\overline{{\alpha }_t}}{x}_0)/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})=\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-\overline{{\alpha }_t}}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0 \end{gathered}$$
又因为由上面 $x_0$和$x_t$的直接关系
$$x_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}ϵ)$$
代入上式 ${\tilde{\mu }}_t=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{(1-{\alpha }_t)}}ϵ)$

这样我们就得到了后向过程$q(x_{t-1}|x_t,x_0)$的均值和方差的公式，为了让$p_{\theta }(x_{t-1}|x_t)$逼近$q(x_{t-1}|x_t,x_0)$
我们固定住${\sum }_{\theta }(x_t,t)={\beta }_t$和${\sum }_{\theta }(x_t,t)={\tilde{\beta }}_t$,${\beta }_t$和${\tilde{\beta }}_t$其实大小是差不多的，但其实用${\tilde{\beta }}_t$更好理解一点，因此只需要预测${\mu }_{\theta }(x_t,t)$，又因为这里面只有$ϵ$是未知的，所以转而预测$ϵ$，因此均值可以写成下面的式子：
$${\mu }_{\theta }(x_t,t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{(1-{\alpha }_t)}}{ϵ}_{\theta }(x_t,t))$$
注：${\tilde{\beta }}_t$方差中没有未知参数

六、采样过程（模型训练完后的预测过程）

$$\begin{gathered} p_{\theta }(x_{t-1}|x_t)=N(x_{t-1};{\mu }_{\theta }(x_t,t),\sum _{\theta }(x_t,t)) \\ {x}_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{(1-{\alpha }_t)}}{ϵ}_{\theta }(x_t,t))+\sqrt{{\tilde{\beta }}_t}zz\sim N(0,I) \end{gathered}$$

这里用z是为了和之前的$ϵ$区别开，都是符合正态分布的随机噪声，迭代1000次

七、优化目标

注意:$q(x_t|x_{t-1})=q(x_t|x_{t-1},x_0)=\frac{q(x_t,x_{t-1},x_0)}{q(x_{t-1},x_0)}=\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)q(x_0)}{q(x_{t-1},x_0)}=\frac{q(x_{t-1}|x_t,x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}$

Code1:二维点demo

import torch
from matplotlib import pyplot as plt
import numpy as np
from sklearn.datasets import make_s_curve
from tqdm import tqdm

s_surve,_=make_s_curve(n_samples=10000,noise=0.1)
s_surve=s_surve[:,[0,2]]/10.0
print("shape of moons:",np.shape(s_surve))
data=s_surve.T
print(data.shape)
fig,ax=plt.subplots()
ax.scatter(*data,color="red",edgecolors="white")
ax.axis("off")

dataset = torch.Tensor(s_surve).float()
#制定每一步的beta
num_steps=100  
betas=torch.linspace(-6,6,num_steps)
betas=torch.sigmoid(betas) *(0.5e-2-1e-5)+1e-5
#计算alpha、alpha_prod、alpha_prod_previous、alpha_bar_sqrt等变量的值
alphas=1-betas
alphas_prod=torch.cumprod(alphas,0)
alphas_prod_prev=torch.cat([torch.tensor([1]).float(),alphas_prod[:-1]],dim=0)
alphas_bar_sqrt=torch.sqrt(alphas_prod)
one_minus_alphas_bar_log=torch.log(1-alphas_prod)
one_minus_alphas_bar_sqrt=torch.sqrt(1-alphas_prod)
assert alphas.shape==alphas_prod.shape==alphas_prod_prev.shape==alphas_bar_sqrt.shape==one_minus_alphas_bar_log.shape==one_minus_alphas_bar_sqrt.shape
print(alphas.shape)

def q_x(x_0,t):
    noise=torch.randn_like(x_0)
    alphas_t=alphas_bar_sqrt[t]
    alphas_1_m_t=one_minus_alphas_bar_sqrt[t]
    return (alphas_t*x_0+alphas_1_m_t*noise)

num_shows=20
fig,axs=plt.subplots(2,10,figsize=(28,3))
plt.rc("text",color="blue")
#共有10000个点，每个店包含两个坐标
#生成100步以内每隔5步加噪声后的图像

for i in range(num_shows):
    j=i//10
    k=i%10
    q_i=q_x(dataset,torch.tensor([i*num_steps//num_shows]))
    axs[j,k].scatter(q_i[:,0],q_i[:,1],color="red",edgecolors="white")
    axs[j,k].axis("off")
    axs[j,k].set_title(str(i))
import torch 
import torch.nn as nn

class MLPDiffusion(nn.Module):
    def __init__(self,n_steps,num_units=128):
        super().__init__()
        self.linears=nn.ModuleList(
            [
            nn.Linear(2,num_units),
            nn.ReLU(),
            nn.Linear(num_units,num_units),
            nn.ReLU(),
            nn.Linear(num_units,num_units),
            nn.ReLU(),
            nn.Linear(num_units,2)
            ]
            
        )
        self.step_embeddings=nn.ModuleList(
            [
                nn.Embedding(n_steps,num_units),
                nn.Embedding(n_steps,num_units),
                nn.Embedding(n_steps,num_units),
            ]
        )
    def forward(self,x,t):
        for idx,embedding_layer in enumerate(self.step_embeddings):
            t_embedding=embedding_layer(t)
            x=self.linears[2*idx](x)
            x+=t_embedding
            x=self.linears[2*idx+1](x)
        x=self.linears[-1](x)
        return x

def diffusion_loss_fn(model,x_0,alphas_bar_sqrt,one_minus_alphas_bar_sqrt,n_steps):
    batch_size=x_0.shape[0]
    #对一个batchsize样本生成随机时刻t
    t=torch.randint(0,n_steps,size=(batch_size//2,))
    t=torch.cat([t,n_steps-1-t],dim=0)
    t=t.unsqueeze(-1)
    #x0系数
    a=alphas_bar_sqrt[t]
    #eps的系数
    aml=one_minus_alphas_bar_sqrt[t]
    #生成随机噪音eps
    e=torch.randn_like(x_0)
    
    #构造模型的输入
    x=x_0*a+e*aml
    
    #送入模型，得到t时刻的随机噪声预测值
    output=model(x,t.squeeze(-1))
    
    #与真实噪声一起计算误差，求平均值
    return (e-output).square().mean()

def p_sample_loop(model,shape,n_steps,betas,one_minus_alphas_bar_sqrt):
    cur_x=torch.randn(shape)
    x_seq=[cur_x]
    for i in reversed(range(n_steps)):
        cur_x=p_sample(model,cur_x,i,betas,one_minus_alphas_bar_sqrt)
        x_seq.append(cur_x)
    return x_seq

def p_sample(model,x,t,betas,one_minus_alphas_bar_sqrt):
    #从x[T]采样t时刻的重构值
    t=torch.tensor([t])
    coeff=betas[t]/one_minus_alphas_bar_sqrt[t]
    eps_theta=model(x,t)
    mean=(1/(1-betas[t]).sqrt())*(x-(coeff*eps_theta))
    z=torch.randn_like(x)
    sigma_t=betas[t].sqrt()
    sample=mean+sigma_t*z
    return sample

seed=1234
class EMA:
    def __init__(self,mu=0.0):
        self.mu=mu
        self.shadow={}
    def register(self,name,val):
        self.shadow[name]=val.clone()
    
    def __call__(self,name,x):
        assert name in self.shadow
        new_average=self.mu*x+(1.0-self.mu)*self.shadow[name]
        self.shadow[name]=new_average.clone()
        return new_average

print("Training model")

batch_size=128
dataloader=torch.utils.data.DataLoader(dataset, batch_size,shuffle=True)
num_epoch=4000
plt.rc("text",color="blue")
model=MLPDiffusion(num_steps)
optimzer=torch.optim.Adam(model.parameters(),lr=1e-3)

for t in range(num_epoch):
    for idx,batch_x in enumerate(dataloader):
        loss=diffusion_loss_fn(model,batch_x,alphas_bar_sqrt,one_minus_alphas_bar_sqrt,num_steps)
        optimzer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(),1.0)
        optimzer.step()
    if(t%100==0):
        print(loss)
        x_seq=p_sample_loop(model,dataset.shape,num_steps,betas,one_minus_alphas_bar_sqrt)
        fig,axs=plt.subplots(1,10,figsize=(28,3))
        for i in range(1,11):
            cur_x=x_seq[i*10].detach()
            axs[i-1].scatter(cur_x[:,0],cur_x[:,1],color="red",edgecolors="white")
            axs[i-1].axis("off")
            axs[i-1].set_title(str(i))

import io
from PIL import Image
imgs=[]
for i in range(100):
    plt.clf()
    q_i=q_x(dataset,torch.tensor([i]))
    plt.scatter(q_i[:,0],q_i[:,1],color="red",edgecolors="white",s=5)
    plt.axis("off")
    
    img_buf=io.BytesIO()
    plt.savefig(img_buf,format="png")
    img=Image.open(img_buf)
    imgs.append(img)

reverse=[]
for i in range(100):
    plt.clf()
    cur_x=x_seq[i].detach()
    plt.scatter(cur_x[:,0],cur_x[:,1],color="red",edgecolors="white",s=5)
    plt.axis("off")
    
    img_buf=io.BytesIO()
    plt.savefig(img_buf,format="png")
    img=Image.open(img_buf)
    imgs.append(img)

Code2:图片Demo

import torch
import torchvision
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torchvision import transforms
from torch.utils.data import DataLoader
import numpy as np
from torch.optim import Adam
from torch import nn
import math
from torchvision.utils import save_image


def show_images(data, num_samples=20, cols=4):
    """ Plots some samples from the dataset """
    plt.figure(figsize=(15,15))
    for i, img in enumerate(data):
        if i == num_samples:
            break
        plt.subplot(int(num_samples/cols) + 1, cols, i + 1)
        plt.imshow(img[0])


def linear_beta_schedule(timesteps, start=0.0001, end=0.02):
    return torch.linspace(start, end, timesteps)


def get_index_from_list(vals, t, x_shape):
    """
    Returns a specific index t of a passed list of values vals
    while considering the batch dimension.
    """
    batch_size = t.shape[0]
    out = vals.gather(-1, t.cpu())
    #print("out:",out)
    #print("out.shape:",out.shape)
    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)

def forward_diffusion_sample(x_0, t, device="cpu"):
    """
    Takes an image and a timestep as input and
    returns the noisy version of it
    """
    noise = torch.randn_like(x_0)
    sqrt_alphas_cumprod_t = get_index_from_list(sqrt_alphas_cumprod, t, x_0.shape)
    sqrt_one_minus_alphas_cumprod_t = get_index_from_list(
        sqrt_one_minus_alphas_cumprod, t, x_0.shape
    )
    # mean + variance
    return sqrt_alphas_cumprod_t.to(device) * x_0.to(device) \
    + sqrt_one_minus_alphas_cumprod_t.to(device) * noise.to(device), noise.to(device)






def load_transformed_dataset(IMG_SIZE):
    data_transforms = [
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.ToTensor(), # Scales data into [0,1]
        transforms.Lambda(lambda t: (t * 2) - 1) # Scale between [-1, 1]
    ]
    data_transform = transforms.Compose(data_transforms)

    train = torchvision.datasets.MNIST(root="./Data",transform=data_transform,train=True)
    test = torchvision.datasets.MNIST(root="./Data", transform=data_transform, train=False)

    return torch.utils.data.ConcatDataset([train, test])

def show_tensor_image(image):
    reverse_transforms = transforms.Compose([
        transforms.Lambda(lambda t: (t + 1) / 2),
        transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC
        transforms.Lambda(lambda t: t * 255.),
        transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),
        transforms.ToPILImage(),
    ])

    #Take first image of batch
    if len(image.shape) == 4:
        image = image[0, :, :, :]
    plt.imshow(reverse_transforms(image))





class Block(nn.Module):
    def __init__(self, in_ch, out_ch, time_emb_dim, up=False):
        super().__init__()
        self.time_mlp =  nn.Linear(time_emb_dim, out_ch)
        if up:
            self.conv1 = nn.Conv2d(2*in_ch, out_ch, 3, padding=1)
            self.transform = nn.ConvTranspose2d(out_ch, out_ch, 4, 2, 1)
        else:
            self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
            self.transform = nn.Conv2d(out_ch, out_ch, 4, 2, 1)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.bnorm1 = nn.BatchNorm2d(out_ch)
        self.bnorm2 = nn.BatchNorm2d(out_ch)
        self.relu  = nn.ReLU()

    def forward(self, x, t):
        #print("ttt:",t.shape)
        # First Conv
        h = self.bnorm1(self.relu(self.conv1(x)))
        # Time embedding
        time_emb = self.relu(self.time_mlp(t))
        # Extend last 2 dimensions
        time_emb = time_emb[(..., ) + (None, ) * 2]
        # Add time channel
        h = h + time_emb
        # Second Conv
        h = self.bnorm2(self.relu(self.conv2(h)))
        # Down or Upsample
        return self.transform(h)


class SinusoidalPositionEmbeddings(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, time):
        device = time.device
        half_dim = self.dim // 2
        embeddings = math.log(10000) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        embeddings = time[:, None] * embeddings[None, :]
        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
        # TODO: Double check the ordering here
        return embeddings


class SimpleUnet(nn.Module):
    """
    A simplified variant of the Unet architecture.
    """
    def __init__(self):
        super().__init__()
        image_channels =1   #灰度图为1，彩色图为3
        down_channels = (64, 128, 256, 512, 1024)
        up_channels = (1024, 512, 256, 128, 64)
        out_dim = 1   #灰度图为1 ，彩色图为3
        time_emb_dim = 32

        # Time embedding
        self.time_mlp = nn.Sequential(
                SinusoidalPositionEmbeddings(time_emb_dim),
                nn.Linear(time_emb_dim, time_emb_dim),
                nn.ReLU()
            )

        # Initial projection
        self.conv0 = nn.Conv2d(image_channels, down_channels[0], 3, padding=1)

        # Downsample
        self.downs = nn.ModuleList([Block(down_channels[i], down_channels[i+1], \
                                    time_emb_dim) \
                    for i in range(len(down_channels)-1)])
        # Upsample
        self.ups = nn.ModuleList([Block(up_channels[i], up_channels[i+1], \
                                        time_emb_dim, up=True) \
                    for i in range(len(up_channels)-1)])

        # Edit: Corrected a bug found by Jakub C (see YouTube comment)
        self.output = nn.Conv2d(up_channels[-1], out_dim, 1)

    def forward(self, x, timestep):
        # Embedd time
        t = self.time_mlp(timestep)
        # Initial conv
        x = self.conv0(x)
        # Unet
        residual_inputs = []
        for down in self.downs:
            x = down(x, t)
            residual_inputs.append(x)
        for up in self.ups:
            residual_x = residual_inputs.pop()
            # Add residual x as additional channels
            x = torch.cat((x, residual_x), dim=1)
            x = up(x, t)
        return self.output(x)




def get_loss(model, x_0, t):
    x_noisy, noise = forward_diffusion_sample(x_0, t, device)
    noise_pred = model(x_noisy, t)
    return F.l1_loss(noise, noise_pred)



@torch.no_grad()
def sample_timestep(x, t):
    """
    Calls the model to predict the noise in the image and returns
    the denoised image.
    Applies noise to this image, if we are not in the last step yet.
    """
    betas_t = get_index_from_list(betas, t, x.shape)
    sqrt_one_minus_alphas_cumprod_t = get_index_from_list(
        sqrt_one_minus_alphas_cumprod, t, x.shape
    )

    sqrt_recip_alphas_t = get_index_from_list(sqrt_recip_alphas, t, x.shape)

    # Call model (current image - noise prediction)
    model_mean = sqrt_recip_alphas_t * (
        x - betas_t * model(x, t) / sqrt_one_minus_alphas_cumprod_t
    )

    posterior_variance_t = get_index_from_list(posterior_variance, t, x.shape)

    if t == 0:
        # As pointed out by Luis Pereira (see YouTube comment)
        # The t's are offset from the t's in the paper
        return model_mean
    else:
        noise = torch.randn_like(x)
        return model_mean + torch.sqrt(posterior_variance_t) * noise

@torch.no_grad()
def sample_plot_image(IMG_SIZE):
    # Sample noise
    img_size = IMG_SIZE
    img = torch.randn((1, 1, img_size, img_size), device=device)   #生成第T步的图片
    plt.figure(figsize=(15,15))
    plt.axis('off')
    num_images = 10
    stepsize = int(T/num_images)

    for i in range(0,T)[::-1]:
        t = torch.full((1,), i, device=device, dtype=torch.long)
        #print("t:",t)
        img = sample_timestep(img, t)
        # Edit: This is to maintain the natural range of the distribution
        img = torch.clamp(img, -1.0, 1.0)
        if i % stepsize == 0:
            plt.subplot(1, num_images, int(i/stepsize)+1)
            plt.title(str(i))
            show_tensor_image(img.detach().cpu())
    plt.show()


if __name__ =="__main__":

    # Define beta schedule
    T = 300
    betas = linear_beta_schedule(timesteps=T)

    # Pre-calculate different terms for closed form
    alphas = 1. - betas
    alphas_cumprod = torch.cumprod(alphas, axis=0)
    # print(alphas_cumprod.shape)
    alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
    # print(alphas_cumprod_prev)
    # print(alphas_cumprod_prev.shape)
    sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
    sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
    sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)
    posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
    # print(posterior_variance.shape)


    IMG_SIZE = 32
    BATCH_SIZE = 16

    data = load_transformed_dataset(IMG_SIZE)
    dataloader = DataLoader(data, batch_size=BATCH_SIZE, shuffle=True, drop_last=True)

    model = SimpleUnet()
    print("Num params: ", sum(p.numel() for p in model.parameters()))

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    optimizer = Adam(model.parameters(), lr=0.001)
    epochs = 1 # Try more!

    for epoch in range(epochs):
        for step, batch in enumerate(dataloader):  #由于batch 是包含标签的所以取batch[0]
            #print(batch[0].shape)
            optimizer.zero_grad()

            t = torch.randint(0, T, (BATCH_SIZE,), device=device).long()
            loss = get_loss(model, batch[0], t)
            loss.backward()
            optimizer.step()

            if epoch % 1 == 0 and step %5== 0:
                print(f"Epoch {epoch} | step {step:03d} Loss: {loss.item()} ")
                sample_plot_image(IMG_SIZE)

参考文献
https://zhuanlan.zhihu.com/p/630354327
https://blog.csdn.net/weixin_45453121/article/details/131223653
https://www.cnblogs.com/risejl/p/17448442.html
https://zhuanlan.zhihu.com/p/569994589?utm_id=0