Update app.py

app.py

@@ -1,14 +1,57 @@
 import streamlit as st
-from PIL import Image
+import cv2 as cv
+import time
+import streamlit as st
 import torch
 from diffusers import StableDiffusionPipeline
-import cv2 as cv
+import os
+import openai
+
+# "stabilityai/stable-diffusion-2-1-base"
+# "CompVis/stable-diffusion-v1-4"
+def create_model(loc = "stabilityai/stable-diffusion-2-1-base", mch = 'cuda'):
+    pipe = StableDiffusionPipeline.from_pretrained(loc, torch_dtype=torch.float16)
+    pipe = pipe.to(mch)
+    return pipe
+
+
+
+openai.api_key = "please-paste-your-API-key-here"
+
+def chatWithGPT(prompt):
+  completion = openai.ChatCompletion.create(
+  model="gpt-3.5-turbo",
+  messages=[
+  {"role": "expert in creating prompts for stable diffusion", "content": prompt}
+  ]
+  )
+  return print(completion.choices[0].message.content)
+
+
+
+t2i = st.checkbox("Text2Image")
+
+if t2i:
+    st.title("Text2Image")
+    t2m_mod = create_model()
+    
+    prom = st.text_input("# Prompt",'')
 
-st.title("Txt2Img")
-pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16")
-pipe.to("cpu")
-prom = st.text_area("Prompt",'')
-if len(prom)>5:
-    image = pipe(prom).images[0]
-    st.image(image)
+    c1,c2,c3 = st.columns([1,1,3])
+    c4,c5 = st.columns(2)
+    with c1:
+      bu_1 = st.text_input("Seed",'999')
+    with c2:
+      bu_2 = st.text_input("Steps",'12')
+    with c3:
+      bu_3 = st.text_input("Number of Images",'1')
+    with c4:
+      sl_1 = st.slider("Width",256,1024,128)
+    with c5:
+      sl_2 = st.slider("hight",256,1024,128)
 
+    create = st.button("Imagine")
+    if create:
+        generator = torch.Generator("cuda").manual_seed(int(bu_1))
+        img = t2m_mod(prom, width=int(sl_1), height=int(sl_2), num_inference_steps=int(bu_2), generator=generator).images[0]
+        st.image(img)