Spaces:
Sleeping
Sleeping
import os | |
import torch | |
import gc | |
import shutil | |
import streamlit as st | |
from PIL import Image | |
from diffusers import DiffusionPipeline, AutoencoderKL | |
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) | |
pipe = DiffusionPipeline.from_pretrained( | |
"stabilityai/stable-diffusion-xl-base-1.0", | |
vae=vae, torch_dtype=torch.float16, variant="fp16", | |
use_safetensors=True | |
) | |
pipe.load_lora_weights("./lora-trained-thumbs-up-pw-0.5-steps-800") | |
pipe.fuse_lora(lora_scale=1.0) | |
pipe.save_pretrained("temp_model") | |
del pipe | |
gc.collect() | |
torch.cuda.empty_cache() | |
pipe = DiffusionPipeline.from_pretrained( | |
"temp_model", | |
vae=vae, torch_dtype=torch.float16, variant="fp16", | |
use_safetensors=True | |
) | |
pipe.load_lora_weights("./lora-trained-thumbs-up-pw-0.5-steps-800-uttam-pw-0.5-steps-1600") | |
_ = pipe.to("cuda") | |
prompt = st.text_area("Enter the prompt!") | |
image = pipe(prompt=prompt, num_inference_steps=25, num_images_per_prompt=1)[:] | |
shutil.rmtree("temp_model", ignore_errors=True) | |
st.image(image) |