[Part 1] Stable Diffusion는 어떻게 동작할까?( + [논문리뷰]High-resolution image synthesis with latent diffusion models)

ChatGPT가 최근 화두가 되어서 비교적 가려졌으나, 이미지 생성에 있어서도 또 하나의 기념비적인 사건이 일어났다. 바로 Stable Diffusion 모델이 등장해서 키워드 위주의 입력 프롬프트로 원하는 사진을 손쉽게 생성할 수 있게 된 것이다. 이런 Stable Diffusion이 파급력을 가진 이유는 높은 이미지 퀄리티와 상대적으로 적은 연산 리소스를 요구하는 하드웨어 사양 때문이다. 즉, 누구나 자신의 PC로 원하는 사진을 뽑을 수 있게 되었기 때문이다.

 

딥러닝 공부를 해본 사람들은 대부분 이미지 생성에 GAN(Generative Adversarial Network)을 사용한다는 것을 배웠을 것이다. 그래서 이런 Diffusion Model에 대해서는 생소할 수 있다. 그래서 이번 시간에는 이렇게 큰 인기를 끌고 있는 Stable Diffusion에 대해 조금 더 심층적으로 살펴보기 위해서 논문 리뷰 겸 Stable Diffusion 본격 파헤치기를 하고자 한다.

 

우선 이 논문은 필자도 방문했었던 CVPR2022에 개제되었던 논문이다. References를 제외하면 8페이지 분량이므로 적당한 길이다. 인용수도 1,000회가 넘었다(2023. 5. 4 기준). 이뿐만 아니라 Stable Diffusion의 대중성에 의해 많은 사람들이 이를 활용하여 원하는 이미지를 생성하고 있다. 그럼 왜 이 Stable Diffusion이 대중성을 갖게 되었는지 그 능력을 살펴보자.

바로 위와 같이 텍스트를 입력하면 이를 해석해서 그림으로 생성해주는 능력이 매우 강력하기 때문이다. 뿐만 아니라, 모델이 매우 경량화되어 일반 그래픽카드 수준의 VRAM으로도 이미지 생성이 가능하다는 점도 한몫했다. 위의 기능은 text to image 기능이지만, image to image나 inpaint와 같이 입력을 이미지로 할 수도 있다. 예를 들면 위의 불타는 말 그림에 추가 텍스트로 'Flying'을 붙이면 불타는 말이 하늘을 나는 그림이 나타난다.

 

그럼 이런 모델들을 사용하기 위해서 필요한 것은 무엇일까?

1. Stable Diffusion 기반의 모델(예: Stable Diffusion v2.0 등)

2. Stable Diffusion 실행을 손쉽게 해주는 인터페이스(예: AUTOMATIC1111의 stable-diffusion-webui 등)

3. (Optional) 로라, VAE, 임베딩, Pose 등의 부가 요소들(예: 특정 스타일을 입히기 위한 LoRA 등)

4. 입력 프롬프트(예: A running horse on fire)

대표적으로는 위의 4가지가 있으며, 부가적인 기능들은 하루가 멀다하고 개발되고 있다. 그리고 위의 4가지를 공유하는 대표적인 사이트들은 아래와 같다.

다양한 모델을 다운받을 수 있는 Civitai(https://civitai.com/)

Civitai와 함께 많이 사용되는 Hugging Face(https://huggingface.co/)

대표적으로 Civitai, Hugging Face가 있으며, 이 사이트들에서는 이미 방대한 데이터로 이미지 생성 모델을 학습하여 공유하고 있으므로 누구나 다운로드하여 원하는 이미지들을 생성할 수 있다. 추후에는 Stable Diffusion을 사용하는 상세 방법들과 꿀팁들을 하나씩 포스팅할 예정이다.

 

그럼 이제 Stable Diffusion의 구조와 작동 원리에 대해 알아보도록 하자. 먼저 큰 관점에서 모델을 살펴보고, 세부적인 모델 아키텍처를 공부해보자.

상위 레벨의 관점에서는 입력 프롬프트를 넣으면 그림이 출력되는 구조이다. 입력 프롬프트는 Encoder를 거쳐 임베딩되어 숫자들의 집합으로 바뀌게 된다. Encoder는 쉽게 말해 '텍스트를 모델이 이해하는 데이터로 바꿔주는 언어 번역기'라고 생각하면 쉽다. 텍스트 정보를 해석해서 이를 텍스트 내에 들어있는 아이디어들을 잘 캐치할 수 있는 숫자 표현형으로 바꿔주는 것이다. 좀 더 세부적으로 들어가면 이 텍스트 인코더는 사실 거대 언어 모델에 사용되는 Transformer다. 즉, 작은 형태의 ChatGPT가 사람의 언어를 입력받아 모델이 이해할 수 있는 숫자들로 바꿔주는 것이다. 이 입력 프롬프트(텍스트)를 해석하는 모델 역시 딥러닝 모델이며, 학습에 CLIP(Contrastive Language-Image Pre-Training) 방법이 사용되는데, 쉽게 말해 그림과 연관된 텍스트의 유사도는 극대화하고, 연관되지 않은 텍스트들과의 유사도는 최소화시키는 훈련을 하는 방법이다(딥러닝에서 이런 학습 프레임워크를 Contrastive Learning이라고 하는데, 이 역시 추후에 다루도록 하겠다). Stable Diffusion에서는 생성된 이미지가 주어진 텍스트와 잘 매칭되도록 하는 작업에 CLIP이 사용된다. Stable Diffusion을 사용해 보신 분들은 입력 프롬프트 작성 창에 'Clip'이라는 단어를 보셔서 익숙하실 것이다. 이때 사용 가능한 단어 각각을 Token이라고 부르는데, 77개의 Token까지 입력할 수 있다. 예를 들어 '아름다운, 불타는, 말'이라고 작성한다면 3개의 Token을 쓴 것이다. 이 각각의 Token은 768차원의 임베딩 벡터로 변환된다.

 

이미지 생성기는 입력받은 숫자들의 집합(이를 임베딩 벡터라고도 부른다)을 받아 그림을 출력한다. 이 역시 2단계로 나눌 수 있다.

 

우선 첫 번째로 Image Information Creator를 살펴보자.

Stable Diffusion에서 Scheduler 및 Step을 설정하는 부분이 바로 이 부분인데, 이는 인코더를 통해 들어온 임베딩 벡터를 조건으로 하여 이미지 정보를 생성하는 과정이다. 재밌는 점은 이미지 정보를 생성할 때, Pixel Space가 아닌, Latent Space에서 동작한다는 점이다. Latent Space를 이해하기 위해서는 VAE에 대해서 잘 알아야 하는데, 다음 포스트에서 다룰 예정이므로 우선은 다음 설명으로 쉽게 이해하고 넘어가자.

만약 어떤 1920x1080 크기의 고화질 사진 속에 초록색 박스가 찍혔다고 가정해 보자. 이 초록색 박스를 포토샵 처리 해서 빨간색으로 바꾸고 싶다. 그리고 만약 이 고화질 사진을 960x540으로 압축하고, 손실 없이 다시 1920x1080으로 키울 수 있는 방법이 있다고 하자.

그럼 원래의 고화질 이미지를 먼저 압축하고, 압축된 이미지를 빨간색으로 바꾼 뒤, 이를 다시 고화질 그림으로 복원하는 것이 큰 이미지 색을 다시 칠하는 것보다 훨씬 쉬울 것이다(위는 단순한 예시였으나, 실제로 이미지 프로세싱 작업이 굉장히 복잡한 경우를 생각해 보면, 작은 이미지가 처리하기 더 쉽다는 것을 이해할 수 있다). 그런데 만약 위의 고화질 사진을 960x540 크기의 작은 이미지가 아니라, 아예 1x256 크기의 벡터로 더 압축해 버릴 수는 없을까? 그리고 이 256개의 벡터 값들 중 색상을 나타내는 값을 바꿔버리는 것이다.

위와 같이 표현된 더 작은 차원을 갖는 벡터로 표현된 공간을 Latent Space라고 부른다. 즉, 원래 데이터를 응축시킨 데이터로 표현한 데이터라고 볼 수 있다. Stable Diffusion에서는 이미지 생성을 위한 정보 생성에 Pixel Space에서 고화질 이미지 정보를 생성하는 것이 아니라, 바로 위의 Latent Space에서 해당 이미지를 잘 나타내는 핵심 정보를 생성한다. 따라서 당연히 이미지 생성 속도를 크게 단축시킬 수 있었고, 이 방법이 핵심적인 아이디어 중 하나다. 이 이미지 정보 생성 과정은 이전 포스트에서 설명한 U-Net 신경망으로 구성되어 있다. 많은 분들이 "Diffusion"이라는 용어가 왜 사용되었는지 궁금하실 텐데, 이 U-Net 신경망이 이미지에 담긴 노이즈를 점차 '확산'시켜 제거하는 과정을 Step 수만큼 반복하기 때문에 이름이 붙여졌다. 어떻게 이미지에 담긴 노이즈를 점차 걷어내는지에 대해서는 이후에 계속 설명하기로 하고, 우선은 각 과정 사이에서 전달되는 정보를 상기해 보자.

먼저 입력 프롬프트를 작성하면 77 Tokens이 인코더로 들어간다(최대 토큰 수는 프레임워크, 모델에 따라 달라질 수도 있다). 이는 텍스트 인코더를 통해 숫자들의 집합인 Word Embeddings이 되어 768차원의 임베딩 벡터 77개가 된다. 이는 Image Information Creator에서 여러 번 반복 처리되며 노이즈가 걷히게 되고, 그 결과 Latent Space에서 응축된 정보 형태로 4x64x64의 Tensor로 출력된다. 그리고 마지막으로 Image Decoder로 들어가 Pixel Space에서 3x512x512의 그림으로 탄생하는 것이다(RGB 채널이므로 3이고, 그림 크기가가 가로x세로 512x512다).

마지막 부분에 속하는 Image Decoder는 Latent Space에서 표현된 4x64x64의 이미지 정보로부터 그림을 그려내는 역할을 한다. 앞의 이미지 정보가 반복적 과정을 거쳐 구해진 것과 다르게, 이 과정은 마지막에 한 번 수행되어 Pixel Space에서 이미지를 그려낸다. 이 Image Decoder는 Autoencoder의 Decoder 구조로, Autoencoder는 앞서 설명한 이미지를 압축-복원하는 역할을 한다고 생각하면 쉽다(U-net에서 Skip connection을 떼버린 구조라고 생각하면 된다. 레이어가 점차 좁혀지므로, 처음 고차원 정보가 저 차원 정보로 압축되었다가, 다시 고차원 정보로 복원되도록 한다. 여기선 저차원 정보를 고차원으로 복원하는 디코더 부분에 해당되며, 자세한 내용은 VAE 포스트에서 VAE, AE에 대해 다루도록 하겠다).

 

여기까지 이해했다면 기본적인 Stable Diffusion 아키텍처에 대해서 이해한 것이다. 이제 논문에서의 실제 아키텍처를 살펴보자.

우리가 입력한 입력 프롬프트(텍스트)는 오른쪽에서 워드 임베딩이 되어 들어온다. 이후 Image Information Creator로 들어가고(Transofrmer이므로, Attention Mechanism이 적용되어 Query, Key, Value가 적혀있는 것도 볼 수 있다), T번 반복되며 노이즈가 점차 확산되며 사라진다. 마지막에 z는 'Latent Vector'를 의미하는데, 이게 앞서 배운 응축된 이미지 정보다. 이 4x64x64 크기의 Tensor로 표현되는 Latent Vector를 Image Decoder에서 받아 Pixel Space에서 VAE의 디코더 구조로 다시 3x512x512 고화질 이미지로 복원하게 되면, 이것이 우리가 원하는 생성된 이미지다.

 

그런데 의문이 든다. 왼쪽 위의 x는 무엇일까? 또한, 'Diffusion Process'는 도대체 어떤 과정이며, 어떻게 학습시켜야 이렇게 노이즈를 점차 확산시켜가며 이미지를 생성할 수 있을까? 앞서 우리는 Stable Diffusion의 큰 아키텍처 흐름을 공부했지만, 이 확산 과정에 대해서는 아직 배우지 않았다. 다음 포스트에서는 이 확산 과정과, 모델 학습 과정에 대해서 살펴보고, 추가로 트랜스포머 계열인 텍스트 인코더의 CLIP에 대해서도 좀 더 자세히 살펴봄으로써, 입력 프롬프트가 그림에 반영되는 원리도 더 자세히 공부해 보자.