[Stable Diffusion모델] 활용한 스크립트를 입력하면 이미지로 변경되는 파이썬 코드

 
Stable Diffusion 오픈소스를 활용해, 
 
스크립트를 입력해 이미지를 만드는 Text to Image 파이썬 코드를 공유한다
 
 

import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk
from diffusers import StableDiffusionPipeline
import torch

# Stable Diffusion 모델 로드
model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(model_id)
pipe = pipe.to(device)

# 이미지 생성 함수
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# 이미지를 UI에 표시하는 함수
def display_image():
    prompt = entry.get()  # 사용자가 입력한 텍스트 가져오기
    if not prompt:
        return
    
    # 이미지 생성 및 표시
    img = generate_image(prompt)
    img.thumbnail((400, 400))  # 이미지 크기 조정
    
    # Pillow 이미지 -> Tkinter에서 사용 가능한 포맷으로 변환
    img_tk = ImageTk.PhotoImage(img)
    panel.config(image=img_tk)
    panel.image = img_tk  # 참조를 유지해야 이미지가 표시됨

# Tkinter 윈도우 설정
root = tk.Tk()
root.title("Stable Diffusion Image Generator")

# 스크립트 입력 필드
entry_label = tk.Label(root, text="Enter script:")
entry_label.pack()

entry = tk.Entry(root, width=50)
entry.pack()

# 이미지 표시 패널
panel = tk.Label(root)
panel.pack()

# 이미지 생성 버튼
generate_button = tk.Button(root, text="Generate Image", command=display_image)
generate_button.pack()

# Tkinter 윈도우 실행
root.mainloop()

 
실행하면 아래와 같은 창이 생성되고

어떤 이미지를 만들 것인지 스크립트를 집어 넣는다

귀엽지 않은 다람쥐가 나왔다. 오픈소스의 한계인 듯

결과는 시중에 있는 플랫폼(Runway, Haiper AI)를 활용하자. 상당헤 퀄리티가 낮다
이미지가 생성되는 속도도 10분가량 소요된다