retooled to use old hugginface pipline instead

app.py

@@ -1,35 +1,49 @@
 # Python
 import streamlit as st
-from PIL import Image
-from openai import OpenAI
-
-
-# Function to load and display image
-def display_image(image_path):
-    image = Image.open(image_path)
-    st.image(image, use_column_width=True)
-
-st.title('Generate Your Image!')
-prompt = st.text_input('Enter a prompt')
-res = []
-
-if prompt: # prompt dalle
-    client = OpenAI()
-    response = client.images.generate(
-        model="dall-e-3",
-        prompt=prompt,
-        size="1024x1024",
-        quality="standard",
-        n=1,
+import torch
+from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
+
+def load_pipelines(device, dtype):
+    prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=dtype).to(device)
+    decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade",  torch_dtype=dtype).to(device)
+    return prior, decoder
+
+def generate_images(prompt, negative_prompt, num_images_per_prompt, device, dtype):
+    with torch.cuda.amp.autocast(dtype=dtype):
+        prior_output = prior(
+            prompt=prompt,
+            height=1024,
+            width=1024,
+            negative_prompt=negative_prompt,
+            guidance_scale=4.0,
+            num_images_per_prompt=num_images_per_prompt,
         )
-    st.write(response)
-    
-    
+        decoder_output = decoder(
+            image_embeddings=prior_output.image_embeddings,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=0.0,
+            output_type="pil",
+        )
+    return decoder_output.images
+
+
+st.title("Image Generator with Diffusers")
+
+@st.cache(allow_output_mutation=True)
+def init_model():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16
+    return load_pipelines(device, dtype), device, dtype
+
+(prior, decoder), device, dtype = init_model()
+
+prompt = st.text_input("Enter a prompt:", "Anthropomorphic cat dressed as a pilot")
+negative_prompt = st.text_input("Enter a negative prompt:", "")
+num_images_per_prompt = st.slider("Number of images per prompt:", 1, 5, 2)
 
-# if images: #display images
-    
-#     image_path1 = 'path_to_image1_based_on_prompt'  # Replace with actual logic
-#     image_path2 = 'path_to_image2_based_on_prompt'  # Replace with actual logic
 
-#     display_image(image_path1)
-#     display_image(image_path2)
+if st.button("Generate"):
+    images = generate_images(prompt, negative_prompt, num_images_per_prompt, device, dtype)
+    for img in images:
+        st.image(img, use_column_width=True)

requirements.txt

@@ -1 +1,3 @@
-openai==1.12.0
+streamlit
+torch
+git+https://github.com/kashif/diffusers.git@wuerstchen-v3