OSASIS（One-Shot Structure-Aware Stylized Image Synthesis）

摘要

本周阅读了一篇关于新型图像风格化的论文《One-Shot Structure-Aware Stylized Image Synthesis》，旨在解决现有GAN模型在风格化过程中难以保持输入图像结构的问题。通过分离图像的结构和语义信息，实现对内容和风格的精确控制。该方法采用了扩散模型和结构保持网络，以及clip方向性损失，以提高风格化图像的原始质量，并在结构上保持方面表现出色。

abstract

This week I read a paper on a new type of Image stylization, “One-Shot Structure-Aware Stylized Image Synthesis”, which aims to solve the problem that existing GAN models have difficulty maintaining input image Structure during the stylization process. By separating the structure and semantic information of the image, the precise control of the content and style is realized. The method uses the diffusion model and structure retention network, as well as the clip directivity loss, to improve the original quality of the stylized image, and performs well in structure retention.

论文摘要

论文中主要提出一种新颖的单发性风格化方法Osasis，在结构保持方面非常强大。OSASIS能够有效地从图像的结构中分离语义，允许它控制对给定输入实现的内容和样式的级别。实验表明对于在训练过程中很少遇到的输入图像优于其他风格化方法。

OSASIS能够使用单个参考图像对输入图像进行风格化，同时稳健地保留输入图像的结构和内容。

在微调期间，跨域损失将逼真的图像（有界的黄色）与其程式化的对应图像（有界的绿色）进行比较。同时，域内损耗测量同一域内方向位移的对齐，用黄色和绿色描绘。重建损失将原始风格图像与重建的对应图像进行比较。直观地说，方向损失的组合保证了每次迭代生成的I_Bⁱⁿ被定位为来自I^style_B的投影向量和I_Aⁱⁿ与CLIP空间中的跨域和域内对应向量共线.

方法

通过逐渐解开图像的结构和语义信息来实现有效的风格化。将结构信息定义为图像的整体概述，图像的综合性进一步解构为内容和样式的组合。

在推理期间控制低级视觉特征(例如纹理和颜色)的数量，称为风格，以及高级视觉特征(例如对象和身份)的数量，称为内容。

采用了两种不同的潜在代码：结构潜在代码x_t0和语义潜在代码Z_sem，Z_sem适当地调节到微调后的DDIM，使我们能够实现这种控制，有效地执行风格化。此外，我们直接对语义潜在代码z_sem进行了优化，用于文本驱动的操作。通过将优化的潜影与微调的DDIM相结合，OSASIS能够生成具有操纵属性的风格化图像。

损失函数

在训练阶段为了解决结构信息丢失问题，论文引入了一个结构保留网络SPN,它利用1*1卷积有效的保留了I_Aⁱⁿ的空间信息和结构完整性。为了生成下一个时间步x_t-1的输出，使用SPN反向DiffAE:
$$\begin{gathered} {\mathbf{x}}_t^{SPN}=SPN(I_A^{\mathrm{in}}) \\ {\mathbf{x}}_t^{\prime }={\mathbf{x}}_t+{\lambda }_{SPN}\ast {\mathbf{x}}_t^{SPN} \\ {\mathbf{x}}_{t-1}=\sqrt{{\alpha }_{t-1}}{f}_{\theta }({\mathbf{x}}_t^{\prime },t,{\mathbf{z}}_{\mathrm{sem}}^{\mathrm{in}})+\sqrt{1-{\alpha }_{t-1}}{ϵ}_{\theta }^B({\mathbf{x}}_t^{\prime },t,{\mathbf{z}}_{\mathrm{sem}}^{\mathrm{in}}) \end{gathered}$$
损失函数：包括跨域损耗、域内损耗和重构损耗。跨域损失的目的是对齐从域A到域B的变化方向，确保从I_Aⁱⁿ到I_Bⁱⁿ的变化与从I_A^style到的变化保持一致I_B^style。虽然跨域损失为模型优化提供了语义信息的变化，但在单独实现时往往会导致意想不到的变化。因此，引入域内损失来提供额外的信息，测量域A和域B内变化的相似性。重建损失为捕获从I_A^style到I_B^style的跨域变化提供了额外的指导。

实验

从FFHQ数据集[12]中随机选择20000张图像，使用每个图像的语义子码z_sem
和随机子码x_T ~N(0, I)(即随机重建)来重建每个图像。我们使用感知相似度损失将重建图像与原始图像进行比较，以确定重建的质量。

上图论文中的方法成功地保留了低密度属性，而其他基线方法却没有做到这一点。

由于基于gan的反演方法的能力有限，基线方法无法从样式图像的结构和语义中解脱出来。这导致结构伪影被转移到输出图像中，而OSASIS仅成功地提取语义。

代码
下面代码的核心思想是将一张目标图像（即 img_ref_B）加噪到某个时间步，然后从这个噪声状态通过扩散模型反向还原多次，并从中挑选最接近原图的一张。

def main():
    args = create_argparser().parse_args()

    # set seed
    random.seed(args.seed)
    np.random.seed(args.seed)
    th.manual_seed(args.seed)

    dist_util.setup_dist()
    logger.configure(dir=args.sample_dir)

    logger.log("creating model and diffusion...")
    model, diffusion = create_model_and_diffusion(
        **args_to_dict(args, model_and_diffusion_defaults().keys())
    )
    model.load_state_dict(
        dist_util.load_state_dict(args.model_path, map_location="cpu")
    )
    model.to(dist_util.dev())
    if args.use_fp16:
        model.convert_to_fp16()
    model.eval()

    logger.log("sampling...")

    device = dist_util.dev()
    transform_256 = transforms.Compose([
        transforms.Resize((256,256)),
        transforms.ToTensor(),
        transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))  
    ])

    # imgs_ref_domainB = glob.glob(f'{args.input_dir}/*')
    imgs_ref_domainB = [f'{args.input_dir}/art_{args.seed}.png']
    percept_loss = lpips.LPIPS(net='vgg').to(device)
    l1 = th.nn.L1Loss(reduction='none')

    # for img_file in tqdm(imgs_ref_domainB):
    #     # file_name = img_file.split('/')[-1]
    # file_name = Path(img_file).name
    
    img_file = os.path.join(args.input_dir, args.img_name)
    img_ref_B = Image.open(img_file).convert('RGB')
    img_ref_B = transform_256(img_ref_B)
    img_ref_B = img_ref_B.unsqueeze(0).to(device)
    img_ref_B_all = img_ref_B.repeat(args.n,1,1,1)

    t_start = int(diffusion.num_timesteps*args.t_start_ratio)
    t_start = th.tensor([t_start], device=device)
    
    # forward DDPM
    xt = diffusion.q_sample(img_ref_B_all.clone(), t_start)

    # reverse DDPM
    indices = list(range(t_start))[::-1]
    for i in indices:
        t = th.tensor([i] * img_ref_B_all.shape[0], device=device)
        with th.no_grad():
            out = diffusion.p_sample(
                model,
                xt,
                t,
                clip_denoised=True,
                denoised_fn=None,
                cond_fn=None,
                model_kwargs=None,
            )
            xt = out["sample"]

        # # reverse DDIM
        # indices = list(range(t_start))[::-1]
        # for i in indices:
        #     t = th.tensor([i] * img_ref_B_all.shape[0], device=device)
        #     with th.no_grad():
        #         out = diffusion.ddim_sample(
        #             model,
        #             xt,
        #             t,
        #             clip_denoised=True,
        #             denoised_fn=None,
        #             cond_fn=None,
        #             model_kwargs=None,
        #         )
        #         xt = out["sample"]

        # compute loss
        # from torchvision.utils import save_image
        #计算 L1 和 LPIPS（感知损失）衡量每张重建图像与原图的相似度。
        # save_image(xt/2+0.5, os.path.join('step1_tmp', file_name))
        l1_loss = l1(xt, img_ref_B.repeat(int(args.n),1,1,1))
        l1_loss = l1_loss.mean(dim=(1,2,3))
        lpips_loss = percept_loss(xt, img_ref_B.repeat(int(args.n),1,1,1)).squeeze()
        loss = 10*l1_loss + lpips_loss

        # pick best image
        img_idx = th.argmin(loss)
        img_ref_A = xt[img_idx]
        
        os.makedirs(args.sample_dir, exist_ok=True)
        utils.save_image(img_ref_A/2+0.5, os.path.join(args.sample_dir, args.img_name))

加权组合后选取最小loss的图像作为最终结果。

结论

OSASIS在风格化中表现出强大的结构感知能力，能够有效地从图像中分离出结构和语义。虽然OSASIS在结构感知的风格化方面取得了重大进展，但仍存在一些限制。OSASIS的一个显著限制是它的训练时间比比较方法要长。这种延长的训练时间是该方法保持结构完整性和适应各种风格的增强能力的一种权衡。OSASIS在保持输入图像的结构完整性方面的鲁棒性、在域外参考风格化方面的有效性以及在文本驱动操作中的适应性使其成为风格化图像合成领域的一种很有
前途的方法。

总结

本周阅读的论文《One-Shot Structure-Aware Stylized Image Synthesis》提出了一种结构感知的单图风格化方法 OSASIS。该方法通过解耦图像的结构与语义信息，实现了风格与内容的独立控制，采用扩散模型和结构保持网络（SPN）在保持结构的同时增强风格表达。其关键在于结合CLIP方向损失、多种感知损失，以及语义潜在编码的优化，使生成图像在风格迁移的同时能最大程度地保留原图结构。实验结果显示，OSASIS在低密度区域保留、跨域风格化和文本引导方面均优于现有方法。