[12] In-Context Edit: Enabling Instructional Image Editing with In-Context Generation in Large Scale Diffusion Transformer

 

[Paper] https://arxiv.org/pdf/2504.20690

[Github] https://github.com/River-Zhang/ICEdit.git

GitHub - River-Zhang/ICEdit: Image editing is worth a single LoRA! 0.1% training data for fantastic image editing! Training rele

 

 

1. Abstract 

• Instruction-based image editing 이란?
  • 사용자가 자연어 prompt로 이미지를 원하는대로 수정할 수 있도록 하는 것
  • 현재는 precision과 efficiency사이에 trade-off가 있음.
  • 현재까지는 두 가지의 접근 방법 존재
    • Fine-tuning: 퀄리티가 좋지만 상당한 계산량과 막대한 데이터셋 필요하다.
    • training-free: 비교적 빠르지만 prompt를 정확히 이해하고 수행하는데 어려움이 있다.

⇒ 이러한 점들을 개선한 large scale Diffusion Transformer 를 활용하여 최소한의 데이터와 파라미터 수정만으로도 prompt 기반 이미지 편집을 정확하고 효율적으로 해내는 방법을 제시한다.

 

 

2. Introduction 

• Instruction-based image editing 의 위 두 가지 접근 방법의 단점을 해결하기 위해 최근에는 Diffusion Transformer 사용
  • Scalable Generation Fidelity: auxiliary module 없이 reference-guided synthesis과 identity preserved editing이 가능 ⇒ 매우 강한 생성력을 가짐 (ex. Flux)
  • Intrinsic Contextual Awareness: 원본과 생성될 이미지 사이의 bidirectional interactions 가능 ⇒ source와 target 이미지 동시처리 가능
  ⇒ 과연 DiT의 강력한 생성력과 문맥 인식 능력만으로, 구조 변경이나 대규모 학습 없이도 Instruction-based image editing의 precision과 efficiency을 모두 달성할 수 있을까?
• 두가지의 중요한 통찰
  • In-context prompts 고안
    • 원본 이미지(좌측)와 편집된 이미지(우측)의 쌍을 묘사하는 문장을 prompt로 구성하여, 한 번의 생성으로 왼쪽에는 원본에 해당하는 그림을, 오른쪽에는 편집된 그림을 동시에 생성시키는 prompt
    • 추가적인 모델 구조 변경이나 tuning 없이도 원본 이미지의 특징을 유지하면서 prompt를 실행
  • 초기 noise가 생성 결과에 큰 영향을 미침
    • 특정 prompt에 잘 맞는 seed값이 있음을 알게 됨
  • precision과 efficiency를 동시에 향상 시키는 전략
    • LoRA-MoE Hybrid Tuning:
      • 대규모 DiT 모델의 일부 모듈만 fine-tuning하고 다양한 prompt(remove, add,…)에 대응하기 위해 LoRA와 MoE(Mixture of Expert) 구조를 결합한 것
      • 태스크별 전문화된 경로를 모델 내에 두고, 입력 내용에 따라 동적으로 적합한 expert 경로가 활성화되도록 함
    • Early Filter Inference Time Scaling
      • VLM을 평가자로 활용하여 생성 초반 단계에서 여러 noise seed들에 대한 후보 결과물을 비교 한 뒤 prompt에 가장 잘 부합하는 seed를 고르는 방법
      • 전체 생성 과정을 모두 거치지 않고도 초반 몇 스텝의 결과로 prompt 수행 여부를 빠르게 판별하고 최적 seed만 최종 생성에 사용하기 때문에, precision과 efficiency을 동시에 얻을 수 있음

 

 

3. Method 

3.1. Exploration of DiT’s In-context Edit Ability

In-Context Generation with Edit Instructions

• 특수한 형태의 prompt 설계
  • A diptych with two side-by-side images of the same scene. On the right, the scene is the same as on the left but {prompt}. (“두 장의 나란한 이미지가 있습니다. 오른쪽에는 왼쪽 이미지와 같지만 거기에 {prompt}가 적용되어 있습니다.”)
  • 왼쪽에는 원본 이미지가 주어지고, 오른쪽은 모델이 생성해야 할 편집된 이미지가 주어짐.
    • 하나의 prompt로 "원본과 편집본"을 동시에 인식 가능 (어떤 부분을 유지하고 무엇을 바꿔야 하는지 등)
    • 실제로 편집해야 할 부분이 활성화 되어 있는 것을 attention map을 통해 확인 가능

 

In-Context Edit Framework

• 두 가지 구조의 실험 (T2I DiT vs Inpainting DiT)
  • Text-to-Image DiT 기반
    • 원본 이미지를 DiT가 받아들일 수 있게 역변환(Inverted Noise)을 수행한 뒤, 이를 위 프롬프트의 좌측 토큰에 삽입하여 이미지 생성에 활용하고 우측은 편집 지시에 따라 생성하는 방법.
  • inpainting DiT
    • 입력 캔버스를 미리 왼쪽 절반은 원본 이미지, 오른쪽 절반은 빈 마스크로 채운 후, 동일한 위 프롬프트를 넣어 오른쪽 마스크 부분만 채우도록 하는 방법

 

3.2. LoRA-MoE Hybrid Fine-tuning

LoRA Tuning

• 데이터셋 (약 5만 장)을 구축하여, DiT 모델 일부의 LoRA fine-tuning 수행 ⇒ 적은 양의 데이터로도 편집 성공률과 화질이 크게 향상 됨
• 그러나, 하나의 LoRA 모듈만으로 모든 편집 작업(add, removal, modification,…)을 향상시키는 데에는 한계가 존재
  • add와 modification은 성격이 다르니 다른 조작이 필요한데 단일 LoRA로 다루기에는 특정 작업에서 높은 실패율을 보임 ⇒ MoE(Mixture of Experts) 접목

Mixture of LoRAs

• MoE 구조는 여러 expert 신경망을 병렬로 두고 입력 특성에 따라 적절한 expert를 선택해 활용함으로써, 다양한 입력 패턴에 대응하여 처리할 수 있도록 함
• prompt 내용에 따라 필요한 expert만 활성화 되므로 매우 효율적
• DiT의 multi-modal attention block 내 output layer에 다수의 LoRA 모듈을 병렬로 삽입하여 복수의 expert LoRA들을 둠(출력 부근에서만 여러 LoRA 경로를 제공)
  • routing classifier
    • 예를 들어) 4개의 LoRA expert가 있다면, routing module이 각 토큰마다 4개 expert의 중요도를 예측하고 그 중 상위 1개만 실제 사용(Top-K=1)하는 sparse MoE 방식을 취함
  ⇒ 다양한 편집 작업에 대한 모델 수용 능력을 높이면서 효율은 그대로 유지

 

3.3. Early Filter Inference Time Scaling

• inference 과정 자체를 개선하여 편집 quality를 높이는 기술
• 초기 노이즈의 상태가 최종 결과에 큰 영향을 주게됨⇒이 초기 단계를 똑똑하게 관리하는 것을 목표로 함
  • VLM을 평가자로 이용 (본 논문에서는 CLIP 이용)
  • 작동 방식
  • 1. 다양한 seed로 부터 이미지를 생성해봄
    2. 초기 몇 step의 노이즈 제거가 진행된 이미지들을 CLIP으로 평가
    3. 가장 유망한 후보를 골라 남기고 나머지 후보들은 버리거나 스케일을 조정
    • 처음부터 올바른 방향으로 시작할 수 있어 효율적

 

4. Experiments 

Implementation Details

• FLUX.1 Fill을 backbone으로 사용
• LoRA fine-tuning 데이터 셋
  • MagicBrush 데이터셋(약 9천 쌍의 편집 샘플)에 OmniEdit 데이터셋의 일부(약 4만 샘플) ⇒ 총 약 5만 장

Evaluation Settings

• 성능평가 데이터셋
  • Emu Edit: GT가 제공되므로, CLIP 점수, DINO 점수, L1 loss 등을 통해 편집된 아웃풋과 GT 간의 유사도를 정량 평가
  • MagicBrush: GT 없음 ⇒ GPT-4 기반 평가를 도입
    • GPT-4에게 원본과 출력, prompt를 보여주고 편집이 제대로 이루어졌는지 점수화하는 방법
    • 추가로 VIE-Score 사용: SC(score compliance)와 PQ(perceptual quality) 두 부분으로 구성
      • SC는 prompt 수행 여부와 편집되지 않아야 하는 부분의 보존 정도를, PQ는 이미지 자체의 품질을 의미함

 

4.1. Comparisons with State-of-the-Art

• MagicBrush & Emu Edit dataset 결과
  • 훨씬 적은 파라미터로 기존 모델과 유사하거나 높은 성능을 보임

DINO score는 이미지 간의 시각적 유사성(visual similarity)을 평가하기 위해 사용하는 지표