AIGC-风格迁移-style Injection in Diffusion-CVPR2024HighLight-论文精度

Style Injection in Diffusion: A Training-free Approach for Adapting Large-scale Diffusion Models for Style Transfer-CVPR2024HighLight

代码：https://github.com/jiwoogit/StyleID
论文：https://jiwoogit.github.io/StyleID_site/

为了解决风格迁移时在推理阶段微调非常耗时的问题，提出操纵自注意力层的特征作为交叉注意力机制的工作方式,目标是通过利用预训练的大规模文本到图像扩散模型的生成能力来解决艺术风格迁移.在生成过程中，用风格形象的key和value来替代内容的key和value。 此外，引入查询保存和注意力温度缩放来减轻原始内容中断的问题，并引入初始潜在自适应实例归一化（AdaIN）来处理不和谐的颜色（未能传输样式颜色）。 实验结果表明，所提出的方法在传统和基于扩散的风格转移基线方面都超越了最先进的方法。

MOTIVATION

基于扩散模型（Diffusion Models, DM）的风格迁移的一般方法及其局限性：

将免训练风格迁移扩展到大规模预训练 DM 上的应用

相应启发：认为 SA 层是转移风格的有效方法。

此外，强调 SA 层在有风格迁移的特性：

语义相似性的保持：在基于SA的风格迁移中，内容图像的patches（查询，query）如果具有语义上的相似性，它们会与相似的风格（键，key）相互作用。这意味着，即使在风格迁移之后，这些内容图像区域之间的关系仍然得以保持。这有助于保持图像内容的完整性和语义连贯性。也就是说，采用 SA 的风格迁移可以保留原始内容，因为具有相似性的内容块往往会从相应的风格图像块接收相似的注意力分数。

强大的特征表示能力：由于大规模扩散模型（Diffusion Model, DM）具有强大的特征表示能力，content image对应的每个查询区域（each patch of the query）都能与具有相似纹理和语义的键（key）建立更高的相似性。换句话说，内容图像中的每个小区域都能够在风格图像中找到具有相似纹理和语义的对应区域，这种相似性有助于模型在进行风格迁移时更好地保留内容图像的语义信息，并根据风格图像的纹理特征进行调整。这种特性允许模型更精确地识别和迁移与内容图像区域相匹配的风格特征。这种相似性鼓励模型基于内容和风格图像之间局部纹理（例如边缘）的相似性来进行风格迁移。

CONTRIBUTION

RELATED WORKS

Diffusion Model-based Neural Style Transfer

Attention-based Image Editing in DM

预训练的文本到图像扩散模型在图像编辑方面取得了显著进展，这些模型被广泛用于各种图像编辑任务。

LDM(Latent Diffusion Model)

潜在扩散模型(LDM)是一种在低维潜在空间中训练的扩散模型，目的是聚焦于数据的语义比特，降低计算成本。其利用预先训练的编码器encoder，将数据集中的全部图像编码为潜在空间 z中的表示。并在潜在空间 z上训练扩散模型，训练时需要预测在时间步t时潜在表示 z t z_t zt​对应的噪声 ϵ \epsilon ϵ。文本利用stable diffusion:
L L D M = E z , ϵ , t [ ∥ ϵ − ϵ θ ( z t , t , y ) ∥ 2 2 ] , L_{\mathrm{LDM}}=\mathbb{E}_{z,\epsilon,t}[\|\epsilon-\epsilon_\theta(z_t,t,y)\|_2^2], LLDM​=Ez,ϵ,t​[∥ϵ−ϵθ​(zt​,t,y)∥22​],

自注意力快

重点关注利用 SA 块（自注意力机制块）来传输style：
Q = W Q ( ϕ ) , K = W K ( ϕ ) , V = W V ( ϕ ) , ϕ o u t = A t t n ( Q , K , V ) = s o f t m a x ( Q K T d ) ⋅ V , Q=W_{Q}(\phi),K=W_{K}(\phi),V=W_{V}(\phi),\\\phi_{\mathrm{out}}=\mathrm{Attn}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d}})\cdot V, Q=WQ​(ϕ),K=WK​(ϕ),V=WV​(ϕ),ϕout​=Attn(Q,K,V)=softmax(d ​QKT​)⋅V,

METHODS

Attention-based Style Injection

输入风格

启发与方法：

潜在特征获取：

风格注入：

内容干扰问题以及查询保留

原因:仅替换键和值可能导致内容干扰( content disruption).因为随着逆扩散过程的进行，注意力值的变化会逐步改变风格化潜在表示的内容,由于注意力值的变化，原始内容的关键信息可能会逐渐丢失，导致生成的图像内容与原始图像内容有所不同。

实现:混合风格后化的潜在表示的query和内容图像的query

Q ~ t c s = γ × Q t c + ( 1 − γ ) × Q t c s , ϕ o u t c s = A t t n ( Q ~ t c s , K t s , V t s ) , \tilde{Q}_{t}^{cs}=\gamma\times Q_{t}^{c}+(1-\gamma)\times Q_{t}^{cs},\\ \phi_{\mathrm{out}}^{\mathrm{cs}}=\mathrm{Attn}(\tilde{Q}_{t}^{cs},K_{t}^{s},V_{t}^{s}), Q~​tcs​=γ×Qtc​+(1−γ)×Qtcs​,ϕoutcs​=Attn(Q~​tcs​,Kts​,Vts​),

Attention Temperature Scaling

注意力图是通过查询（query）和键（key）特征之间的缩放点积来计算的。

为了解决这个问题，作者测量了注意力图的标准差，同时消除了基于注意力的风格注入。

计算在应用 softmax 之前的attention map,即计算查询和键之间的缩放点积，以得到注意力图,实验验证了这种风格注入往往会降低整个时间步上注意力图的标准差

为了更清晰地校正注意力图，我们引入了注意力温度缩放参数。 具体来说，我们将进行softmax之前的注意力图(attention map)乘以大于1的恒定温度缩放参数τ：

A t t n τ ( Q t c s ~ , K t s , V t s ) = s o f t m a x ( τ Q t c s ~ ( K t s ) T d ) ⋅ V t s , τ > 1 \mathrm{Attn}_{\tau}(\tilde{Q_{t}^{cs}},K_{t}^{s},V_{t}^{s})=\mathrm{softmax}(\frac{\tau \tilde{Q_{t}^{cs}}(K_{t}^{s})^{T}}{\sqrt{d}})\cdot V_{t}^{s},\tau>1 Attnτ​(Qtcs​~​,Kts​,Vts​)=softmax(d ​τQtcs​~​(Kts​)T​)⋅Vts​,τ>1

使用 τ = 1.5 作为默认设置，这是整个时间步长的平均比率。 如图 (b) 所示，确认它有效地校准了注意力图的标准差，使其与其原始值相似.softmax之后的注意力图将比其原始值更清晰

Initial Latent AdaIN

仅使用基于注意力的样式注入进行的样式迁移通常无法捕获给定样式的色调。

分析了 DM 的另一个重要部分:initial latent noise

为了利用两种初始潜在(initial latents)中的有价值信息，作者采用AdaIN（自适应实例归一化）来调节初始潜在，以有效传递色调信息。

z T c s = σ ( z T s ) ( z T c − μ ( z T c ) σ ( z T c ) ) + μ ( z T s ) , z_T^{cs}=\sigma(z_T^s)\left(\frac{z_T^c-\mu(z_T^c)}{\sigma(z_T^c)}\right)+\mu(z_T^s), zTcs​=σ(zTs​)(σ(zTc​)zTc​−μ(zTc​)​)+μ(zTs​),

Experiment

实验设置

作者在LAION数据集上使用了Stable Diffusion 1.4预训练模型，并采用了DDIM采样方法，总共进行了50个时间步（t = {1, …, 50}）的迭代。于超参数的默认设置，作者使用了γ = 0.75和τ = 1.5。

评估协议

AIGC常见图像质量评估可以跳转到这篇博客~
AIGC-常见图像质量评估MSE、PSNR、SSIM、LPIPS、FID、CSFD，余弦相似度

ArtFID

为了避免传统风格迁移方法可能过度拟合 Style Loss 的问题，作者采用了 ArtFID 作为评估指标，它综合考虑了内容和风格的保持，并与人类判断高度一致。

A r t F I D = ( 1 + L P I P S ) ⋅ ( 1 + F I D ) ArtFID = (1 + LPIPS) · (1 + FID) ArtFID=(1+LPIPS)⋅(1+FID)

LPIPS 测量风格化图像和相应内容图像之间的内容保真度

FID 评估风格化图像和相应风格图像之间的风格保真度

dataset

使用 MSCOCO 数据集的内容图像和 WikiArt 数据集的风格图像进行评估。 所有输入图像均经过中心裁剪，分辨率为 512 × 512。 此外，为了进行定量比较，从每个数据集中随机选择 20 个内容图像和 40 个风格图像，产生 800 个风格化图像，如 StyTR2所做的那样。

Content Feature Structural Distance-CFSD内容特征结构距离

在风格迁移评估中，内容保真度通常依赖于LPIPS距离，该指标使用了在ImageNet数据集上预训练的AlexNet模型的特征空间，这使得LPIPS对纹理有偏见。图像的风格信息可能会影响LPIPS分数，因为它偏向于纹理特征。为了减少风格信息对评估的影响，作者引入了CFSD，这是一种只考虑图像块之间空间相关性的新距离度量。

CFSD的计算步骤：

CFSD公式：

C F S D = 1 h w ∑ i = 1 h w D K L ( S i c ∣ ∣ S i c s ) , \mathrm{CFSD}=\frac{1}{hw}\sum_{i=1}^{hw}D_{\mathrm{KL}}(S_{i}^{c}||S_{i}^{cs}), CFSD=hw1​i=1∑hw​DKL​(Sic​∣∣Sics​),

#F:\AIGCprojects\StyleID\evaluation\eval_artfid.py
def compute_patch_simi(path_to_stylized, path_to_content, batch_size, device, num_workers=1):
    """Computes the distance for the given paths.

    Args:
        path_to_stylized (str): Path to the stylized images.
        path_to_style (str): Path to the style images. [注：这里应该为 path_to_content，修正为 path_to_content]
        batch_size (int): Batch size for computing activations.
        content_metric (str): Metric to use for content distance. Choices: 'lpips', 'vgg', 'alexnet' [注：缺少 content_metric 参数]
        device (str): Device for computing activations.
        num_workers (int): Number of threads for data loading.

    Returns:
        (float) FID value. [注：文档中写的是 FID value，但函数名为 compute_patch_simi，可能存在混淆，需要核对和确认]

    """
    device = torch.device('cuda') if device == 'cuda' and torch.cuda.is_available() else torch.device('cpu')

    # 根据路径获取图像路径并排序以匹配样式化图像与对应的内容图像
    stylized_image_paths = get_image_paths(path_to_stylized, sort=True)
    content_image_paths = get_image_paths(path_to_content, sort=True)

    # 确保样式化图像和内容图像数量相等
    assert len(stylized_image_paths) == len(content_image_paths), 'Number of stylized images and number of content images must be equal.'

    # 定义图像转换方法
    style_transforms = ToTensor()
    
    # 创建样式化图像的数据集和数据加载器
    dataset_stylized = ImagePathDataset(stylized_image_paths, transforms=style_transforms)
    dataloader_stylized = torch.utils.data.DataLoader(dataset_stylized,
                                                      batch_size=batch_size,
                                                      shuffle=False,
                                                      drop_last=False,
                                                      num_workers=num_workers)

    # 创建内容图像的数据集和数据加载器
    dataset_content = ImagePathDataset(content_image_paths, transforms=style_transforms)
    dataloader_content = torch.utils.data.DataLoader(dataset_content,
                                                     batch_size=batch_size,
                                                     shuffle=False,
                                                     drop_last=False,
                                                     num_workers=num_workers)
    
    # 初始化用于计算距离的度量类
    metric = image_metrics.PatchSimi(device=device).to(device)

    dist_sum = 0.0
    N = 0
    pbar = tqdm(total=len(stylized_image_paths))
    
    # 遍历样式化图像和内容图像的批次
    for batch_stylized, batch_content in zip(dataloader_stylized, dataloader_content):
        # 在不计算梯度的上下文中进行操作，节省内存和计算资源
        with torch.no_grad():
            # 计算当前批次的距离
            batch_dist = metric(batch_stylized.to(device), batch_content.to(device))
            N += batch_stylized.shape[0]
            dist_sum += torch.sum(batch_dist)

        pbar.update(batch_stylized.shape[0])

    pbar.close()

    return dist_sum / N

def compute_cfsd(path_to_stylized, path_to_content, batch_size, device, num_workers=1):
    """Computes CFSD for the given paths.

    Args:
        path_to_stylized (str): Path to the stylized images.
        path_to_content (str): Path to the content images.
        batch_size (int): Batch size for computing activations.
        device (str): Device for computing activations.
        num_workers (int): Number of threads for data loading.

    Returns:
        (float) CFSD value.
    """
    print('Compute CFSD value...')

    # 计算 Patch Similarity，该函数返回样式化图像和内容图像的距离值
    simi_val = compute_patch_simi(path_to_stylized, path_to_content, 1, device, num_workers)
    
    # 将距离值保留四位小数
    simi_dist = f'{simi_val.item():.4f}'
    return simi_dist

#evaluation\image_metrics.py
class PatchSimi(nn.Module):

    def __init__(self, device=None):
        # 初始化函数
        super(PatchSimi, self).__init__()
        # 加载预训练的 VGG19 模型，并移到指定设备上进行评估
        self.model = models.vgg19(pretrained=True).features.to(device).eval()
        # 指定层名称和替换名称的映射
        self.layers = {"11": "conv3"}
        # 图像归一化的均值和标准差
        self.norm_mean = (0.485, 0.456, 0.406)
        self.norm_std = (0.229, 0.224, 0.225)
        # KL 散度损失函数
        self.kld = torch.nn.KLDivLoss(reduction='batchmean')
        
        self.device = device

    def get_feats(self, img):
        features = []
        # 遍历 VGG19 模型的各层并提取特征
        for name, layer in self.model._modules.items():
            img = layer(img)
            if name in self.layers:
                features.append(img)
        return features
    
    def normalize(self, input):
        # 图像归一化处理
        return transforms.functional.normalize(input, self.norm_mean, self.norm_std)

    def patch_simi_cnt(self, input):
        b, c, h, w = input.size()
        # 转置和重塑特征
        input = torch.transpose(input, 1, 3)
        features = input.reshape(b, h*w, c).div(c)
        feature_t = torch.transpose(features, 1, 2)
        # 计算内容图像的特征相似度
        patch_simi = F.log_softmax(torch.bmm(features, feature_t), dim=-1)
        return patch_simi.reshape(b, -1)

    def patch_simi_out(self, input):
        b, c, h, w = input.size()
        # 转置和重塑特征
        input = torch.transpose(input, 1, 3)
        features = input.reshape(b, h*w, c).div(c)
        feature_t = torch.transpose(features, 1, 2)
        # 计算样式化图像的特征相似度
        patch_simi = F.softmax(torch.bmm(features, feature_t), dim=-1)
        return patch_simi.reshape(b, -1)

    def forward(self, input, target):
        src_feats = self.get_feats(self.normalize(input))
        target_feats = self.get_feats(self.normalize(target))
        init_loss = 0.
        # 计算各层的 KL 散度并求和作为初始损失值
        for idx in range(len(src_feats)):
            init_loss += F.kl_div(self.patch_simi_cnt(src_feats[idx]), self.patch_simi_out(target_feats[idx]), reduction='batchmean')

效果

与传统风格迁移方法的比较：时间快，效果好

作者的方法在保持内容图像的结构信息方面表现更好，同时能够有效地迁移风格: