[Python] Google Colab 에서 stable_diffusion 2 사용하기

안녕하세요.

오늘은 Google Colab 에서 huggingface의 Diffusers 라이브러리를 사용하여 stable diffusion 2를 간단하게 사용해보겠습니다.

출처 :

https://huggingface.co/stabilityai/stable-diffusion-2

 

1. huggingface 에서 필요한 stable_diffusion 모델을 다운받고, 추가적으로 필요한 라이브러리들을 다운로드 합니다.

!pip install --upgrade -qq git+https://github.com/huggingface/diffusers.git transformers accelerate scipy xformers

 

2.  Diffusiers 라이브러리에서 pipeline 과 scheduler 모듈을 가져오고

from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler

model_id = "stabilityai/stable-diffusion-2"

# Use the Euler scheduler here instead
scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, revision="fp16", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt, height=768, width=768).images[0]
    
image.save("astronaut_rides_horse.png")

image

어려운 인공지능을 다음과 같이 코드 몇줄로 간단한게 사용할 수 있습니다.

stable_diffusion 2 로 만든 이미지