[토이 프로젝트] Stable Diffusion 텍스트로 원하는 이미지 생성해보기(초간단 colab)

stable diffusion은 텍스트로 이미지를 생성하는 모델입니다. 최근 gan을 능가하는 성능을 보이고 있으며, pretrained model이 오픈되어 많은 사람들이 활용해볼 수 있는데요. 토이 프로젝트를 진행하기 앞서 간단하게 테스트를 해보았습니다.

 

 

colab 링크를 클릭하셔서 몇가지 설정을 하시고 명령문을 하나하나 실행하면서 이미지를 생성해봅니다.

 

먼저, 런타임 유형을 gpu로 바꿔줍니다.

 

 

colab 파일을 새로 만들고 다음 명령어를 입력하여 라이브러리를 설치해줍니다

pip install --upgrade diffusers transformers scipy

huggingface에 가입을 한 후에 settings를 클릭한 후 token을 복사합니다.

 

 

허깅페이스 로그인을 위해 위에서 복사한 token을 입력합니다.

!huggingface-cli login

 

Model Card를 클릭하면 이 모델을 사용하기 위한 동의서가 나오는데 저는 이미 동의를 수락한 상태입니다. 동의를 수락하면 다음과 같은 상태가 될 것입니다.

 

모델을 불러오는 부분이 조금 오래 걸리고 prompt를 입력하면 해당하는 그림이 저장됩니다.

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"


pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True)
pipe = pipe.to(device)

prompt = "a photo of an astronaut riding a horse on mars"
with autocast("cuda"):
    image = pipe(prompt, guidance_scale=7.5).images[0]  
    
image.save("astronaut_rides_horse.png")

 

해당하는 이미지를 출력해줍니다.

import matplotlib.pyplot as plt
plt.imshow(image)