!pip install -q -U einops datasets matplotlib tqdm


import math
from inspect import isfunction
from functools import partial                                 


%matplotlib inline
import matplotlib.pyplot as plt
from tqdm.auto import tqdm
from einops import rearrange, reduce
from einops.layers.torch import Rearrange


import torch
from torch import nn, einsum
import torch.nn.functional as F

import torch


# cosine_beta_schedule函數(shù)用于創(chuàng)建一個余弦退火beta調(diào)度。
# 這種調(diào)度方法基于余弦函數(shù)，并且可以調(diào)整隨時間的衰減速率。
def cosine_beta_schedule(timesteps, s=0.008):
    steps = timesteps + 1  # 計算總的步數(shù)，需要比時間步多一個，以便計算alpha的累積乘積
    x = torch.linspace(0, timesteps, steps)  # 創(chuàng)建從0到timesteps的均勻分布的張量
    # 計算alpha的累積乘積，使用一個余弦變換，并平方來計算當(dāng)前步的alpha值
    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]  # 歸一化，確保初始值為1
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])  # 計算每個時間步的beta值
    return torch.clip(betas, 0.0001, 0.9999)  # 對beta值進(jìn)行裁剪，避免過大或過小


# linear_beta_schedule函數(shù)用于創(chuàng)建一個線性退火beta調(diào)度。
# 這意味著beta值將從beta_start線性增加到beta_end。
def linear_beta_schedule(timesteps):
    beta_start = 0.0001  # 定義起始beta值
    beta_end = 0.02  # 定義結(jié)束beta值
    return torch.linspace(beta_start, beta_end, timesteps)  # 創(chuàng)建一個線性分布的beta值數(shù)組


# quadratic_beta_schedule函數(shù)用于創(chuàng)建一個二次退火beta調(diào)度。
# 這意味著beta值將根據(jù)二次函數(shù)變化。
def quadratic_beta_schedule(timesteps):
    beta_start = 0.0001  # 定義起始beta值
    beta_end = 0.02  # 定義結(jié)束beta值
    # 創(chuàng)建一個線性分布的數(shù)組，然后將其平方以生成二次分布，最后再次平方以計算beta值
    return torch.linspace(beta_start**0.5, beta_end**0.5, timesteps) ** 2


# sigmoid_beta_schedule函數(shù)用于創(chuàng)建一個sigmoid退火beta調(diào)度。
# 這意味著beta值將根據(jù)sigmoid函數(shù)變化，這是一種常見的激活函數(shù)。
def sigmoid_beta_schedule(timesteps):
    beta_start = 0.0001  # 定義起始beta值
    beta_end = 0.02  # 定義結(jié)束beta值
    betas = torch.linspace(-6, 6, timesteps)  # 創(chuàng)建一個從-6到6的線性分布，用于sigmoid函數(shù)的輸入
    # 應(yīng)用sigmoid函數(shù)，并根據(jù)beta_start和beta_end調(diào)整其范圍和位置
    return torch.sigmoid(betas) * (beta_end - beta_start) + beta_start

# import torch  # 假設(shè)在代碼的其他部分已經(jīng)導(dǎo)入了torch庫


# 定義前向擴(kuò)散函數(shù)
# x_start: 初始數(shù)據(jù)，例如一批圖像
# t: 擴(kuò)散的時間步，表示當(dāng)前的擴(kuò)散階段
# noise: 可選參數(shù)，如果提供，則使用該噪聲數(shù)據(jù)；否則，將生成新的隨機(jī)噪聲
def q_sample(x_start, t, noise=None):
    if noise is None:
        noise = torch.randn_like(x_start)  # 如果未提供噪聲，則生成一個與x_start形狀相同的隨機(jī)噪聲張量
        
 # 提取對應(yīng)于時間步t的α的累積乘積的平方根
    sqrt_alphas_cumprod_t = extract(sqrt_alphas_cumprod, t, x_start.shape)
    # 提取對應(yīng)于時間步t的1-α的累積乘積的平方根
    sqrt_one_minus_alphas_cumprod_t = extract(
        sqrt_one_minus_alphas_cumprod, t, x_start.shape
    )
  # 返回前向擴(kuò)散的結(jié)果，該結(jié)果是初始數(shù)據(jù)和噪聲的線性組合
    # 系數(shù)sqrt_alphas_cumprod_t和sqrt_one_minus_alphas_cumprod_t分別用于縮放初始數(shù)據(jù)和噪聲
    return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise      

# take time step
for noise in [10,20,40,80 100]:
  t = torch.tensor([40])
  get_noisy_image(x_start, t)

import torch.nn as nn
from einops import rearrange  # 假設(shè)已經(jīng)導(dǎo)入了einops庫中的rearrange函數(shù)
from torch_utils import exists  # 假設(shè)已經(jīng)定義了exists函數(shù)，用于檢查對象是否存在


# 定義一個基礎(chǔ)的Block類，該類將作為神經(jīng)網(wǎng)絡(luò)中的一個基本構(gòu)建模塊
class Block(nn.Module):
    def __init__(self, dim, dim_out, groups=8):
        super().__init__()
        # 一個2D卷積層，卷積核大小為3x3，邊緣填充為1，從輸入維度dim到輸出維度dim_out
        self.proj = nn.Conv2d(dim, dim_out, 3, padding=1)
        # GroupNorm層用于歸一化，分組數(shù)為groups
        self.norm = nn.GroupNorm(groups, dim_out)
        # 使用SiLU（也稱為Swish）作為激活函數(shù)
        self.act = nn.SiLU()
        
    def forward(self, x, scale_shift=None):
        x = self.proj(x)  # 應(yīng)用卷積操作
        x = self.norm(x)  # 應(yīng)用歸一化操作
        # 如果scale_shift參數(shù)存在，則對歸一化后的數(shù)據(jù)進(jìn)行縮放和位移操作
        if exists(scale_shift):
            scale, shift = scale_shift
            x = x * (scale + 1) + shift
        x = self.act(x)  # 應(yīng)用激活函數(shù)
        return x  # 返回處理后的數(shù)據(jù)


# 定義一個ResnetBlock類，用于構(gòu)建殘差網(wǎng)絡(luò)中的基本塊
class ResnetBlock(nn.Module):
    """https://arxiv.org/abs/1512.03385"""
    
    def __init__(self, dim, dim_out, *, time_emb_dim=None, groups=8):
        super().__init__()
        # 如果time_emb_dim存在，定義一個小型的多層感知器（MLP）網(wǎng)絡(luò)
        self.mlp = (
            nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, dim_out))
            if exists(time_emb_dim)
            else None
        )
        # 定義兩個順序的基礎(chǔ)Block模塊
        self.block1 = Block(dim, dim_out, groups=groups)
        self.block2 = Block(dim_out, dim_out, groups=groups)
        # 如果輸入維度dim和輸出維度dim_out不同，則使用1x1卷積進(jìn)行維度調(diào)整
        # 否則使用Identity層（相當(dāng)于不做任何處理）
        self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
        
    def forward(self, x, time_emb=None):
        h = self.block1(x)  # 通過第一個Block模塊
        # 如果存在時間嵌入向量time_emb且存在mlp模塊，則將其應(yīng)用到h上
        if exists(self.mlp) and exists(time_emb):
            time_emb = self.mlp(time_emb)  # 通過MLP網(wǎng)絡(luò)
            # 重整time_emb的形狀以匹配h的形狀，并將結(jié)果加到h上
            h = rearrange(time_emb, "b c -> b c 1 1") + h
        h = self.block2(h)  # 通過第二個Block模塊
        return h + self.res_conv(x)  # 將Block模塊的輸出與調(diào)整維度后的原始輸入x相加并返回

import torch
from torch import nn
from einops import rearrange
import torch.nn.functional as F


# 定義一個標(biāo)準(zhǔn)的多頭注意力（Multi-Head Attention）機(jī)制的類
class Attention(nn.Module):
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        # 根據(jù)維度的倒數(shù)平方根來縮放查詢（Query）向量
        self.scale = dim_head ** -0.5
        # 頭的數(shù)量（多頭中的"多"）
        self.heads = heads
        # 計算用于多頭注意力的隱藏層維度
        hidden_dim = dim_head * heads
        # 定義一個卷積層將輸入的特征映射到QKV（查詢、鍵、值）空間
        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
        # 定義一個卷積層將多頭注意力的輸出映射回原特征空間
        self.to_out = nn.Conv2d(hidden_dim, dim, 1)


    def forward(self, x):
        # 獲取輸入的批量大小、通道數(shù)、高度和寬度
        b, c, h, w = x.shape
        # 使用to_qkv卷積層得到QKV，并將其分離為三個組件
        qkv = self.to_qkv(x).chunk(3, dim=1)
        # 將QKV重排并縮放查詢向量
        q, k, v = map(
            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
        )
        q = q * self.scale
        # 使用愛因斯坦求和約定計算查詢和鍵之間的相似度得分
        sim = einsum("b h d i, b h d j -> b h i j", q, k)
        # 從相似度得分中減去最大值以提高數(shù)值穩(wěn)定性
        sim = sim - sim.amax(dim=-1, keepdim=True).detach()
        # 應(yīng)用Softmax函數(shù)獲取注意力權(quán)重
        attn = sim.softmax(dim=-1)
        # 使用注意力權(quán)重對值進(jìn)行加權(quán)
        out = einsum("b h i j, b h d j -> b h i d", attn, v)
        # 將輸出重新排列回原始的空間形狀
        out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w)
        # 返回通過輸出卷積層的結(jié)果
        return self.to_out(out)


# 定義一個線性注意力（Linear Attention）機(jī)制的類
class LinearAttention(nn.Module):
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        # 根據(jù)維度的倒數(shù)平方根來縮放查詢（Query）向量
        self.scale = dim_head ** -0.5
        # 頭的數(shù)量
        self.heads = heads
        # 計算用于多頭注意力的隱藏層維度
        hidden_dim = dim_head * heads
        # 定義一個卷積層將輸入的特征映射到QKV空間
        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
        # 定義一個順序容器包含卷積層和組歸一化層將輸出映射回原特征空間
        self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1),
                                    nn.GroupNorm(1, dim))


    def forward(self, x):
        # 獲取輸入的批量大小、通道數(shù)、高度和寬度
        b, c, h, w = x.shape
        # 使用to_qkv卷積層得到QKV，并將其分離為三個組件
        qkv = self.to_qkv(x).chunk(3, dim=1)
        # 將QKV重排，應(yīng)用Softmax函數(shù)并縮放查詢向量
        q, k, v = map(
            lambda t: 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
        # 計算上下文矩陣，是鍵和值的加權(quán)組合
        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)
        # 返回經(jīng)過輸出順序容器處理的結(jié)果
        return self.to_out(out)

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)
        return embeddings

class Unet(nn.Module):
    # 初始化函數(shù)，定義U-Net網(wǎng)絡(luò)的結(jié)構(gòu)和參數(shù)
    def __init__(
            self,
            dim,  # 基本隱藏層維度
            init_dim=None,  # 初始層維度，如果未提供則會根據(jù)dim計算得出
            out_dim=None,  # 輸出維度，如果未提供則默認(rèn)為輸入圖像的通道數(shù)
            dim_mults=(1, 2, 4, 8),  # 控制每個階段隱藏層維度倍增的倍數(shù)
            channels=3,  # 輸入圖像的通道數(shù)，默認(rèn)為3
            with_time_emb=True,  # 是否使用時間嵌入，這對于某些生成模型可能是必要的
            resnet_block_groups=8,  # ResNet塊中的組數(shù)
            use_cnotallow=True,  # 是否使用ConvNeXt塊而不是ResNet塊
            convnext_mult=2,  # ConvNeXt塊的維度倍增因子
    ):
        super().__init__()  # 調(diào)用父類構(gòu)造函數(shù)
        # 確定各層維度
        self.channels = channels
        init_dim = default(init_dim, dim // 3 * 2)  # 設(shè)置或計算初始層維度
        self.init_conv = nn.Conv2d(channels, init_dim, 7, padding=3)  # 初始卷積層，使用7x7卷積核和padding
        dims = [init_dim, *map(lambda m: dim * m, dim_mults)]  # 計算每個階段的維度
        in_out = list(zip(dims[:-1], dims[1:]))  # 創(chuàng)建輸入輸出維度對
        # 根據(jù)use_convnext選擇塊類
        if use_convnext:
            block_klass = partial(ConvNextBlock, mult=convnext_mult)
        else:
            block_klass = partial(ResnetBlock, groups=resnet_block_groups)
        # 時間嵌入層
        if with_time_emb:
            time_dim = dim * 4  # 時間嵌入的維度
            self.time_mlp = nn.Sequential(  # 時間嵌入的多層感知機(jī)
                SinusoidalPositionEmbeddings(dim),  # 正弦位置嵌入
                nn.Linear(dim, time_dim),  # 線性變換
                nn.GELU(),  # GELU激活函數(shù)
                nn.Linear(time_dim, time_dim),  # 再一次線性變換
            )
        else:
            time_dim = None
            self.time_mlp = None
        # 下采樣層
        self.downs = nn.ModuleList([])
        self.ups = nn.ModuleList([])
        num_resolutions = len(in_out)  # 解析的層數(shù)
        # 構(gòu)建下采樣模塊
        for ind, (dim_in, dim_out) in enumerate(in_out):
            is_last = ind >= (num_resolutions - 1)  # 是否為最后一層
            self.downs.append(  # 添加下采樣塊
                nn.ModuleList(
                    [
                        block_klass(dim_in, dim_out, time_emb_dim=time_dim),  # 卷積塊
                        block_klass(dim_out, dim_out, time_emb_dim=time_dim),  # 卷積塊
                        Residual(PreNorm(dim_out, LinearAttention(dim_out))),  # 殘差連接和注意力模塊
                        Downsample(dim_out) if not is_last else nn.Identity(),  # 下采樣或恒等映射
                    ]
                )
            )
        # 中間層（瓶頸層）
        mid_dim = dims[-1]
        # 中間層（瓶頸層）
        # 第一個中間卷積塊
        self.mid_block1 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)
        # 中間層的注意力模塊
        self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim)))
        # 第二個中間卷積塊
        self.mid_block2 = block_klass(mid_dim, mid_dim, time_emb_dim=time_dim)


        # 構(gòu)建上采樣模塊
        for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
            is_last = ind >= (num_resolutions - 1)  # 是否是最后一次上采樣，減2是因為我們需要留出一個輸出層
            self.ups.append(
                nn.ModuleList(
                    [
                        # 卷積塊，這里輸入維度翻倍是因為上采樣過程中會與編碼器階段的相應(yīng)層進(jìn)行拼接
                        block_klass(dim_out * 2, dim_in, time_emb_dim=time_dim),
                        # 卷積塊
                        block_klass(dim_in, dim_in, time_emb_dim=time_dim),
                        # 殘差和注意力模塊
                        Residual(PreNorm(dim_in, LinearAttention(dim_in))),
                        # 上采樣或恒等映射
                        Upsample(dim_in) if not is_last else nn.Identity(),
                    ]
                )
            )


        # 設(shè)置或計算輸出維度，如果未提供則默認(rèn)為輸入圖像的通道數(shù)
        out_dim = default(out_dim, channels)
        # 最后的卷積層，將輸出維度變換到期望的輸出維度
        self.final_conv = nn.Sequential(
            block_klass(dim, dim),  # 卷積塊
            nn.Conv2d(dim, out_dim, 1)  # 1x1卷積，用于輸出維度變換
        )


    # 前向傳播函數(shù)
    def forward(self, x, time):
        # 初始卷積層
        x = self.init_conv(x)
        # 如果存在時間嵌入層，則將時間編碼
        t = self.time_mlp(time) if exists(self.time_mlp) else None
        # 用于存儲各個階段的特征圖
        h = []


        # 下采樣過程
        for block1, block2, attn, downsample in self.downs:
            x = block1(x, t)  # 應(yīng)用卷積塊
            x = block2(x, t)  # 應(yīng)用卷積塊
            x = attn(x)  # 應(yīng)用注意力模塊
            h.append(x)  # 存儲特征圖以便后續(xù)的拼接
            x = downsample(x)  # 應(yīng)用下采樣或恒等映射


        # 中間層或瓶頸層
        x = self.mid_block1(x, t)  # 第一個中間卷積塊
        x = self.mid_attn(x)  # 中間層的注意力模塊
        x = self.mid_block2(x, t)  # 第二個中間卷積塊


        # 上采樣過程
        for block1, block2, attn, upsample in self.ups:
            # 拼接特征圖和對應(yīng)的編碼器階段的特征圖
            x = torch.cat((x, h.pop()), dim=1)
            x = block1(x, t)  # 應(yīng)用卷積塊
            x = block2(x, t)  # 應(yīng)用卷積塊
            x = attn(x)  # 應(yīng)用注意力模塊
            x = upsample(x)  # 應(yīng)用上采樣或恒等映射


        # 最后的輸出層，輸出最終的特征圖或圖像
        return self.final_conv(x)

import torch
import torch.nn.functional as F


# 定義損失函數(shù)，它評估去噪模型的性能
def p_losses(denoise_model, x_start, t, noise=None, loss_type="l1"):
    if noise is None:
        noise = torch.randn_like(x_start)  # 如果未提供噪聲，則生成一個與x_start形狀相同的隨機(jī)噪聲張量


    # 使用q_sample函數(shù)生成帶有噪聲的數(shù)據(jù)x_noisy，這模擬了擴(kuò)散模型的前向過程
    x_noisy = q_sample(x_start=x_start, t=t, noise=noise)
    # 使用去噪模型對噪聲數(shù)據(jù)x_noisy進(jìn)行預(yù)測，試圖恢復(fù)加入的噪聲
    predicted_noise = denoise_model(x_noisy, t)


    # 根據(jù)指定的損失類型計算損失
    if loss_type == 'l1':  # 如果損失類型為L1損失
        loss = F.l1_loss(noise, predicted_noise)  # 使用L1損失函數(shù)計算真實噪聲和預(yù)測噪聲之間的差異
    elif loss_type == 'l2':  # 如果損失類型為L2損失（均方誤差損失）
        loss = F.mse_loss(noise, predicted_noise)  # 使用均方誤差損失函數(shù)計算真實噪聲和預(yù)測噪聲之間的差異
    elif loss_type == "huber":  # 如果損失類型為Huber損失
        loss = F.smooth_l1_loss(noise, predicted_noise)  # 使用Huber損失函數(shù)，這是L1和L2損失的結(jié)合，對異常值不那么敏感
    else:
        raise NotImplementedError()  # 如果指定了未實現(xiàn)的損失類型，則拋出異常


    return loss  # 返回計算得到的損失值

if __name__=="__main__":
    for epoch in range(epochs):
        for step, batch in tqdm(enumerate(dataloader), desc='Training'):
          optimizer.zero_grad()
          batch = batch[0]


          batch_size = batch.shape[0]
          batch = batch.to(device)
          # 國內(nèi)版啟用這段，注釋上面兩行
          # batch_size = batch[0].shape[0]
          # batch = batch[0].to(device)


          # Algorithm 1 line 3: sample t uniformally for every example in the batch
          t = torch.randint(0, timesteps, (batch_size,), device=device).long()


          loss = p_losses(model, batch, t, loss_type="huber")


          if step % 50 == 0:
            print("Loss:", loss.item())


          loss.backward()
          optimizer.step()


          # save generated images
          if step != 0 and step % save_and_sample_every == 0:
            milestone = step // save_and_sample_every
            batches = num_to_groups(4, batch_size)
            all_images_list = list(map(lambda n: sample(model, batch_size=n, channels=channels), batches))
            all_images = torch.cat(all_images_list, dim=0)
            all_images = (all_images + 1) * 0.5
            # save_image(all_images, str(results_folder / f'sample-{milestone}.png'), nrow = 6)
            currentDateAndTime = datetime.now()
            torch.save(model,f"train.pt")