mar.py

from functools import partial
import math
import numpy as np
import scipy.stats as stats
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
from timm.models.vision_transformer import DropPath, Mlp

from pixelloss import PixelLoss
from visualize import visualize_patch


def mask_by_order(mask_len, order, bsz, seq_len):
    """
    根据给定的顺序生成掩码。

    该函数根据指定的顺序和长度，为批量中的每个序列生成一个布尔掩码。
    掩码中前mask_len个位置为True，其余位置为False。

    参数:
        mask_len (int or torch.Tensor): 要掩盖的标记数量。可以是一个标量或一个形状为 (bsz,) 的张量。
        order (torch.Tensor): 顺序张量，形状为 (bsz, seq_len)，表示每个序列中标记的排列顺序。
        bsz (int): 批量大小（batch size）。
        seq_len (int): 序列长度（sequence length）。

    返回:
        torch.Tensor: 生成的掩码张量，形状为 (bsz, seq_len)，数据类型为布尔型（torch.bool）。
    """
    # 初始化一个全零的掩码张量，形状为 (bsz, seq_len)
    masking = torch.zeros(bsz, seq_len).cuda()
    # 将掩码张量展平为一维，并在指定的位置（由order[:, :mask_len.long()]指定）填充1
    # src=torch.ones(bsz, seq_len).cuda() 生成一个全1的张量，用于填充
    # 最终，masking中前mask_len.long()个位置为1，其余为0
    masking = torch.scatter(masking, dim=-1, index=order[:, :mask_len.long()], src=torch.ones(bsz, seq_len).cuda()).bool()
    # 返回生成的布尔掩码
    return masking


class Attention(nn.Module):
    """
    自注意力机制层（Attention Layer）。

    参数:
        dim (int): 输入和输出的维度。
        num_heads (int): 多头注意力机制中的头数，默认为8。
        qkv_bias (bool): 查询、键和值线性变换是否使用偏置，默认为False。
        qk_scale (float, optional): 键查询缩放因子。如果为None，则使用 head_dim ** -0.5。默认为None。
        attn_drop (float): 注意力层的dropout概率，默认为0.0。
        proj_drop (float): 投影层的dropout概率，默认为0.0。
    """
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        # 多头数
        self.num_heads = num_heads
        # 每个注意力头的维度
        head_dim = dim // num_heads
       # 缩放因子，如果未指定，则使用 head_dim ** -0.5
        self.scale = qk_scale or head_dim ** -0.5

        # 查询、键和值线性变换层
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        # 注意力层的dropout
        self.attn_drop = nn.Dropout(attn_drop)
        # 输出投影层
        self.proj = nn.Linear(dim, dim)
        # 输出投影层的dropout
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        """
        前向传播函数。

        参数:
            x (torch.Tensor): 输入张量，形状为 (batch_size, sequence_length, dim)。

        返回:
            torch.Tensor: 自注意力机制的输出，形状为 (batch_size, sequence_length, dim)。
        """
        # 获取批量大小、序列长度和维度
        B, N, C = x.shape
        # 计算查询、键和值张量
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        # 分离查询、键和值
        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)

        with torch.cuda.amp.autocast(enabled=False):
            # 计算缩放点积注意力权重
            attn = (q.float() @ k.float().transpose(-2, -1)) * self.scale

        # 对注意力权重进行缩放和平滑处理
        attn = attn - torch.max(attn, dim=-1, keepdim=True)[0]
        # 应用softmax激活函数
        attn = attn.softmax(dim=-1)
        # 应用dropout
        attn = self.attn_drop(attn)

        # 计算注意力输出
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        # 应用输出投影层
        x = self.proj(x)
        # 应用输出投影层的dropout
        x = self.proj_drop(x)
        return x


class Block(nn.Module):
    """
    Transformer块（Transformer Block），由自注意力机制和前馈神经网络组成。

    参数:
        dim (int): 输入和输出的维度。
        num_heads (int): 多头注意力机制中的头数。
        mlp_ratio (float): MLP（多层感知机）隐藏层维度与输入维度的比率，默认为4.0。
        qkv_bias (bool): 查询、键和值线性变换是否使用偏置，默认为False。
        qk_scale (float, optional): 键查询缩放因子。如果为None，则使用 head_dim ** -0.5。默认为None。
        proj_drop (float): 投影层的dropout概率，默认为0.0。
        attn_drop (float): 注意力层的dropout概率，默认为0.0。
        drop_path (float): DropPath的概率，用于随机深度正则化，默认为0.0。
        act_layer (nn.Module): 激活函数层，默认为nn.GELU（高斯误差线性单元）。
        norm_layer (nn.Module): 归一化层，默认为nn.LayerNorm。
    """
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, proj_drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        # 第一个归一化层
        self.norm1 = norm_layer(dim)
        # 自注意力机制层
        self.attn = Attention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=proj_drop)
        # DropPath用于随机深度正则化，如果drop_path > 0，则使用DropPath，否则使用恒等层
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        # 第二个归一化层
        self.norm2 = norm_layer(dim)
        # MLP隐藏层维度
        mlp_hidden_dim = int(dim * mlp_ratio)
        # MLP层
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=proj_drop)

    def forward(self, x):
        """
        前向传播函数。

        参数:
            x (torch.Tensor): 输入张量，形状为 (batch_size, sequence_length, dim)。

        返回:
            torch.Tensor: Transformer块的输出，形状为 (batch_size, sequence_length, dim)。
        """
        # 自注意力机制 + DropPath + 残差连接
        x = x + self.drop_path(self.attn(self.norm1(x)))
        # MLP + DropPath + 残差连接
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


class MAR(nn.Module):
    """
    多尺度自回归模型（Multi-scale Autoregressive Model, MAR）。

    该模型通过多尺度自回归的方式生成数据，适用于图像生成等任务。
    它结合了自回归生成和条件嵌入，能够在不同尺度上生成数据。

    参数:
        seq_len (int): 序列长度，表示模型处理的最大序列长度。
        patch_size (int): 补丁大小，将输入图像分割成的小块尺寸。
        cond_embed_dim (int): 条件嵌入的维度，用于条件输入的嵌入向量维度。
        embed_dim (int): 嵌入的维度，模型中隐藏层的维度。
        num_blocks (int): Transformer块的数量，模型中堆叠的Transformer层数。
        num_heads (int): 多头注意力机制中的头数，用于自注意力机制的头数。
        attn_dropout (float): 注意力层的dropout概率，默认为0.0。
        proj_dropout (float): 投影层的dropout概率，默认为0.0。
        num_conds (int): 条件数量，默认为1。
        guiding_pixel (bool): 是否使用引导像素，默认为False。
        grad_checkpointing (bool): 是否启用梯度检查点，默认为False。
    """
    def __init__(self, seq_len, patch_size, cond_embed_dim, embed_dim, num_blocks, num_heads, attn_dropout, proj_dropout,
                 num_conds=1, guiding_pixel=False, grad_checkpointing=False
                 ):
        super().__init__()

        # 序列长度
        self.seq_len = seq_len
        # patch大小
        self.patch_size = patch_size

        # 条件数量
        self.num_conds = num_conds
        # 是否使用引导像素
        self.guiding_pixel = guiding_pixel

        # 梯度检查点标志
        self.grad_checkpointing = grad_checkpointing

        # 可变掩码比例
        # 使用截断正态分布生成掩码比例，均值为1.0，标准差为0.25，范围在[-4, 0]之间
        self.mask_ratio_generator = stats.truncnorm(-4, 0, loc=1.0, scale=0.25)

        # 网络结构
        # 掩码标记参数，形状为 (1, 1, embed_dim)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # patch嵌入层
        self.patch_emb = nn.Linear(3 * patch_size ** 2, embed_dim, bias=True)
        # patch嵌入层的LayerNorm
        self.patch_emb_ln = nn.LayerNorm(embed_dim, eps=1e-6)
        # 条件嵌入层
        self.cond_emb = nn.Linear(cond_embed_dim, embed_dim, bias=True)
        if self.guiding_pixel:
            # 如果使用引导像素，则初始化像素投影层
            self.pix_proj = nn.Linear(3, embed_dim, bias=True)
        # 可学习的位置嵌入
        self.pos_embed_learned = nn.Parameter(torch.zeros(1, seq_len+num_conds+self.guiding_pixel, embed_dim))

        # Transformer块列表
        self.blocks = nn.ModuleList([
            Block(embed_dim, num_heads, mlp_ratio=4.,
                  qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6),
                  proj_drop=proj_dropout, attn_drop=attn_dropout)
            for _ in range(num_blocks)
        ])
        # LayerNorm层
        self.norm = nn.LayerNorm(embed_dim, eps=1e-6)

        # 初始化权重
        self.initialize_weights()

        if self.guiding_pixel:
            # 如果使用引导像素，则初始化PixelLoss
            self.guiding_pixel_loss = PixelLoss(
                c_channels=cond_embed_dim,
                width=128,
                depth=3,
                num_heads=4,
                r_weight=5.0
            )

    def initialize_weights(self):
        """
        初始化模型权重。
        """
        # 初始化掩码标记和位置嵌入参数
        torch.nn.init.normal_(self.mask_token, std=.02)
        torch.nn.init.normal_(self.pos_embed_learned, std=.02)

        # 初始化nn.Linear和nn.LayerNorm的权重
        self.apply(self._init_weights)

    def _init_weights(self, m):
        """
        初始化nn.Linear和nn.LayerNorm的权重。

        参数:
            m (nn.Module): 要初始化的模块。
        """
        if isinstance(m, nn.Linear):
            # 使用xavier_uniform初始化线性层的权重
            torch.nn.init.xavier_uniform_(m.weight)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
            if m.weight is not None:
                nn.init.constant_(m.weight, 1.0)

    def patchify(self, x):
        """
        将输入图像分割成patch。

        参数:
            x (torch.Tensor): 输入图像，形状为 (batch_size, channels, height, width)。

        返回:
            torch.Tensor: 分割后的补丁，形状为 (batch_size, patches, patch_dim)。
        """
        # 获取批量大小、通道数、高度和宽度
        bsz, c, h, w = x.shape
        # 获取patch大小
        p = self.patch_size
        # 计算patch数量
        h_, w_ = h // p, w // p

        # 重塑张量以分割成patch
        x = x.reshape(bsz, c, h_, p, w_, p)
        x = torch.einsum('nchpwq->nhwcpq', x)
        x = x.reshape(bsz, h_ * w_, c * p ** 2)
        # 返回分割后的patch
        return x  # [n, l, d]

    def unpatchify(self, x):
        """
        将patch重新组合成图像。

        参数:
            x (torch.Tensor): 补丁张量，形状为 (batch_size, patches, patch_dim)。

        返回:
            torch.Tensor: 重新组合后的图像，形状为 (batch_size, channels, height, width)。
        """
        bsz = x.shape[0]
        p = self.patch_size
        # 计算高度和宽度
        h_, w_ = int(np.sqrt(self.seq_len)), int(np.sqrt(self.seq_len))

        # 重塑张量以重新组合patch
        x = x.reshape(bsz, h_, w_, 3, p, p)
        x = torch.einsum('nhwcpq->nchpwq', x)
        x = x.reshape(bsz, 3, h_ * p, w_ * p)
        # 返回重新组合后的图像
        return x  # [n, 3, h, w]

    def sample_orders(self, bsz):
        """
        生成随机排列的顺序。

        参数:
            bsz (int): 批量大小。

        返回:
            torch.Tensor: 生成的顺序张量，形状为 (batch_size, seq_len)。
        """
        orders = torch.argsort(torch.rand(bsz, self.seq_len).cuda(), dim=1).long()
        return orders

    def random_masking_uniform(self, x, orders):
        """
        均匀随机掩码生成。

        参数:
            x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len, embed_dim)。
            orders (torch.Tensor): 顺序张量，形状为 (batch_size, seq_len)。

        返回:
            torch.Tensor: 生成的掩码张量，形状为 (batch_size, seq_len)。
        """
        bsz, seq_len, embed_dim = x.shape
        # 随机生成要掩盖的标记数量
        num_masked_tokens = np.random.randint(seq_len) + 1
        mask = torch.zeros(bsz, seq_len, device=x.device)
        mask = torch.scatter(mask, dim=-1, index=orders[:, :num_masked_tokens],
                             src=torch.ones(bsz, seq_len, device=x.device))
        return mask

    def random_masking(self, x, orders):
        """
        随机掩码生成，根据掩码比例生成掩码。

        参数:
            x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len, embed_dim)。
            orders (torch.Tensor): 顺序张量，形状为 (batch_size, seq_len)。

        返回:
            torch.Tensor: 生成的掩码张量，形状为 (batch_size, seq_len)。
        """
        bsz, seq_len, embed_dim = x.shape
        # 生成掩码比例
        mask_rates = self.mask_ratio_generator.rvs(bsz)
        # 计算要掩盖的标记数量
        num_masked_tokens = torch.Tensor(np.ceil(seq_len * mask_rates)).cuda()
        # 生成扩展的索引
        expanded_indices = torch.arange(seq_len, device=x.device).expand(bsz, seq_len)
        # 对顺序进行排序
        sorted_orders = torch.argsort(orders, dim=-1)
        # 生成掩码
        mask = (expanded_indices < num_masked_tokens[:, None]).float()
        # 应用掩码
        mask = torch.scatter(torch.zeros_like(mask), dim=-1, index=sorted_orders, src=mask)

        return mask

    def predict(self, x, mask, cond_list):
        """
        前向预测函数，用于生成模型的预测结果。

        参数:
            x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len, patch_dim)。
                              其中，batch_size 是批量大小，seq_len 是序列长度，patch_dim 是每个补丁的维度。
            mask (torch.Tensor): 掩码张量，形状为 (batch_size, seq_len)。
                                  用于指示哪些位置需要被掩盖（masked）。
            cond_list (List[torch.Tensor]): 条件输入列表，包含用于生成的条件信息。

        返回:
            List[torch.Tensor]: 预测结果列表，包含中间条件和四个方向的条件（上、右、下、左）。
        """
        # 将输入patch嵌入到隐藏空间
        x = self.patch_emb(x)

        # 将条件嵌入从之前的生成器中前置到输入张量
        for i in range(self.num_conds):
            x = torch.cat([self.cond_emb(cond_list[i]).unsqueeze(1), x], dim=1)

        # 如果使用引导像素，则将引导像素前置到输入张量
        if self.guiding_pixel:
            x = torch.cat([self.pix_proj(cond_list[-1]).unsqueeze(1), x], dim=1)

        # 应用掩码，将被掩盖的位置替换为掩码标记
        mask_with_cond = torch.cat([torch.zeros(x.size(0), self.num_conds+self.guiding_pixel, device=x.device), mask], dim=1).bool()
        x = torch.where(mask_with_cond.unsqueeze(-1), self.mask_token.to(x.dtype), x)

        # 添加位置嵌入
        x = x + self.pos_embed_learned
        # 应用LayerNorm
        x = self.patch_emb_ln(x)

        # 应用Transformer块
        if self.grad_checkpointing and not torch.jit.is_scripting() and self.training:
            for block in self.blocks:
                # 使用梯度检查点
                x = checkpoint(block, x)
        else:
            for block in self.blocks:
                # 应用Transformer块
                x = block(x)
        # 应用LayerNorm
        x = self.norm(x)

        # 返回五个条件：中间、上、右、下、左
        middle_cond = x[:, self.num_conds+self.guiding_pixel:] # 中间条件
        bsz, seq_len, c = middle_cond.size()
        h = int(np.sqrt(seq_len)) # 计算高度
        w = int(np.sqrt(seq_len)) # 计算宽度

        # 上条件
        top_cond = middle_cond.reshape(bsz, h, w, c)
        top_cond = torch.cat([torch.zeros(bsz, 1, w, c, device=top_cond.device), top_cond[:, :-1]], dim=1)
        top_cond = top_cond.reshape(bsz, seq_len, c)

        # 右条件
        right_cond = middle_cond.reshape(bsz, h, w, c)
        right_cond = torch.cat([right_cond[:, :, 1:], torch.zeros(bsz, h, 1, c, device=right_cond.device)], dim=2)
        right_cond = right_cond.reshape(bsz, seq_len, c)

        # 下条件
        bottom_cond = middle_cond.reshape(bsz, h, w, c)
        bottom_cond = torch.cat([bottom_cond[:, 1:], torch.zeros(bsz, 1, w, c, device=bottom_cond.device)], dim=1)
        bottom_cond = bottom_cond.reshape(bsz, seq_len, c)

        # 左条件
        left_cond = middle_cond.reshape(bsz, h, w, c)
        left_cond = torch.cat([torch.zeros(bsz, h, 1, c, device=left_cond.device), left_cond[:, :, :-1]], dim=2)
        left_cond = left_cond.reshape(bsz, seq_len, c)

        # 返回预测结果列表
        return [middle_cond, top_cond, right_cond, bottom_cond, left_cond]

    def forward(self, imgs, cond_list):
        """
        前向传播函数，用于训练阶段。

        参数:
            imgs (torch.Tensor): 输入图像张量，形状为 (batch_size, channels, height, width)。
                                 其中，batch_size 是批量大小，channels 是通道数，height 和 width 分别是图像的高度和宽度。
            cond_list (List[torch.Tensor]): 条件输入列表，包含用于生成的条件信息。

        返回:
            Tuple[torch.Tensor, List[torch.Tensor], torch.Tensor]: 
                - patches (torch.Tensor): 分割后的补丁，形状为 (batch_size * patches_per_image, 3, patch_size, patch_size)。
                - cond_list_next (List[torch.Tensor]): 下一级的条件列表。
                - guiding_pixel_loss (torch.Tensor): 引导像素损失。
        """
        # patchify to get gt
        # 将输入图像分割成patch以获取标签
        patches = self.patchify(imgs)

        # mask tokens
        # 生成随机排列的顺序
        orders = self.sample_orders(bsz=patches.size(0))
        if self.training:
            # 在训练模式下，使用随机掩码
            mask = self.random_masking(patches, orders)
        else:
            # uniform random masking for NLL computation
            # 在非训练模式下，使用均匀随机掩码
            mask = self.random_masking_uniform(patches, orders)

        # guiding pixel
        if self.guiding_pixel:
            # 计算引导像素（对图像的高度和宽度维度求均值）
            guiding_pixels = imgs.mean(-1).mean(-1)
            # 计算引导像素损失
            guiding_pixel_loss = self.guiding_pixel_loss(guiding_pixels, cond_list)
            # 将引导像素添加到条件列表中
            cond_list.append(guiding_pixels)
        else:
            # 否则，损失为0
            guiding_pixel_loss = torch.Tensor([0]).cuda().mean()

        # get condition for next level
        # 获取下一级的条件
        cond_list_next = self.predict(patches, mask, cond_list)

        # only keep those conditions and patches on mask
        for cond_idx in range(len(cond_list_next)):
            # 重塑条件张量
            cond_list_next[cond_idx] = cond_list_next[cond_idx].reshape(cond_list_next[cond_idx].size(0) * cond_list_next[cond_idx].size(1), -1)
            # 仅保留被掩盖的位置的条件
            cond_list_next[cond_idx] = cond_list_next[cond_idx][mask.reshape(-1).bool()]

        # 重塑patch张量
        patches = patches.reshape(patches.size(0) * patches.size(1), -1)
        # 仅保留被掩盖位置的patch
        patches = patches[mask.reshape(-1).bool()]
        # 重塑patch形状为 (batch_size, 3, patch_size, patch_size)
        patches = patches.reshape(patches.size(0), 3, self.patch_size, self.patch_size)

        # 返回patch、下一级条件和引导像素损失
        return patches, cond_list_next, guiding_pixel_loss

    def sample(self, cond_list, num_iter, cfg, cfg_schedule, temperature, filter_threshold, next_level_sample_function,
               visualize=False):
        """
        生成函数，用于从条件输入中采样生成图像。

        参数:
            cond_list (List[torch.Tensor]): 条件输入列表，包含用于生成的条件信息。
            num_iter (int): 采样迭代次数。
            cfg (float): 无分类器指导的权重，用于控制生成的多样性。
            cfg_schedule (str): cfg调度策略，可以是"linear"或其他自定义策略。
            temperature (float): 温度参数，控制生成的多样性。
            filter_threshold (float): 过滤阈值，用于筛选生成的样本。
            next_level_sample_function (callable): 下一级的采样函数，用于生成下一个位置的样本。
            visualize (bool): 是否可视化生成过程，默认为False。

        返回:
            torch.Tensor: 生成的结果，形状为 (batch_size, channels, height, width)。
        """
        if cfg == 1.0:
            # 如果cfg为1.0，批量大小为条件列表中第一个张量的批量大小
            bsz = cond_list[0].size(0)
        else:
            # 否则，批量大小减半，用于实现无分类器指导
            bsz = cond_list[0].size(0) // 2

        # sample the guiding pixel
        if self.guiding_pixel:
            # 采样引导像素
            sampled_pixels = self.guiding_pixel_loss.sample(cond_list, temperature, cfg, filter_threshold)
            if not cfg == 1.0:
                # 如果cfg不为1.0，则拼接引导像素
                sampled_pixels = torch.cat([sampled_pixels, sampled_pixels], dim=0)
            # 将引导像素添加到条件列表中
            cond_list.append(sampled_pixels)

        # 初始化掩码张量，形状为 (batch_size, seq_len)
        mask = torch.ones(bsz, self.seq_len).cuda()
        # 初始化patch张量，形状为 (batch_size, seq_len, 3 * patch_size**2)
        patches = torch.zeros(bsz, self.seq_len, 3 * self.patch_size**2).cuda()
        # 生成随机排列的顺序
        orders = self.sample_orders(bsz)
        # 确定迭代次数，确保不超过序列长度
        num_iter = min(self.seq_len, num_iter)

        # sample image
        for step in range(num_iter):
            # 克隆当前patch
            cur_patches = patches.clone()

            if not cfg == 1.0:
                # 如果cfg不为1.0，则拼接patch
                patches = torch.cat([patches, patches], dim=0)
                # 拼接掩码
                mask = torch.cat([mask, mask], dim=0)

            # 获取下一级的条件
            cond_list_next = self.predict(patches, mask, cond_list)

            # 计算下一个迭代的掩码比例
            mask_ratio = np.cos(math.pi / 2. * (step + 1) / num_iter)
            # 计算要掩盖的标记数量
            mask_len = torch.Tensor([np.floor(self.seq_len * mask_ratio)]).cuda()

            # 确保至少有一个标记被掩盖
            mask_len = torch.maximum(torch.Tensor([1]).cuda(),
                                     torch.minimum(torch.sum(mask, dim=-1, keepdims=True) - 1, mask_len))

            # 生成下一个迭代的掩码
            mask_next = mask_by_order(mask_len[0], orders, bsz, self.seq_len)
            if step >= num_iter - 1:
                # 如果是最后一个迭代，则预测当前掩码位置
                mask_to_pred = mask[:bsz].bool()
            else:
                # 否则，预测被掩盖的位置
                mask_to_pred = torch.logical_xor(mask[:bsz].bool(), mask_next.bool())
            # 更新掩码
            mask = mask_next
            if not cfg == 1.0:
                # 如果cfg不为1.0，则拼接掩码
                mask_to_pred = torch.cat([mask_to_pred, mask_to_pred], dim=0)

            # 仅保留被预测位置的下一级条件
            for cond_idx in range(len(cond_list_next)):
                cond_list_next[cond_idx] = cond_list_next[cond_idx][mask_to_pred.nonzero(as_tuple=True)]

            # 线性调度cfg
            if cfg_schedule == "linear":
                cfg_iter = 1 + (cfg - 1) * (self.seq_len - mask_len[0]) / self.seq_len
            else:
                # 否则，使用cfg
                cfg_iter = cfg
            # 采样patch
            sampled_patches = next_level_sample_function(cond_list=cond_list_next, cfg=cfg_iter,
                                                         temperature=temperature, filter_threshold=filter_threshold)
            # 重塑patch形状
            sampled_patches = sampled_patches.reshape(sampled_patches.size(0), -1)

            if not cfg == 1.0:
                # 如果cfg不为1.0，则拆分掩码
                mask_to_pred, _ = mask_to_pred.chunk(2, dim=0)

            # 更新当前patch
            cur_patches[mask_to_pred.nonzero(as_tuple=True)] = sampled_patches.to(cur_patches.dtype)
            # 更新patch
            patches = cur_patches.clone()

            if visualize:
                # 可视化生成过程
                visualize_patch(self.unpatchify(patches))

        # 将patch重新组合成图像
        patches = self.unpatchify(patches)

        # 返回生成的图像
        return patches