파이썬 diffusers 모듈로 stable diffusion 사용해보기 예제

안녕하세요.

이번 시간에서는 Python에서 diffusers 모듈을 이용하여

허깅페이스에서 스테이블 디퓨전 모델을 쉽게 불러와서

이미지를 원하는대로 생성해보는 예제를 다루어보도록 하겠습니다.

 

 

1. 모듈 설치

먼저, 필요한 모듈을 다음 명령어를 통하여 설치해보도록 하겠습니다.

!pip install diffusers transformers accelerate torch

 

 

2. 모델 불러오기 및 이미지 생성

모듈 설치 후, 아래의 코드 스니펫으로 허깅페이스에서 모델을 불러와

이미지 생성을 진행해볼 수 있습니다.

(초기 실행 시에는 모델 다운로드에 수 분 정도가 소요될 수 있습니다.)

import torch
from diffusers import AutoPipelineForText2Image

# 1. 모델과 파이프라인 불러오기 (GPU 사용)
# Stable Diffusion 1.5 버전 사용 예시
pipe = AutoPipelineForText2Image.from_pretrained(
    "sd-legacy/stable-diffusion-v1-5",
    torch_dtype=torch.float16
).to("cuda")

# 2. 그리고 싶은 그림을 텍스트로 설명하기 (프롬프트), 단 해당 모델은 영어로 작성해야 함
prompt = "A dog driving a sports car"

# 3. 이미지 생성하기
image = pipe(prompt=prompt).images[0]

# 4. 생성된 이미지 저장하기
image.save("my_image.png")

 

스포츠 카를 운전하는 강아지 이미지를 위 코드로 생성한 예시는 다음과 같습니다.

 

 

간혹, AutoPipelineForText2Image 로 파이프라인을 불러올 수 없는 예시가 있는데요.

이런 경우는 허깅페이스 공식 페이지 내에서 해당 모델을 검색하여

알맞는 파이프라인을 통하여 모델을 불러와주시면 마찬가지로 사용이 가능합니다.

 

이 글이 디퓨전 모델을 파이썬 내에서 직접 활용하는 과정에 도움이 되셨기를 기원합니다.

감사합니다.