Update app.py

app.py

@@ -1,78 +1,67 @@
 import streamlit as st
 import requests
 import tempfile
+import numpy as np
 
 # Set page config
-st.set_page_config(page_title='Image Generation', page_icon=':camera:', layout='wide')
+st.set_page_config(page_title='Image Generation', page_icon=':pizza:', layout='wide')
 
-st.title('Image Generation')
+st.title('Pizza Image Generation')
 
-# Initialize session state
+# Initialize session state for storing user inputs
 if 'dynamic_prompt' not in st.session_state:
     st.session_state['dynamic_prompt'] = ""
-if 'image_url' not in st.session_state:
-    st.session_state['image_url'] = ""
-if 'image_generated' not in st.session_state:
-    st.session_state['image_generated'] = False
+if 'image_data' not in st.session_state:
+    st.session_state['image_data'] = None
 
-# Define your toppings options
-toppings_options = ["Mushroom", "Tomatoes", "Onion", "Arugula", "Bacon", "Basil", "Broccoli", "Cheese", "Chicken", "Corn", "Ham", "Olives", "Pepperoni", "Peppers", "Pineapple"]
+# Define pizza options
+pizza_names = ["Crispy Chicken", "Chicken Pesto", "Deluxe Pepperoni", "Truffle Mushroom",
+               "Ultimate Cheese", "Spicy Veggie Ranch", "Classic Chicken Ranch",
+               "BBQ Chicken Ranch", "Spicy Chicken Ranch", "Very Veggie",
+               "Super Supreme", "Classic Pepperoni", "Margherita", "Cheeky Chicken",
+               "Chicken Super Supreme", "Chicken BBQ Supreme", "Chicken Fajita",
+               "Chicken Shawerma", "Other"]
 
-# Main content
-with st.form(key='prompt_form'):
-    # Image generation input
-    size = st.selectbox('Select size of the images', ('256x256', '512x512', '1024x1024'))
-    num_images = st.selectbox('Select Number of Images to be Generated', ('1', '2', '3', '4', '5'))
+pizza_sizes = ["Small", "Medium", "Large", "Extra Large"]
+
+toppings_options = ["Arugula", "Bacon", "Basil", "Broccoli", "Cheese", "Chicken", "Corn", 
+                    "Ham", "Mushroom", "Olives", "Onion", "Pepperoni", "Peppers", 
+                    "Pineapple", "Tomatoes"]
+
+bg_color_options = ['White', 'Black', 'Light Blue', 'Light Green', 'Light Yellow', 'Light Pink']
+
+orientation_options = ["Top View", "Side View", "Lateral View"]
 
-    # Pizza order input
-    pizza_names = ["Crispy Chicken", "Chicken Pesto", "..."]  # Add the rest of the pizza names
-    name = st.selectbox("Name of Pizza", pizza_names)
-    size_pizza = st.selectbox("Size of Pizza", ["Small", "Medium", "Large", "Extra Large"])
-    toppings = st.multiselect("Toppings", toppings_options, default=["Mushroom", "Tomatoes", "Onion"])
-    prompt = st.text_input(label='Enter additional Descriptions ')
-    bg_color_options = ['White', 'Black', 'Light Blue', 'Light Green', 'Light Yellow', 'Light Pink']
-    bg_color = st.selectbox('Select Background Color', bg_color_options)
-    orientation = ["Top View", "Side View", "Lateral View"]
-    orientation_checkbox = st.multiselect("Orientation", orientation, default=["Top View"])
+# UI Components for Pizza Preferences
+with st.form(key='prompt_form'):
+    # Pizza selection
+    selected_pizza = st.selectbox("Choose a Pizza Type", pizza_names)
+    selected_size = st.selectbox("Choose Pizza Size", pizza_sizes)
     
+    # Toppings selection
+    selected_toppings = st.multiselect("Choose Toppings", toppings_options, default=["Mushroom", "Tomatoes", "Onion"])
+
+    # Additional customization
+    additional_description = st.text_input("Enter additional descriptions (optional)")
+    background_color = st.selectbox("Choose Background Color", bg_color_options)
+    orientation = st.multiselect("Select Orientation", orientation_options, default=["Top View"])
+
     generate_prompt_button = st.form_submit_button(label='Generate Prompt')
 
-# Generate the dynamic prompt
+# Processing user input to generate dynamic prompt
 if generate_prompt_button:
-    st.session_state['dynamic_prompt'] = f"Imagine you are a marketer who wants to post an image on social media for a {size_pizza} {name} Pizza with the following toppings: {', '.join(toppings)}. I want to generate an image for the given description in {', '.join(orientation_checkbox)} with a background color of {bg_color}."
+    st.session_state['dynamic_prompt'] = f"A {selected_size} {selected_pizza} Pizza with {', '.join(selected_toppings)}. Additional description: {additional_description}. Background color: {background_color}. Orientation: {', '.join(orientation)}."
     st.write("Prompt:", st.session_state['dynamic_prompt'])
 
-# Display the button to generate the image
+# Rest of your code for image generation and display...
 with st.form(key='image_form'):
     generate_image_button = st.form_submit_button(label='Generate Image')
 
-
-
 if generate_image_button and st.session_state['dynamic_prompt']:
-    # Integrate with your custom model to generate the image
-    # Assuming 'my_custom_model_generate_image' is a function that returns the image URL
-    st.session_state['image_url'] = my_custom_model_generate_image(st.session_state['dynamic_prompt'])
-    st.session_state['image_generated'] = st.session_state['image_url'] is not None
-
-    # Display the generated image if the URL is valid
-    if st.session_state['image_url']:
-        try:
-            st.image(st.session_state['image_url'], caption="Generated Image", use_column_width=True)
-        except Exception as e:
-            st.error(f"Failed to load image from URL: {e}")
-
-
-
+    # Generate the image using your custom model
+    # Assuming 'my_custom_model_generate_image' is a function that returns the image data
+    st.session_state['image_data'] = my_custom_model_generate_image(st.session_state['dynamic_prompt'])
 
-# Download button for the generated image
-if st.session_state['image_generated']:
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
-        response = requests.get(st.session_state['image_url'])
-        with open(tmp_file.name, "wb") as f:
-            f.write(response.content)  # Save the image data
-        st.download_button(
-            label="Download Image",
-            data=tmp_file,
-            file_name="generated_image.png",
-            mime="image/png"
-        )
+    # Display the generated image if available
+    if st.session_state['image_data'] is not None:
+        st.image(st.session_state['image_data'], caption="Generated Image", use_column_width=True)