Upload 3 files

app.py

@@ -0,0 +1,83 @@
+import diffusers
+import torch
+import os
+import time
+
+import streamlit as st
+
+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 LoRa attn layer weights to unet attn layers
+    unet.load_attn_procs(LoRa_DIR)
+
+    # Load pipeline
+    pipeline = DiffusionPipeline.from_pretrained(
+        MODEL_REPO,
+        unet=unet,
+        revision=None,
+        torch_dtype=torch.float32,
+    )
+
+    return pipeline
+
+
+def elapsed_time(fn, *args):
+    start = time.time()
+    output = fn(*args)
+    end = time.time()
+
+    elapsed = f'{end - start:.2f}'
+
+    return elapsed, output
+
+
+def main():
+
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+    st.title("BAYC Text to IMAGE generator")
+    st.write(f"Stable diffusion model is fine-tuned by lora using dataset {DATASET_REPO}")
+
+    sample = Image.open(SAMPLE_IMAGE)
+    st.image(sample, caption="An ape with solid gold fur and beanie")
+
+    elapsed, pipeline = elapsed_time(load_pipeline_w_lora)
+    st.write(f"Model is loaded in {elapsed} seconds!")
+
+    prompt = st.text_input(
+        label="Write prompt to generate your unique BAYC image! (e.g. An ape with golden fur)")
+
+    num_images = st.slider("Number of images to generate", 1, 10, 1)
+
+    seed = st.slider("Seed for images", 1, 10000, 1)
+
+    if prompt and num_images and seed:
+        st.write(f"Generating {num_images}BAYC image with prompt {prompt}...")
+
+        generator = torch.Generator(device=device).manual_seed(seed)
+        images = []
+        for img_idx in 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}")
+
+if __name__ == '__main__':
+    main()

weights/pytorch_lora_weights.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4927a8337a888edff7d844cac082c141726ac60fe8ad065c9143ec0987a9297
+size 1080571