[Text-to-Image Generation] ControlNet

AI가 만들어준 그림으로 미술대회에서 1등을 수상하다

혹시 여러분께서는 작년에 뉴스 기사를 읽거나 유튜브를 보던 중 아래 그림을 본 기억이 있으신가요? 이 신비한 그림은 2022년 8월 29일 미국 콜로라도 주립박람회 미술대회 디지털 아트 부문에서 1등을 차지한 그림입니다. 하지만 그림을 제출한 Jason Allen은 수상 발표를 통해 직접 그림을 그린 것이 아닌 텍스트로 이미지를 생성하는 미드저니(Midjourney) AI 모델로 그림을 생성했다고 밝히며 많은 논란이 있었습니다.

 

사진 1. 미국 콜로라도 주립박람회 미술대회 디지털 아트 부문 1등을 차지한 'Space Opera Theater' [1]

사실 완벽해보이는 위 그림은 단순히 미드저니 AI 모델을 한 번 사용한 것으로 만들어진 사진은 아닙니다. Jason Allen은 원하는 이미지를 생성하기 위한 적절한 텍스트 프롬프트를 찾는데 약 80시간 정도의 시간을 할애했다고 전했습니다. 또 포토샵을 통한 리터칭과 다른 AI 모델을 활용하여 수정 작업을 거쳤다고 하며, 아직까지는 사용자가 정확히 원하는 이미지를 생성하려면 많은 시간과 노력이 필요하다고 볼 수 있습니다.

구조적인 가이드를 제공하여 이미지를 생성하는 ControlNet

그래서 최근 이 분야에서는 사용자가 원하는 이미지를 생성할 수 있도록, 텍스트와 별개의 새로운 조건을 추가하는 확장 모듈에 대한 연구가 진행되고 있습니다. 이러한 새로운 조건에는 Canny edge, Sketch, Segmentation, Depth map, Color, Style 등이 있습니다. 이번 글에서는 미드저니와 유사하게 텍스트를 통해 이미지를 생성하면서 성능이 좋고 무료로 공개되어 있는 모델인 스테이블 디퓨전(Stable Diffusion)에 확장 모듈 연구인 ControlNet을 소개하려고 합니다. 아래는 Canny edge를 새로운 조건으로 추가하는 ControlNet 모듈을 스테이블 디퓨전과 함께 사용했을 때 생성된 이미지 결과입니다.

 

사진 2. 원본 이미지에서 추출한 Canny edge를 사용하여 생성한 이미지 [2]

새로운 조건을 사용하여 이미지를 생성하는 방법은 간단합니다. 먼저 왼쪽 위의 원본 이미지(Source image)로 부터 Canny edge detection 알고리즘을 통해 왼쪽 아래의 Canny edge를 추출합니다. 추출된 Canny edge는 미리 학습된 ControlNet의 입력으로 전달되어 생성할 이미지의 구조에 대한 가이드(Structure Guidance)를 스테이블 디퓨전에 제공합니다. 마지막으로 스테이블 디퓨전은 이미지의 구조에 대한 가이드를 바탕으로 결과 이미지를 생성하게 됩니다.

 

그럼 ControlNet의 구성을 자세히 살펴보겠습니다. 아래 사진에서 (a)는 Stable Diffusion의 U-Net block이며, (b)에서 새로운 조건인 c를 입력으로 받는 파란색 점선 박스가 ControlNet 입니다. "trainable copy"라고 명시된 파란색 박스는 바로 왼쪽의 U-Net block과 동일한 구조를 사용하며 이미 훈련된 왼쪽의 파라미터를 그대로 복사하여 ControlNet 학습에 사용합니다. 또 안정적인 초기 학습을 위해 zero convolution을 추가하여 기존 레이어로 더해지는 ControlNet의 출력값이 0에서 시작하여 점차 학습되도록 합니다.

 

사진 3. ControlNet [2]

다음으로 ControlNet이 적용된 Stable Diffusion의 구조를 살펴보겠습니다. 아래 사진에서 왼쪽은 기존 Stable Diffusion 이며, 오른쪽은 새로 추가된 ControlNet 입니다. 새로운 조건을 학습할 때 ControlNet의 파라미터는 항상 학습되지만, Stable Diffusion U-Net의 파라미터는 오버피팅을 고려하여 데이터셋의 크기가 1M 이상인 경우에만 함께 학습을 진행합니다. 반대로 1M 보다 작은 경우에는 학습을 하지않고 Freeze하여 사용합니다.

 

사진 4. Stable Diffusion with ControlNet [2]

학습에 사용하는 데이터셋은 어떻게 구성하였을까요? 위 사진에서 Stable Diffusion, ControlNet의 입력을 보면 이미지-텍스트 페어와 함께 우리가 학습하려고 하는 컨디션 데이터까지 필요하다는 것을 알 수 있습니다. 하지만 이미지와 이를 설명하는 텍스트, 그리고 학습하려는 컨디션이 같이 들어있는 데이터셋을 구하기는 쉽지 않습니다. 그래서 본 논문의 저자는 이미지-컨디션이 있는 데이터셋에서 BLIP이라는 별도 모델로 이미지를 설명하는 텍스트를 생성하여 데이터셋으로 사용했다고 합니다. BLIP으로 텍스트를 생성하는 방법은 좌측의 이미지 인코더로 추출한 이미지 임베딩을 우측 텍스트 디코더의 크로스 어텐션(Cross Attention)에 전달합니다. BLIP-1 모델은 아래 사진과 같으며 현재는 BLIP-2 모델까지 공개 되어있습니다.

 

사진 5. BLIP Captioning [3]

이제 결과를 살펴보겠습니다. 이미지를 생성할 때 사용한 텍스트 종류는 Default Prompt, Automatic Prompt, User Prompt로 총 세 가지가 있습니다. Default Prompt는 더 좋은 품질의 이미지를 생성하기 위해 "a professional, detailed, high-quality image"를 사용하였고, Automatic Prompt는 원본 이미지를 입력으로 BLIP 모델을 통해 생성한 텍스트에 Default Prompt를 추가하여 사용했다고 합니다. 마지막으로 User Prompt는 저자가 이미지를 보고 직접 작성한 텍스트를 사용했습니다. 어떤 텍스트를 사용하더라도 결과 이미지의 품질이 준수한 것을 확인 할 수 있습니다. 추가적인 예시는 논문 또는 깃허브를 참고해주시기 바랍니다.

 

사진 6. Canny edge를 사용하여 생성한 이미지 [2]

사진 7. Hough line을 사용하여 생성한 이미지 [2]

사진 8. Cartoon line drawing을 사용하여 생성한 이미지 [2]

앞으로의 연구 방향

현재 공개된 연구 중 새로운 조건을 추가하여 이미지 생성 시 구조적인 가이드를 제공하는 모델은 ControlNet 외에 T2I-Adapter, CoAdapter, Composer가 있습니다. 이 중 ControlNet, T2I-Adapter, CoAdapter는 모델 코드와 학습 파라미터가 공개되어있고, 허깅페이스(Huggingface)의 웹 데모를 간단히 테스트를 해볼 수도 있습니다.

 

우리가 이 분야의 연구에서 한 가지 더 주목해볼만한 점은 이미지 생성 외에도 Style Transfer, Inpainting 등 여러 활용처가 있습니다. 지금까지 연구가 활발히 진행되고 있기 때문에, 앞으로의 모델은 어떤 것까지 할 수 있게 될 지 기대가 됩니다.

참고자료

[1] AI Art Wins Competition And Sparks Controversy

https://medium.com/mlearning-ai/ai-art-wins-fine-arts-competition-and-sparks-controversy-882f9b4df98c

AI Art Wins Competition And Sparks Controversy

[2] ControlNet: Adding Conditional Control to Text-to-Image Diffusion Models

https://arxiv.org/abs/2302.05543

Adding Conditional Control to Text-to-Image Diffusion Models

[3] BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation

https://arxiv.org/abs/2201.12086

BLIP: Bootstrapping Language-Image Pre-training for Unified Vision-Language Understanding and Generation