pianoprincesjy/VAEfinetuned

from transformers import AutoTokenizer
from diffusers import StableDiffusionPipeline
import torch
from IPython.display import display  # display 함수 불러오기

# Stable Diffusion 2.1 모델을 불러옵니다
model_id = "stabilityai/stable-diffusion-2-1"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cuda")  # GPU 사용 설정

# 저장된 VAE 가중치 불러오기
vae_path = "pytorch_model.bin"  # 파인튜닝한 VAE 가중치 파일 경로
vae_weights = torch.load(vae_path)  # VAE 가중치 불러오기

# 파이프라인의 VAE 모델에 가중치를 덮어씌우기
pipe.vae.load_state_dict(vae_weights)  # 파이프라인의 VAE에 가중치 적용
pipe.vae.to("cuda")  # GPU로 이동

# 프롬프트 설정 및 이미지 생성
prompt = "a vector illustration of a man"
image = pipe(prompt).images[0]

# Jupyter Notebook 상에서 이미지를 바로 표시
display(image)