模型代码

models/__init__.py

@@ -0,0 +1,2 @@
+from models.ddm import *
+from models.restoration import *

models/ddm.py

@@ -0,0 +1,260 @@
+import os
+import time
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.backends.cudnn as cudnn
+import utils
+from models.unet import DiffusionUNet
+import torch.distributed as dist
+from torch.utils.tensorboard import SummaryWriter
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+
+def data_transform(X):
+    return 2 * X - 1.0
+
+
+def inverse_data_transform(X):
+    return torch.clamp((X + 1.0) / 2.0, 0.0, 1.0)
+
+
+class EMAHelper(object):
+    def __init__(self, mu=0.9999):
+        self.mu = mu
+        self.shadow = {}
+
+    def register(self, module):
+        if isinstance(module, nn.DataParallel) or isinstance(module, nn.parallel.DistributedDataParallel):
+            module = module.module
+        for name, param in module.named_parameters():
+            if param.requires_grad:
+                self.shadow[name] = param.data.clone()
+
+    def update(self, module, device):
+        if isinstance(module, nn.DataParallel) or isinstance(module, nn.parallel.DistributedDataParallel):
+            module = module.module
+        for name, param in module.named_parameters():
+            if param.requires_grad:
+                self.shadow[name].data = (1. - self.mu) * param.data + self.mu * self.shadow[name].data.to(device)
+
+    def ema(self, module):
+        if isinstance(module, nn.DataParallel) or isinstance(module, nn.parallel.DistributedDataParallel):
+            module = module.module
+        for name, param in module.named_parameters():
+            if param.requires_grad:
+                param.data.copy_(self.shadow[name].data)
+
+    def ema_copy(self, module):
+        if isinstance(module, nn.DataParallel) or isinstance(module, nn.parallel.DistributedDataParallel):
+            inner_module = module.module
+            module_copy = type(inner_module)(inner_module.config).to(inner_module.config.device)
+            module_copy.load_state_dict(inner_module.state_dict())
+            module_copy = nn.DataParallel(module_copy)
+        else:
+            module_copy = type(module)(module.config).to(module.config.device)
+            module_copy.load_state_dict(module.state_dict())
+        self.ema(module_copy)
+        return module_copy
+
+    def state_dict(self):
+        return self.shadow
+
+    def load_state_dict(self, state_dict):
+        self.shadow = state_dict
+
+
+def get_beta_schedule(beta_schedule, *, beta_start, beta_end, num_diffusion_timesteps):
+    def sigmoid(x):
+        return 1 / (np.exp(-x) + 1)
+
+    if beta_schedule == "quad":
+        betas = (np.linspace(beta_start ** 0.5, beta_end ** 0.5, num_diffusion_timesteps, dtype=np.float64) ** 2)
+    elif beta_schedule == "linear":
+        betas = np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == "const":
+        betas = beta_end * np.ones(num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == "jsd":  # 1/T, 1/(T-1), 1/(T-2), ..., 1
+        betas = 1.0 / np.linspace(num_diffusion_timesteps, 1, num_diffusion_timesteps, dtype=np.float64)
+    elif beta_schedule == "sigmoid":
+        betas = np.linspace(-6, 6, num_diffusion_timesteps)
+        betas = sigmoid(betas) * (beta_end - beta_start) + beta_start
+    else:
+        raise NotImplementedError(beta_schedule)
+    assert betas.shape == (num_diffusion_timesteps,)
+    return betas
+
+
+def noise_estimation_loss(model, x0, t, e, b):
+    a = (1 - b).cumprod(dim=0).index_select(0, t).view(-1, 1, 1, 1)
+    x = x0[:, 3:, :, :] * a.sqrt() + e * (1.0 - a).sqrt()
+    output = model(torch.cat([x0[:, :3, :, :], x], dim=1), t.float())
+    return (e - output).square().sum(dim=(1, 2, 3)).mean(dim=0)
+
+
+class DenoisingDiffusion(object):
+    def __init__(self, config, test=False):
+        super().__init__()
+        self.config = config
+        self.device = config.device
+        self.writer = SummaryWriter(config.data.tensorboard)
+        self.model = DiffusionUNet(config)
+        self.model.to(self.device)
+        if test:
+            self.model = torch.nn.DataParallel(self.model)
+        else:
+            self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[config.local_rank],
+                                                                   output_device=config.local_rank)
+        self.ema_helper = EMAHelper()
+        self.ema_helper.register(self.model)
+
+        self.optimizer = utils.optimize.get_optimizer(self.config, self.model.parameters())
+        self.scheduler = CosineAnnealingLR(self.optimizer, T_max=config.training.n_epochs)
+        self.start_epoch, self.step = 0, 0
+
+        betas = get_beta_schedule(
+            beta_schedule=config.diffusion.beta_schedule,
+            beta_start=config.diffusion.beta_start,
+            beta_end=config.diffusion.beta_end,
+            num_diffusion_timesteps=config.diffusion.num_diffusion_timesteps,
+        )
+
+        betas = self.betas = torch.from_numpy(betas).float().to(self.device)
+        self.num_timesteps = betas.shape[0]
+
+    def load_ddm_ckpt(self, load_path, ema=False):
+        checkpoint = utils.logging.load_checkpoint(load_path, None)
+        self.start_epoch = checkpoint['epoch']
+        self.step = checkpoint['step']
+        self.model.load_state_dict(checkpoint['state_dict'], strict=True)
+        self.optimizer.load_state_dict(checkpoint['optimizer'])
+        self.ema_helper.load_state_dict(checkpoint['ema_helper'])
+        self.scheduler.load_state_dict(checkpoint['scheduler'])
+        if ema:
+            self.ema_helper.ema(self.model)
+        print("=> loaded checkpoint '{}' (epoch {}, step {})".format(load_path, checkpoint['epoch'], self.step))
+
+    def train(self, DATASET):
+        cudnn.benchmark = True
+        train_loader, val_loader = DATASET.get_loaders()
+        pretrained_model_path = self.config.training.resume + '.pth.tar'
+        if os.path.isfile(pretrained_model_path):
+            self.load_ddm_ckpt(pretrained_model_path)
+        dist.barrier()
+        # 训练
+        for epoch in range(self.start_epoch, self.config.training.n_epochs):
+            if (epoch == 0) and dist.get_rank() == 0:
+                utils.logging.save_checkpoint({
+                    'epoch': epoch + 1,
+                    'step': self.step,
+                    'state_dict': self.model.state_dict(),
+                    'optimizer': self.optimizer.state_dict(),
+                    'ema_helper': self.ema_helper.state_dict(),
+                    'config': self.config,
+                    'scheduler': self.scheduler.state_dict()
+                }, filename=self.config.training.resume + '_' + str(epoch))
+                utils.logging.save_checkpoint({
+                    'epoch': epoch + 1,
+                    'step': self.step,
+                    'state_dict': self.model.state_dict(),
+                    'optimizer': self.optimizer.state_dict(),
+                    'ema_helper': self.ema_helper.state_dict(),
+                    'config': self.config,
+                    'scheduler': self.scheduler.state_dict()
+                }, filename=self.config.training.resume)
+            if dist.get_rank() == 0:
+                print('=> current epoch: ', epoch)
+            data_start = time.time()
+            data_time = 0
+            train_loader.sampler.set_epoch(epoch)
+            for i, (x, y) in enumerate(train_loader):
+                x = x.flatten(start_dim=0, end_dim=1) if x.ndim == 5 else x
+                n = x.size(0)
+                data_time += time.time() - data_start
+                self.model.train()
+                self.step += 1
+
+                x = x.to(self.device)
+                x = data_transform(x)
+                e = torch.randn_like(x[:, 3:, :, :])
+                b = self.betas
+
+                # antithetic sampling
+                t = torch.randint(low=0, high=self.num_timesteps, size=(n // 2 + 1,)).to(self.device)
+                t = torch.cat([t, self.num_timesteps - t - 1], dim=0)[:n]
+                loss = noise_estimation_loss(self.model, x, t, e, b)
+                current_lr = self.optimizer.param_groups[0]['lr']
+
+                if self.step % 10 == 0:
+                    print(
+                        'rank: %d, step: %d, loss: %.6f, lr: %.6f, time consumption: %.6f' % (
+                            dist.get_rank(), self.step, loss.item(), current_lr, data_time / (i + 1)))
+
+                # 更新参数
+                self.optimizer.zero_grad()
+                loss.backward()
+                self.optimizer.step()
+                self.ema_helper.update(self.model, self.device)
+                data_start = time.time()
+
+                if self.step % self.config.training.validation_freq == 0:
+                    self.model.eval()
+                    self.sample_validation_patches(val_loader, self.step)
+
+                if (self.step % 100 == 0) and dist.get_rank() == 0:
+                    self.writer.add_scalar('train/loss', loss.item(), self.step)
+                    self.writer.add_scalar('train/lr', current_lr, self.step)
+
+            self.scheduler.step()
+            # 保存模型
+            if (epoch % self.config.training.snapshot_freq == 0) and dist.get_rank() == 0:
+                utils.logging.save_checkpoint({
+                    'epoch': epoch + 1,
+                    'step': self.step,
+                    'state_dict': self.model.state_dict(),
+                    'optimizer': self.optimizer.state_dict(),
+                    'ema_helper': self.ema_helper.state_dict(),
+                    'config': self.config,
+                    'scheduler': self.scheduler.state_dict()
+                }, filename=self.config.training.resume + '_' + str(epoch))
+                utils.logging.save_checkpoint({
+                    'epoch': epoch + 1,
+                    'step': self.step,
+                    'state_dict': self.model.state_dict(),
+                    'optimizer': self.optimizer.state_dict(),
+                    'ema_helper': self.ema_helper.state_dict(),
+                    'config': self.config,
+                    'scheduler': self.scheduler.state_dict()
+                }, filename=self.config.training.resume)
+
+    def sample_image(self, x_cond, x, last=True, patch_locs=None, patch_size=None):
+        skip = self.config.diffusion.num_diffusion_timesteps // self.config.sampling.sampling_timesteps
+        seq = range(0, self.config.diffusion.num_diffusion_timesteps, skip)
+        if patch_locs is not None:
+            xs = utils.sampling.generalized_steps_overlapping(x, x_cond, seq, self.model, self.betas, eta=0.,
+                                                              corners=patch_locs, p_size=patch_size, device=self.device)
+        else:
+            xs = utils.sampling.generalized_steps(x, x_cond, seq, self.model, self.betas, eta=0., device=self.device)
+        if last:
+            xs = xs[0][-1]
+        return xs
+
+    def sample_validation_patches(self, val_loader, step):
+        image_folder = os.path.join(self.config.data.val_save_dir, str(self.config.data.image_size))
+        with torch.no_grad():
+            if dist.get_rank() == 0:
+                print(f"Processing a single batch of validation images at step: {step}")
+            for i, (x, y) in enumerate(val_loader):
+                x = x.flatten(start_dim=0, end_dim=1) if x.ndim == 5 else x
+                break
+            n = x.size(0)
+            x_cond = x[:, :3, :, :].to(self.device)  # 条件图像
+            x_cond = data_transform(x_cond)
+            x = torch.randn(n, 3, self.config.data.image_size, self.config.data.image_size, device=self.device)
+            x = self.sample_image(x_cond, x)
+            x = inverse_data_transform(x)
+            x_cond = inverse_data_transform(x_cond)
+
+            for i in range(n):
+                utils.logging.save_image(x_cond[i], os.path.join(image_folder, str(step), f"{i}_cond.png"))
+                utils.logging.save_image(x[i], os.path.join(image_folder, str(step), f"{i}.png"))

models/restoration.py

@@ -0,0 +1,59 @@
+import torch
+import utils
+import os
+from tqdm import tqdm
+
+
+def data_transform(X):
+    return 2 * X - 1.0
+
+
+def inverse_data_transform(X):
+    return torch.clamp((X + 1.0) / 2.0, 0.0, 1.0)
+
+
+class DiffusiveRestoration:
+    def __init__(self, diffusion, config):
+        super(DiffusiveRestoration, self).__init__()
+        self.config = config
+        self.diffusion = diffusion
+
+        # 判断预训练模型是否存在
+        pretrained_model_path = self.config.training.resume + '.pth.tar'
+        assert os.path.isfile(pretrained_model_path), ('pretrained diffusion model path is wrong!')
+        self.diffusion.load_ddm_ckpt(pretrained_model_path, ema=True)
+        self.diffusion.model.eval()
+        self.diffusion.model.requires_grad_(False)
+
+    def restore(self, val_loader, r=None):
+        image_folder = self.config.data.test_save_dir
+        with torch.no_grad():
+            for i, (x, y) in tqdm(enumerate(val_loader)):
+                print(f"=> starting processing image named {y}")
+                x = x.flatten(start_dim=0, end_dim=1) if x.ndim == 5 else x
+                x_cond = x[:, :3, :, :].to(self.diffusion.device)
+                x_output = self.diffusive_restoration(x_cond, r=r)
+                x_output = inverse_data_transform(x_output)
+                utils.logging.save_image(x_output, os.path.join(image_folder, f"{y[0]}.png"))
+
+    def diffusive_restoration(self, x_cond, r=None):
+        p_size = self.config.data.image_size
+        h_list, w_list = self.overlapping_grid_indices(x_cond, output_size=p_size, r=r)
+        corners = [(i, j) for i in h_list for j in w_list]
+        x = torch.randn(x_cond.size(), device=self.diffusion.device)
+        x_output = self.diffusion.sample_image(x_cond, x, patch_locs=corners, patch_size=p_size)
+        return x_output
+
+    def overlapping_grid_indices(self, x_cond, output_size, r=None):
+        _, c, h, w = x_cond.shape
+        r = 16 if r is None else r
+        h_list = [i for i in range(0, h - output_size + 1, r)]
+        w_list = [i for i in range(0, w - output_size + 1, r)]
+        return h_list, w_list
+
+    def web_restore(self, image, r=None):
+        with torch.no_grad():
+            image_cond = image.to(self.diffusion.device)
+            image_output = self.diffusive_restoration(image_cond, r=r)
+            image_output = inverse_data_transform(image_output)
+            return image_output

models/unet.py

@@ -0,0 +1,331 @@
+import math
+import torch
+import torch.nn as nn
+
+
+def get_timestep_embedding(timesteps, embedding_dim):
+    assert len(timesteps.shape) == 1
+
+    half_dim = embedding_dim // 2
+    emb = math.log(10000) / (half_dim - 1)
+    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
+    emb = emb.to(device=timesteps.device)
+    emb = timesteps.float()[:, None] * emb[None, :]
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+    if embedding_dim % 2 == 1:  # zero pad
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+def nonlinearity(x):
+    # swish
+    return x*torch.sigmoid(x)
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(
+            x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=0)
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0, 1, 0, 1)
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
+                 dropout, temb_channels=512):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(in_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        self.temb_proj = torch.nn.Linear(temb_channels,
+                                         out_channels)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels,
+                                                     out_channels,
+                                                     kernel_size=3,
+                                                     stride=1,
+                                                     padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                    out_channels,
+                                                    kernel_size=1,
+                                                    stride=1,
+                                                    padding=0)
+
+    def forward(self, x, temb):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x+h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # 自注意力
+        b, c, h, w = q.shape
+        q = q.reshape(b, c, h*w)
+        q = q.permute(0, 2, 1).contiguous()   # b,hw,c
+        k = k.reshape(b, c, h*w)  # b,c,hw
+        w_ = torch.bmm(q, k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b, c, h*w)
+        w_ = w_.permute(0, 2, 1).contiguous()    # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v, w_)
+        h_ = h_.reshape(b, c, h, w)
+
+        h_ = self.proj_out(h_)
+
+        return x+h_
+
+
+class DiffusionUNet(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        ch, out_ch, ch_mult = config.model.ch, config.model.out_ch, tuple(config.model.ch_mult)
+        num_res_blocks = config.model.num_res_blocks
+        attn_resolutions = config.model.attn_resolutions
+        dropout = config.model.dropout
+        in_channels = config.model.in_channels * 2 if config.data.conditional else config.model.in_channels
+        resolution = config.data.image_size
+        resamp_with_conv = config.model.resamp_with_conv
+
+        self.ch = ch
+        self.temb_ch = self.ch*4
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        # timestep embedding
+        self.temb = nn.Module()
+        self.temb.dense = nn.ModuleList([
+            torch.nn.Linear(self.ch,
+                            self.temb_ch),
+            torch.nn.Linear(self.temb_ch,
+                            self.temb_ch),
+        ])
+
+        # 下采样
+        self.conv_in = torch.nn.Conv2d(in_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+ch_mult
+        self.down = nn.ModuleList()
+        block_in = None
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # 上采样
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            skip_in = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                if i_block == self.num_res_blocks:
+                    skip_in = ch*in_ch_mult[i_level]
+                block.append(ResnetBlock(in_channels=block_in+skip_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)  # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x, t):
+        assert x.shape[2] == x.shape[3] == self.resolution
+
+        # timestep embedding
+        temb = get_timestep_embedding(t, self.ch)
+        temb = self.temb.dense[0](temb)
+        temb = nonlinearity(temb)
+        temb = self.temb.dense[1](temb)
+
+        # 下采样
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # 上采样
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](
+                    torch.cat([h, hs.pop()], dim=1), temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+# net = DiffusionUNet()