diffusion model (九) EmuEdit技术小结

Paper: https://emu-edit.metademolab.com/assets/emu_edit.pdf

Project web: https://emu-edit.metademolab.com/

Code: have not opensource

背景

在发布Emu后，近日，META又发布了两个非常惊艳的工作：EmuEdit、EmuVideo。文本将对EmuEdit相关技术进行总结。

1 核心思想

作者将intruction-base image editing任务建模为生成任务，并用diffusion model进行求解。核心创新点有两个

• 详细定义了instruction-based image edit处理的任务，并设计了一个高效高质量的数据构建方法。
• 为提升模型对instruction的理解能力，引入learnable task embedding，能较好的解决上述问题。并且提出task inversion的训练方法，只需少量数据就能有效将模型扩展到新的task（类似textual inversion的思想）。

2 方法

2.1 方法建模

前面提到，作者将一系列的intruction-base image editing任务建模为生成任务，并用diffusion model来求解。具体来看intruction-base image editing任务做的是这么一件事：给定一张参考图片和一段表述文本，输出符合上述两个条件的图片。从上述描述可知：intruction-base image editing的训练数据应当至少是一个三元组$\mathcal{D}=\{(c_T^i,c_I^i,x^i)|i=1,\cdots N\}$，其中

$c_I$: 参考图片(condition of image)

$c_T$: 参考文本（condition of text）

$x$: 目标图片

这样，基于diffusion model的优化目标可建模为：
$$\underset{\theta }{\min }{\mathbb{E}}_{y,ϵ,t}[\Vert ϵ-{ϵ}_{\theta }(z_t,t,E(c_I),c_T){\Vert }_2^2{]}^2\text{\hspace{1em}(1)}$$
和经典的classifier-free有所区别的是，此处多了一个参考图片的condition $E(c_I)$。条件融入的方法上，

• 作者参考Instructpix2pix将image condition融入到输入层(在通道维度进行concat）。

• 参考classifier-free将text condition的融入在cross-attention。

通过实验，作者发现用上述方法训练的模型对task的理解不够准确如下图所示。为此，作者引入learnable task embedding来增强模型对task的理解。此时的优化目标建模为：
$$\underset{\theta ,v_1,\dots ,v_k}{\min }{\mathbb{E}}_{\hat{y},ϵ,t}[\Vert ϵ-{ϵ}_{\theta }(z_t,t,E(c_I),c_T,v_i){\Vert }_2^2]\text{\hspace{1em}(2)}$$
为了求解上述目标方程，构造的训练数据集的每一个元素应当是一个四元组${\mathcal{D}}_{new}=\{(c_T^i,c_I^i,x^i,v_j^i)|i=1,\cdots N\}$，$v_j$为这条数据所属的task类别。并且此时的diffusion model的噪声预测模型多了一个task embedding $v_i$的条件。作者的融入方式是将其与time-step的embedding进行相加，共同融入到cross-attention中。这样设计还保留了可扩展性：当有一个新的task时，可以将优化目标转化为
$$\underset{v_{\mathrm{new}}}{\min }{\mathbb{E}}_{y,ϵ,t}[\Vert ϵ-{ϵ}_{\theta }(z_t,t,E(c_I),c_T,v_{new}){\Vert }_2^2\text{\hspace{1em}(3)}$$
此时训练的参数仅为新增的task embedding，其它参数都freeze。作者将其称之为task inversion(类似textual inversion)。

在用户层的推理阶段，用户无需输入task index，作者基于Flan-T5-XL训练了一个task index预测模型，来根据用户输入的instruction预测出相应的task index。

从实现原理上，上述方法不难想到。论文取得的卓越的效果取决于训练的数据集。下面来看作者是如何用一种高效的方法构建高质量的数据集。

2.2 数据工程

前文提到，训练一个image-edit diffusion model训练数据至少是一个三元组$\mathcal{D}=\{(c_T^i,c_I^i,x^i)|i=1,\cdots N\}$ (其中$c_I$: 参考图片(condition of image) $c_T$: 参考文本（condition of text）$x$: 目标图片)。手动构建数据集的成本非常大，开源数据规模又不够大，一些规模大的合成数据多样性和质量又不高，因此需要探寻如何用cheap的方法来构建一个高质量、大规模、高多样的image-edit数据集。为了结合task inversion，新构建的数据集应当是一个四元组${\mathcal{D}}_{new}=\{(c_T^i,c_I^i,x^i,v_j^i)|i=1,\cdots N\}$, $v_j$为这条数据所属的task类别。

2.2.1 image-edit任务类别定义

作者将image-edit分为了三大类，分别是Region-based Editing、Free-From Editing、Vision tasks，每个大类中有若干小类。下图展示了每一个image-edit任务所做的事

2.2.2 指令集生成

任务定义：已知image caption和编辑任务，输出满足编辑任务新的caption

• 输入：image caption + edit任务

• 输出：edit instruction， edit instruction应当包含：1）edit指令；2）edit的目标(edited object)；3）新的image caption；4）原始目标（original object）（7.2节提到有这个字段，但在7.1中的示例没有，实际上应当要有这个字段，否则后续的mask提取无法进行）

举个例子（对于add的image-edit任务）

输入：{“image_caption”: “Beautiful cat with mojito sitting in a cafe on the street”, “task”: “Add”}

输出：{“edit”: “include a hat”, “edited object”: “hat”, “output”: “Beautiful cat wearing a hat with mojito sitting in a cafe on the street”, “original object”: “cat”}

作者用context learning的方法来实现上述任务的目标。作者构建的prompt方案如下：（作者所用的LLM是微调了的70B LLama2，我用chatgpt尝试了一下，也能实现类似的效果）

import random 
import torch 
from random import choice, shuffle

few_shot_examples = [
    """
    [INST]User: "Beautiful cat with mojito sitting in a cafe on the street"[/INST] 
    Assistant: {"edit": "include a hat", "edited object": "hat", "output": "Beautiful cat wearing a hat with mojito sitting in a cafe on the street"}""",
    """
    [INST]User: "robot playing chess at home."[/INST] 
    Assistant: {"edit": "add a cheerful smiling face.", "edited object": "robot", "output": "robot playing chess at home with a cheerful smiling face."} """,
    """
    [INST]User: "A cute creature sits at the beach."[/INST] 
    Assistant: {"edit": "set a dog besides the creature", "edited object": "dog", "output": "A cute creature and a dog sit at the beach."} """,
    """
    [INST]User: "Superhero on the street in sunny day working on his tablet."[/INST] 
    Assistant: {"edit": "put a vintage tie on the superhero.", "edited object": "tie", "output": "Superhero with a vintage tie on the street in sunny day working on his tablet."} """,
    """
    [INST]User: "Picture clouds, birds, the wind, foliage, rainbow, hill, art, pair, guy"[/INST] 
    Assistant: {"edit": "together with a dog on the left", "edited object": "dog", "output": "Picture clouds, birds, the wind, foliage, rainbow, hill, art, pair, guy, dog on the left"} """,
    """
    [INST]User: "horse on a red Boat Near Mountains During Golden Hour"[/INST] 
    Assistant: {"edit": "give the horse sunglassess", "edited object": "sunglassess", "output": "horse with sunglassess on a red Boat Near Mountains During Golden Hour"} """,
    """
    [INST]User: "An animal family on studio background."[/INST] 
    Assistant: {"edit": "make them hold a teddy bear.", "edited object": "teddy bear", "output": "An animal family holding a teddy bear on studio background."} """,
    """
    [INST]User: "Baked Salmon With Bell Peppers"[/INST] 
    Assistant: {"edit": "insert kale pesto to the dish", "edited object": "kale pesto", "output": "Baked Salmon With Kale Pesto And Bell Peppers"} """,
    """
    [INST]User: "An airplaine is flying in the sky in rainy day."[/INST] 
    Assistant: {"edit": "add flowers in the windows", "edited object": "flowers", "output": "An airplaine with flowers in the windows is flying in the sky in rainy day."} """,
    """
    [INST]User: "photo of mountains and trees"[/INST] 
    Assistant: {"edit": "position a castle between the trees", "edited object": "castle", "output": "photo of mountains, trees and castle between the trees"} """,
    """
    [INST]User: "Little bunny in the park"[/INST] 
    Assistant: {"edit": "Make the bunny play with a kite.", "edited object": "kite", "output": "Little bunny playing with a kite in the park"} """,
    """
    [INST]User: "Attic Bedroom With Large Ceilings"[/INST] 
    Assistant: {"edit": "decorate the room with beautiful chandeliers", "edited object": "chandeliers", "output": "Attic Bedroom With Beautiful Chandeliers on Large Ceilings"} """,
    """
    [INST]User: "Wedding rings and yellow flower on a red background"[/INST]  
    Assistant: {"edit": "place a bird by the yellow flower", "edited object": "bird", "output": "Wedding rings, a bird, and yellow flower on a red background"} """,
    """
    [INST]User: "Tree Near the lake in the morning"[/INST] 
    Assistant: {"edit": "Give it autumn leaves on top", "edited object": "leaves", "output": "Tree with autumn leaves on top Near the lake in the morning"} """,
    """
    [INST]User: "robot and alien sitting on hanging bridge at daytime"[/INST]  
    Assistant: {"edit": "make them hold three books.", "edited object": "threebooks", "output": "robot and alien holding three books while sitting on hanging bridge at daytime"} """,
    """
    [INST]User: "Skogafoss waterfall in the south of Iceland"[/INST]  
    Assistant: {"edit": "Set a colorful rainbow in the backhground!", "edited object": "rainboe", "output": "Skogafoss waterfall with a colorful rainbow in the south of Iceland"} """,
    """
    [INST]User: "Polar Bear with rubber gloves pushing shopping carts"[/INST]  
    Assistant: {"edit": "Make it wear a coat", "edited object": "coat", "output": "Polar Bear with a coat pushing shopping carts"}
    """
    ]


def get_content_instruction(new_prompt): 
    optional_verbs = choice(["include", "place", "position", "set", "incorporate", "alongside", " give", "put", "insert", "together with", "with", "make", "integrate", "have", "append", " make", "add", "include"]) 
     # system message # 
    system_message = f"""
        <<SYS>>You are an assistant that only speaks JSON. Do not write normal text. The assistant answer is 
            JSON with the following string fields: 'edit', 'edited object','output'. 
            Here is the latest conversation between Assistant and User.
        <</SYS>>
        """
     # introduction message 
    intro_message = f"""
    [INST]User: Hi, My job to take a given caption ('input') and to output the following: an instruction for {optional_verbs} an object to the image ('edit'), the object to {optional_verbs} ('edited object'), and the caption with the object ('output'). Please help me do it.
        I will give you the 'input', and you will help. When you reply, use the following format: {"edit": '<instruction>', 'edited object': '<object>', 'output': '<caption>'}
    [/INST]
    Assistant: Sure, I'd be happy to help! Please provide the actual input caption you'd like me to read and I'll assist you with writing an instruction to {optional_verbs} an object to the image, writing the added object and writing the caption with the object."""  
    random.seed(torch.randint(1 << 32, ()).item())
    shuffle(few_shot_examples)
    few_shot_examples = few_shot_examples[:int(len(few_shot_examples) * 0.6)] 
    prompt = system_message + intro_message + "".join(few_shot_examples)  # add the test prompt 
    prompt = prompt + f"[INST]User: {new_prompt}[/INST]"
    return prompt

2.2.3 图片对的生成

通过上面的步骤我们拿到了4元组$(c_T,c_I,x,v_j)$,中的$c_T,c_I,v_j$,其中$c_T$还有很多附加信息：

如：编辑的对象，新的image caption，如：

{“edit”: “include a hat”, “edited object”: “hat”, “output”: “Beautiful cat wearing a hat with mojito sitting in a cafe on the street”, “original object”: “cat”}

此处需要进行的是根据上面的条件，得到对应的图片pair ($x$)。

任务目标：根据输入图片、instruction信息生成对应的图片pair ($x$)并且除了编辑的区域，$x$与$c_I$的差异应当尽可能的小。
$$\begin{array}{rl}& \max \mathrm{SIM}(\mathrm{Instruction},c_I^{\mathrm{Edit}})\\ & \min \mathrm{Dist}(x,c_I^{\mathrm{notEdit}})\end{array}\text{\hspace{1em}(4)}$$

为了解决上述的任务目标，作者提出一种mask-based attention control的方法(相当于DiffEdit和P2P的结合)。具体分为以下几个步骤：

已知条件：

• $\mathrm{Ca}{\mathrm{p}}_{\mathrm{ori}}$: image caption 。 example：Beautiful cat with mojito sitting in a cafe on the street
• $\mathrm{Im}{\mathrm{g}}_{\mathrm{ori}}:$image caption用DM生成的图片
• $\mathrm{Ca}{\mathrm{p}}_{\mathrm{edit}}$:编辑后的image caption。Beautiful cat wearing a hat with mojito sitting in a cafe on the street
• $\mathrm{Ob}{\mathrm{j}}_{\mathrm{ori}}$image caption的原始目标（original object）。cat
• $\mathrm{Ob}{\mathrm{j}}_{\mathrm{edit}}$编辑目标(edited object)：hat

STEP1: 提取mask。将$\mathrm{Ob}{\mathrm{j}}_{\mathrm{ori}}$与$\mathrm{Im}{\mathrm{g}}_{\mathrm{ori}}$送入到SAM+DINO模型中，得到3类mask

1. 精确的mask，有sam+dino的生成
2. 将1中的mask进行膨胀，在进行高斯模糊，作为新的mask
3. 取第一步mask的bounding box作为新的mask

SETP2: 通过mask-based attention control进行图片生成。具体为：先用P2P的cross-attention control的方法将common token的对应的attention map进行注入，随后用diffedit的根据mask融合方法进行融合。

STEP3: 图片Filter。通过上述步骤得到3个目标图片，留存最好的一个。filter的规则为

1. 用CLIP filtering metrics，留存最相关的一个
2. 留存edit image与input image在深度图上的L1 距离最小的图片。
3. …

每一类Edit方法的详细的数据构造细节见论文7.2.3

最后得到的各类训练数据比例如下：

3 结果

EmuEdit在多个测试集取得了SOTA。并且作者公开了一个新的基于EmuEdit的benchmark：https://huggingface.co/datasets/facebook/emu_edit_test_set

一些惊艳结果：