Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
-import freegpt
 from freeGPT import Client
+from PIL import Image
+from io import BytesIO
 
 # Initialize the client with your API key
 client = Client()
@@ -11,6 +12,9 @@ st.write("Welcome to our ChatBot! Ask me anything.")
 # Create a text area for user input
 user_input = st.text_area("Enter your message here: ", height=100)
 
+# Create a dropdown menu to select the mode
+mode_dropdown = st.selectbox("Select Mode", ["Text", "Image"], index=0)
+
 # Create a submit button
 submit_button = st.button("Submit")
 
@@ -18,11 +22,22 @@ submit_button = st.button("Submit")
 def handle_form():
     global user_input
     if user_input != "":
-        # Generate response using the freeGPT model
-        response = client.create_completion("MODEL", user_input)
-        
-        # Print the response
-        st.write(response)
+        # Check the selected mode
+        if mode_dropdown == "Text":
+            # Generate response using the freeGPT model
+            response = client.create_completion("gpt3", user_input)
+            
+            # Print the response
+            st.write(response)
+        elif mode_dropdown == "Image":
+            # Generate image using the freeGPT model
+            img_bytes = client.create_image("prodia", user_input)
+            
+            # Convert bytes to image object
+            img = Image.open(BytesIO(img_bytes))
+            
+            # Show the image
+            st.image(img, width=None, height=None)
     else:
         st.write("Please enter a valid message.")
 
@@ -30,4 +45,4 @@ def handle_form():
 submit_button.onclick(handle_form)
 
 # Display the form
-st.form(user_input, submit_button)
+st.form(user_input, mode_dropdown, submit_button)