Update app.py

app.py

@@ -2,6 +2,9 @@ import streamlit as st
 from g4f.client import Client
 import sqlite3
 import google.generativeai as genai
+from diffusers import DiffusionPipeline
+import matplotlib.pyplot as plt
+import torch
 # import pyttsx3
 # import pyperclip
 
@@ -22,6 +25,24 @@ try:
     conn.commit()
 except Exception as e:
     st.error(f"An error occurred: {e}")
+    
+
+def generate_image(pipe, prompt, params):
+  img = pipe(prompt, **params).images
+
+  num_images = len(img)
+  if num_images>1:
+    fig, ax = plt.subplots(nrows=1, ncols=num_images)
+    for i in range(num_images):
+      ax[i].imshow(img[i]);
+      ax[i].axis('off');
+
+  else:
+    fig = plt.figure()
+    plt.imshow(img[0]);
+    plt.axis('off');
+  plt.tight_layout()
+    
 
 # Streamlit app
 def main():
@@ -35,7 +56,8 @@ def main():
         models = {
             "🚀 Airoboros 70B": "airoboros-70b",
             "🔮 Gemini Pro": "gemini-pro",
-            "⚡ GPT-4 Turbo": "gpt-4-turbo"
+            "📷 StabilityAI": "stabilityai/stable-diffusion-xl-base-1.0"
+            
         }
 
         columns = st.columns(3)  # Split the layout into three columns
@@ -105,8 +127,14 @@ def main():
 
                 except Exception as e:
                     st.error(f"An error occurred: {e}")
-            elif selected_model == "gpt-4-turbo":
-                st.write("Please use the default model for text generation.")
+                    
+            elif selected_model == "stabilityai/stable-diffusion-xl-base-1.0" or user_input == '/image:':
+                pipe = DiffusionPipeline.from_pretrained(models, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+                pipe.to("cuda")
+
+                params = {'num_inference_steps': 100, 'num_images_per_prompt': 2}
+                generate_image(pipe, prompt, params)
+                
             else:
                 try:
                     client = Client()