Spaces:
Runtime error
Runtime error
import diffusers | |
import torch | |
import os | |
import time | |
import accelerate | |
import streamlit as st | |
from stqdm import stqdm | |
from diffusers import DiffusionPipeline, UNet2DConditionModel | |
from PIL import Image | |
MODEL_REPO = 'OFA-Sys/small-stable-diffusion-v0' | |
LoRa_DIR = 'weights' | |
DATASET_REPO = 'VESSL/Bored_Ape_NFT_text' | |
SAMPLE_IMAGE = 'weights/Sample.png' | |
def load_pipeline_w_lora() : | |
# Load pretrained unet from huggingface | |
unet = UNet2DConditionModel.from_pretrained( | |
MODEL_REPO, | |
subfolder="unet", | |
revision=None | |
) | |
# Load pipeline | |
pipeline = DiffusionPipeline.from_pretrained( | |
MODEL_REPO, | |
unet=unet, | |
revision=None, | |
torch_dtype=torch.float32, | |
) | |
# Load LoRa attn layer weights to unet attn layers | |
print('LoRa layers loading...') | |
pipeline.unet.load_attn_procs(LoRa_DIR) | |
print('LoRa layers loaded') | |
pipeline.set_progress_bar_config(disable=True) | |
return pipeline | |
def elapsed_time(fn, *args): | |
start = time.time() | |
output = fn(*args) | |
end = time.time() | |
elapsed = f'{end - start:.2f}' | |
return elapsed, output | |
st.title("BAYC Text to IMAGE generator") | |
st.write(f"Stable diffusion model is fine-tuned by lora using dataset {DATASET_REPO}") | |
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") | |
st.write("Loading models...") | |
elapsed, pipeline = elapsed_time(load_pipeline_w_lora) | |
st.write(f"Model is loaded in {elapsed} seconds!") | |
pipeline = pipeline.to(device) | |
sample = Image.open(SAMPLE_IMAGE) | |
st.image(sample, caption="Example image with prompt <An ape with solid gold fur and beanie>") | |
with st.form(key="information", clear_on_submit=True): | |
prompt = st.text_input( | |
label="Write prompt to generate your unique BAYC image! (e.g. An ape with golden fur)") | |
num_images = st.number_input(label="Number of images to generate", min_value=1, max_value=10) | |
seed = st.number_input(label="Seed for images", min_value=1, max_value=10000) | |
submitted = st.form_submit_button(label="Submit") | |
if submitted : | |
st.write(f"Generating {num_images} BAYC image with prompt <{prompt}>...") | |
generator = torch.Generator(device=device).manual_seed(seed) | |
images = [] | |
for img_idx in stqdm(range(num_images)): | |
generated_image = pipeline(prompt, num_inference_steps=30, generator=generator).images[0] | |
images.append(generated_image) | |
st.write("Done!") | |
st.image(images, width=150, caption=f"Generated Images with <{prompt}>") | |