[Project] 이미지 생성 모델 만들기

모델 호출 및 이미지 생성 코드 기록

venv 생성하는 법도 모르고 헤매고, image_processor 에러 때문에 버벅이던 어제보다는 오늘이 매우 뿌듯하다. 지금부터가 진짜 시작이다. 가장 기본적인 text to image 모델 부터 API 연결하고 이후에 모델을 어떻게 지속적으로 training 할지 계획을 세워보자.

import os
from diffusers import DiffusionPipeline
import torch

# load both base & refiner
base = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)

# mps 는 Metal performers shaders 의 약자로 Apple silicon Mac 은 이걸로 해야함. 그 외는 cuda
base.to("mps")
refiner = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=base.vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
)
refiner.to("mps")

# Define how many steps and what % of steps to be run on each experts (80/20) here
n_steps = 30
high_noise_frac = 0.8

prompt = "어쩌구 저쩌구"

# run both experts
image = base(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_end=high_noise_frac,
    output_type="latent",
).images
image = refiner(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_start=high_noise_frac,
    image=image,
).images[0]

# 이미지 경로와 기본 파일명 설정
image_dir = os.path.expanduser("~/경로")
base_filename = "image.png"

# 파일명에 추가할 숫자 초기화
num = 0

# 최종 파일명 생성
final_path = os.path.join(image_dir, base_filename)

# 파일이 이미 존재하는 경우, 새로운 파일명 생성
while os.path.exists(final_path):
    num += 1
    # 파일명과 확장자 분리
    name, ext = os.path.splitext(base_filename)
    # 새로운 파일명 생성
    new_filename = f"{name}({num}){ext}"
    final_path = os.path.join(image_dir, new_filename)

# 최종 경로에 이미지 저장
image.save(final_path)
print(f"Image saved to {final_path}")