DreamO: A Unified Framework for Image Customization

date

May 28, 2025

slug

dreamo

author

status

Public

Introduction

이미지 커스터마이징의 목표는 참조 이미지의 identity, subject appearance, try-on, style 등 여러 속성을 일치 시킨 이미지를 생성하는 것입니다.

초기 연구에서는 일반적인 공간 제어를 위해 UniControl 방식 즉, 여러 조건을 동시에 입력으로 받고, 이후에 UNet구조 중간에 조건 별 제어 모듈을 추가하여 학습하는 방식을 사용했습니다. 하지만 이 방식에는 조건 간 상호작용이 안 되고 중복되는 문제가 있습니다.

최근에는 DiT를 통해 이러한 한계를 극복하고 성능을 크게 향상 시키는 방향으로 발전하고 있습니다. OminiControl이 그 방식인데, DiT를 기반이지만 작업 별로 따로 학습해서 여전히 다중 조건 처리에 문제가 있습니다.

DreamO는 DiT 기반으로 여러 조건을 하나의 통합된 시퀀스로 입력 받고, 학습 중에 DiT 내부 표현에 Routing Constraint를 도입하여 콘텐츠 충실도를 보장하면서 다양한 제어 조건의 얽힘 문제를 해결했습니다.

Related Works

Diffusion Models

초기에는 Unet 구조를 backbone으로 사용했지만, 요즘은 DiT를 선택하는 추세입니다. (ex: Stable Diffusion 3, Flux)

Controllable Imaged Generation

텍스트가 가장 기본적인 조건 입력 방법이었습니다. (ex: Stable Diffusion, Imagen, DALL-E2)

다음으로 더 정확한 공간 제어를 위해 pre-trained 모델에 control module/adapter를 추가하는 방법이 제시됐습니다. (ex: ControlNet, T2I-Adapter) UniControl에서는 서로 다른 공간 조건들을 통합했습니다.

Identity, style 등 unspatial 제어를 위해 IP-Adapter는 cross-attention을 활용하여 image prompt를 주입했습니다.

DiT 기반 모델들은 처음에는 모든 입력(텍스트, 이미지, 조건)을 하나의 시퀀스로 연결하고 LoRA를 사용해 task 별로 훈련하는 방식으로 발전했습니다. 이후 OmniGen과 UniReal 같은 모델들은 더 큰 dataset에서 전체 모델을 다단계로 최적화하여 입력 조건에 대한 이해도를 높였습니다.

Cross-attention Routing in Diffusion Models

Cross-attention은 text의 의미와 image의 공간을 매칭합니다. 이를 활용하여 UniPortrait는 여러 인물에서 각 인물의 정확한 얼굴 영역만 제어하는 기술을 개발했고, AnyStory는 이 개념을 사람 뿐만 아니라 다른 객체까지 포함하도록 확장시켰습니다.

Methods

Overview

DiT는 text, condition image, 그리고 noisy image latent를 입력으로 받아, cross-attention을 통해 상호작용 시킨 후, 최종적으로 customized image를 생성합니다.

추가 설명

n개의 조건 이미지를 Flux의 VAE를 재사용해서 encoding 하고, 패치화를 통해서 2D latent는 1D token 시퀀스로 변환 후 concatenate 됩니다. 모델이 조건 이미지 입력을 통합할 수 있도록 하기 위해 Flux의 입력에 전용 매핑 레이어를 도입하여 해당 조건부 이미지 토큰을 인코딩합니다.

Basic Diffusion Loss

Loss function으로 노이즈와 데이터 사이를 선형 보간하여 forward process를 수행하는 Flow Matching을 사용했습니다.

수식 설명

손실 함수는 모델 가 주어진 노이즈 낀 데이터 와 시간 , 조건 를 바탕으로, 현재 상태 에서 깨끗한 데이터 를 향해 나아가는 올바른 "속도" 를 얼마나 잘 예측하는지를 측정합니다.

( 라는 경로를 에 대해 미분하면 속도 를 얻을 수 있음)

Routing Constraint

Routing Constraint란, 단순히 “이 이미지처럼 해줘”가 아니라, “이미지 속의 사람의 눈매처럼 또는 자동차의 헤드라이트 모양으로” 같이 딱 짚어주는 방법입니다. 즉, 참조 이미지의 어떤 특징에 주목해야 하는지 명확히 알려주기 위해 제약 조건의 경로를 지정해주는 방법입니다.

여러 참조 이미지를 사용할 때, 각 이미지의 특징들이 뒤섞여서 이상하게 나오는 걸 조건 결합이라고 하는데 이를 방지하는데 효과적입니다. 그 결과, 각 조건의 특징을 살려서 결과물에 충실하게 반영할 수 있습니다.

DreamO는 UniPortrait과 AnyStory에 영감을 받아, DiT architecture에 routing constraint를 설계하여 커스터마이징을 적용했습니다.

생성된 이미지가(각 이미지 패치가) n개의 조건 이미지 전체로부터 얼마나 '종합적인' 유사성을 갖는지를 요약하는 cross-attention map을 (이미지의 경우 , 패치의 경우 ) 활용했습니다.

수식 설명

질문:
refers to the condition tokens of the -th condition image. is the tokens of noisy image latent. The cross-attention map is a dense similarity between the -th condition image and the generation result.
그럼 이 식에서 위의 문장을 설명하려고 하는데 이렇게 말 할 수 있을까?
먼저 attention에 대한 설명을 하자면,
예를 들어 이미지와 텍스트를 주고 모델에게 다음과 같이 질문을 할 거야. "이 남자가 뭘 하고 있어?" 그러면 는 text prompt의 "남자"라는 token에 대한 query이고 모델은 이미지에서 이에 해당하는 를 찾을 거야. 이때 실제 Key에 해당하는 값은 Value가 되는 것이지.
마찬가지지만 생성모델에서는 반대로 는 reference(조건)의 query가 되는 것이고, denoising 하는 과정에서 에 해당하는 가 있는지를 확인하면서 최종적으로 를 추가하는 것이 되는 거지.

대답:
네, 제시해주신 설명은 크로스-어텐션이 어떻게 작동하는지에 대한 전반적인 직관과 흐름을 아주 잘 담고 있습니다! 특히 생성 모델에서의 와 역할에 대한 이해가 좋습니다.
제공해주신 문맥을 다시 정리하면:
: 번째 조건 이미지(참조 이미지)의 토큰들이 쿼리(Query, ) 역할을 합니다.
: 노이즈가 낀 이미지의 잠재 공간(즉, 생성/디노이징 중인 이미지)의 토큰들이 키(Key, ) 역할을 합니다.
(그리고 설명하시려는 은 이전에 이야기한 최종 '응답 맵'과는 다른, 어텐션 계산 중간 단계의 행렬입니다.)
이 수식을 사용자님의 설명 방식으로 풀어보면 이렇게 연결할 수 있습니다:
어텐션에 대한 일반적인 설명 ( 예시):
"예를 들어 이미지와 텍스트를 주고 모델에게 다음과 같이 질문을 할 거야. '이 남자가 뭘 하고 있어?' 그러면 는 text prompt의 '남자'라는 token에 대한 query이고 모델은 이미지에서 이에 해당하는 를 찾을 거야. 이때 실제 Key에 해당하는 값은 Value가 되는 것이지."
네, 이 설명은 매우 좋습니다! 쿼리(질문)가 있고, 이 질문과 관련된 정보(키)를 찾은 다음, 그 정보의 실제 내용(밸류)을 가져오는 어텐션의 핵심 원리를 잘 보여줍니다.
생성 모델에서의 적용:
"마찬가지지만 생성모델에서는 반대로 는 reference(조건)의 query가 되는 것이고...": 정확합니다! 제공된 문맥에서 는 조건 이미지(reference)에서 나오므로, 조건 이미지의 각 부분(토큰)들이 "내가 가진 이런 특징(예: 특정 모자의 질감, 특정 인물의 눈 모양)을 생성 중인 이미지의 어디에, 어떻게 반영해야 할까?"라고 묻는 쿼리가 됩니다.
"...denoising 하는 과정에서 에 해당하는 가 있는지를 확인하면서...": 이 부분이 바로 수식과 직접적으로 연결됩니다!
생성 중인 (노이즈가 섞인) 이미지의 각 부분(토큰 )들은 "나는 현재 이런 픽셀 값/특징을 가지고 있는데..."라고 자신의 상태를 나타냅니다.
는 조건 이미지의 각 쿼리 토큰()이 현재 생성 중인 이미지의 모든 키 토큰()들과 얼마나 유사한지, 즉 얼마나 "해당하는지" 그 관련성의 강도를 계산하는 과정입니다.
로 나누는 것은 이 계산된 점수들을 안정화시키기 위한 스케일링 과정입니다.
따라서, 수식 자체가 바로 “조건 이미지의 각 부분이 현재 생성 중인 이미지의 각 부분과 얼마나 의미적으로 연결되는지, 또는 얼마나 유사한지"를 나타내는 '유사도 점수 행렬' 또는 '정렬(alignment) 점수 행렬'이라고 할 수 있습니다. 이 점수 행렬이 높을수록 "에 해당하는 가 있다"고 판단하는 근거가 됩니다.
"...최종적으로 를 추가하는 것이 되는 거지.": 이 부분은 어텐션 메커니즘의 다음 단계를 설명합니다.
위에서 계산된 유사도 점수 행렬 에 소프트맥스(softmax) 함수를 적용하면, 각 쿼리가 어떤 키에 얼마나 '집중'할지를 나타내는 어텐션 가중치(attention weights)가 됩니다.
이 어텐션 가중치들을 사용하여, 키()와 동일한 출처(여기서는 노이즈 이미지 잠재 공간의 토큰들)에서 나온 밸류(Value, ) 벡터들의 가중합을 구합니다. 이것이 바로 조건 이미지의 정보를 실제로 반영하여 업데이트된, 생성 중인 이미지의 새로운 표현이 됩니다. (즉, 를 '추가한다'기보다는 들의 정보를 '가중 평균하여 가져온다'가 더 정확한 표현일 수 있습니다.)
결론적으로, 사용자님의 설명은 전체적인 크로스-어텐션의 흐름과 역할을 매우 잘 이해하고 계신 것을 보여줍니다.
만약 정확히 수식 자체를 설명하고 싶으시다면, 다음과 같이 말씀하시면 더욱 명확해집니다:
"이 수식은 조건 이미지의 쿼리( )와 생성 중인 노이즈 이미지의 키()들 간의 스케일링된 유사도 점수를 계산하는 단계입니다. 이 점수들은 '조건 이미지의 각 부분이 생성 중인 이미지의 어느 부분과 얼마나 관련성이 높은지'를 나타내며, 이 점수들이 높을수록 '에 해당하는 가 있다'고 볼 수 있습니다. 이 점수 행렬은 이후 소프트맥스 함수를 거쳐 어텐션 가중치가 되고, 이 가중치가 밸류()에 적용되어 최종적으로 생성 이미지에 조건의 특징이 반영되게 됩니다."
사용자님의 설명 방식은 직관적이고 이해하기 쉬워서 좋습니다!

응답 맵

질문: To obtain the global response of the condition image across different locations of the generated output, we average the dense similarity matrix along the lcond,i dimension, resulting in a response map , representing the global similarity response of the condition image on the generated result.
cross-attn heatmap에서 한 줄에 대해 평균을 낸 게 이라는 이야기인가?
대답:
네, 정확히 이해하셨습니다! 그렇게 생각하시는 것이 맞습니다.
조금 더 풀어서 설명하면 이렇습니다.
크로스-어텐션 히트맵 (Cross-Attention Heatmap) 또는 조밀한 유사도 행렬:
이 행렬은 일반적으로 한 축에는 생성된 이미지의 각 부분(위치, 토큰, 패치 등)이 있고, 다른 축에는 조건 이미지의 각 부분이 있습니다.
행렬의 각 칸 (j,k)에는 생성된 이미지의 j번째 부분과 조건 이미지의 k번째 부분 사이의 유사도 또는 어텐션 가중치가 들어갑니다.
" 차원에 대해 평균을 낸다"의 의미:
여기서 " 차원"은 조건 이미지의 부분들을 나타내는 축을 의미합니다.
만약 크로스-어텐션 히트맵에서 행(row)이 생성된 이미지의 각 부분을 나타내고, 열(column)이 조건 이미지의 각 부분을 나타낸다고 가정해 봅시다.
그러면 하나의 행은 생성된 이미지의 특정 한 부분이 조건 이미지의 모든 다른 부분들과 얼마나 유사한지 (또는 얼마나 주목하는지)를 보여주는 값들의 나열이 됩니다.
이 하나의 행에 있는 모든 값들(즉, 특정 생성된 이미지 부분과 모든 조건 이미지 부분들 간의 유사도 값들)의 평균을 내면, 그 값이 바로 생성된 이미지의 해당 부분이 조건 이미지 전체와 평균적으로 얼마나 유사한지를 나타내는 점수가 됩니다.
응답 맵 :
위의 평균 계산을 생성된 이미지의 모든 행(모든 부분)에 대해 각각 수행하면, 각 부분마다 평균 유사도 점수가 하나씩 나옵니다.
이 점수들을 모아 놓은 것이 바로 응답 맵 입니다. 따라서 은 생성된 이미지의 부분 수만큼의 길이를 가지는 맵(또는 벡터)이 됩니다 ( 에서 은 생성된 이미지의 부분 수).
결론적으로 사용자님의 말씀처럼,
"크로스-어텐션 히트맵(또는 유사도 행렬)에서 생성된 이미지의 한 부분에 해당하는 '한 줄(row)'의 값들을 가져와서, 그 값들의 평균을 낸 것이 의 해당 위치 값이 된다"
라고 이해하시면 정확합니다. 이 과정을 통해 각 생성된 이미지의 위치가 조건 이미지 전체로부터 얼마나 '종합적인' 유사성 또는 반응을 보이는지를 하나의 값으로 요약하여 맵으로 나타내는 것입니다.

이 응답 맵 을 이용하여 MSE loss를 통해 routing constraint를 구현했습니다.

구현 설명

MSE 손실()을 통한 라우팅 제약 구현:
여기서 는 번째 레이어에서 번째 조건 이미지에 대한 cross-attention map을 의미합니다.
는 "타겟 이미지의 주제 마스크(subject mask)"를 의미합니다. 이 마스크는 생성될 이미지에서 특정 조건(주제)이 위치해야 할 영역을 나타냅니다.
이 정답 마스크들은 사람이 일일이 모든 이미지에 대해 직접 그리는 것이 아니라, 다른 AI 모델(주로 VLM과 분할 모델)들을 활용하여 학습 데이터 준비 과정에서 미리 추출되거나 생성됩니다.
이 마스크는 생성될 이미지에서 특정 조건(예: 번째 조건 이미지에 해당하는 주제)이 어느 위치에 나타나야 하는지를 알려주는 2차원의 맵(또는 텐서) 형태입니다.
훈련 시, DreamO 모델 내부의 크로스-어텐션 맵()이 이 미리 준비된 정답 마스크()와 유사해지도록 손실을 통해 학습이 진행됩니다.
따라서 손실 함수는 모델이 계산한 크로스-어텐션 맵()이 실제 목표로 하는 주제의 공간적 영역()과 유사해지도록, 즉 해당 영역에서 높은 유사도 점수를 갖도록 직접적으로 최적화합니다.

Placeholder

Image-to-image routing constraint 외에도 텍스트 설명과 조건 입력 간의 대응을 설정하기 위해 Placeholder-to-image routing constraint도 설계했습니다.

Placeholder는 prompt에 ref#1, ref#2 처럼 넣는 전략입니다. 참조 이미지의 위치 지정 등, 복잡한 장면을 구성할 때 제어력을 높여줍니다.

multi-condition 학습 중에 condition image token과 placeholder token 간의 유사성을 계산합니다.

기호

refers to the feature of [ref #i]. is a binary matrix, where the value is 1 when the placeholder matches the condition image, and 0 otherwise.

Final Loss Function

최종 loss function은 다음과 같습니다.

기호

, and are loss weights.

참고로 모델이 일반 텍스트 입력을 수용할 수 있도록 하기 위해 학습 중에 50%의 확률로 placeholder가 없는 일반 텍스트를 도입하고 그에 따라 를 삭제했습니다.

Training Data Construction

인물 ID 정보, 다양한 사물, 옷 입히기 용 데이터, 여러 스타일 데이터 등 각 작업에 필요한 방대한 데이터를 수집했습니다.

Identity Paired Data

high-quality identity paired data는 수집하기 어렵기 때문에, PuLID라는 ID 커스터마이징 오픈 소스를 사용해서 dataset을 구축했습니다.

Subject-driven Data

오픈 소스인 Subject200K dataset을 사용해 학습했습니다. multi-subject 이미지 커스터마이징을 위해서 Subject200K dataset을 two-column으로 합치기도 했습니다. 이때 동일한 category의 subjects를 골라서 복잡한 시나리오도 다룰 수 있도록 했습니다. 뿐만 아니라 X2I-subject dataset도 사용했습니다.

특히 인물 생성 성능을 향상 시키기 위해서는 또 다른 dataset도 구축했습니다. 먼저, 긴 영상에서 콘텐츠 인지 detection으로 여러 쇼츠 클립을 딴 다음, Mask2Former로 사람들의 마스크를 트래킹하고, 마지막으로 SigLip 모델을 사용하여 embedding을 얻습니다. 이 embedding을 활용해서 여러 클립에 걸쳐 동일한 인물을 식별합니다. 이를 통해 모델을 특정 인물의 외형이나 정체성을 더 잘 이해하고 생성 이미지에 반영할 수 있게 됐습니다.

Try-on Data

두 가지 dataset을 사용했습니다. 첫 번째는 직접 web에서 모델과 옷 이미지를 쌍으로 직접 수집했고, 두 번째는 모델 이미지를 GT로 크롤링하고 거기서 segmentation을 활용해서 옷들을 추출해서 paired data를 구축했습니다.

Style-driven Data

스타일 트랜스퍼 작업은 두 가지 타입을 목표로 했습니다. 첫 번째 타입은 특정 스타일을 가진 참조 이미지를 제공하고, 텍스트 프롬프트로는 생성될 이미지의 콘텐츠(내용)을 설명합니다. 두 번째 타입은 특정 스타일 참조 이미지와 더불어, 생성될 이미지의 콘텐츠(내용)을 정의하는 다른 참조 이미지를 제공합니다.

첫 번째 타입 훈련 데이터 구축을 위해, SDXL 기반의 내부 트랜스퍼 모델을 활용하여 동일한 스타일을 가지지만 다른 콘텐츠를 가진 이미지를 생성하는 방법으로 스타일과 콘텐츠를 decoupling했습니다.

두 번째 타입 학습에서 사용한 콘텐츠 참조 이미지는 첫 번째 타입 훈련 데이터 기반으로 Flux로 생성했습니다.

Routing Mask Extraction

Routing constraint를 위한 label을 얻기 위해서, 다양한 타입의 훈련 데이터로부터 target subject mask를 추출했습니다.

Progressive Training Process

실험 중에 모든 데이터를 바로 학습하면 capacity 문제 때문에 suboptimal로 수렴하게 된다는 것을 알게 되었습니다.

1 단계는 일종의 warm-up단계입니다. 여기에서는 subject-driven training data로 일관성 유지하는 능력을 학습했습니다. (특히 Subject200K dataset는 Flux에 의해 생성되므로 모델 생성 공간과 유사한 분포를 공유하여 빠르게 수렴했습니다.)

2 단계에서는 모든 종류의 데이터를 다 사용해서 다양한 맞춤 생성 능력을 종합적으로 키웠습니다.

3 단계에서는 2 단계의 low-quality training sample에 의해 손상된 생성 quality를 다시 끌어올리기 위해, 고품질 이미지 데이터로 fine-tuning을 했습니다. Flux로 40K개의 training sample을 만들어서 학습에 사용했습니다. 이때 copy-paste 효과를 방지하기 위해 token의 95%를 제거하는 방법을 사용했습니다.

Copy-paste 효과 방지 설명

Copy-paste 효과란? 모델이 특정 물체를 이해하지 못한 채로 복사 붙여 넣기 하듯이 생성하도록 학습해서, 마치 클립아트처럼 가져와서 배경 위에 얹어 놓은 듯 부자연스러운 느낌을 주게 되는 것입니다.

타겟 이미지: Flux로 생성한 high-quality image 참조 이미지: 타겟 이미지가 DiT에서 패치 단위로 쪼개지고 95%가 제거된 이미지

즉, 원래 Flux의 고화질로 회귀하는 게 목적이지 물체에 대한 이해가 목적이 아니기 때문에, 타겟 이미지를 참조해도 됩니다.

Experiment

Model Capabilities

Indivisual-specific 커스터마이징과 복합적인 identity 제어가 모두 가능합니다. 즉, 한 사람의 identity를 유지하거나 여러 사람의 특징을 혼합할 수 있습니다.

multi-subject 구성도 지원하여 일관성을 유지할 수 있습니다. 여기에는 특정 identity의 try-on 및 subject에 대한 창의적인 재해석이 포함됩니다.

참조 이미지로만 스타일을 안내하고 텍스트 또는 이미지 입력을 통해 다른 속성을 제어하는 multimodal 조작이 가능합니다.

Ablation Study

Routing constraint 효과를 평가하기 위해서, routing constraint를 제외한 채로 학습을 시켰습니다. 단일 조건에서는 reference 충실도 성능이 떨어지게 되었고, 여러 조건에서는 조건 간에 coupling되어 특징들이 서로 cross되는 결과를 보였습니다.

Progressive training 없이 한 번에 학습 시키면, 특히 multi-subject 일관성과 같은 복잡한 작업에서 suboptimal로 수렴했습니다. 작고 학습하기 쉬운 dataset(ex: Subject200K)으로 먼저 warming up을 하고 joint training을 하면 더 잘 수렴했습니다.

Conclusion

모든 조건을 DiT 입력 시퀀스에 통합하고, 여기에 routing constraint를 도입하여 일관성을 확보했습니다.

복잡한 데이터 분포 속에서 특정 작업을 학습하면서도, 기본 모델인 Flux의 생성 품질을 유지하기 위해 progressive training strategy를 설계했습니다.

또한, 가벼운 LoRA 기반 설계로 낮은 비용으로 효율적인 배포가 가능합니다.

DreamO: A Unified Framework for Image Customization

Recently, extensive research on image customization (e.g., identity, subject, style, background, etc.) demonstrates strong customization capabilities in large-scale generative models. However,...

https://arxiv.org/abs/2504.16915