app.py

import streamlit as st
import requests  # For downloading images
import tempfile

# Set page config
st.set_page_config(page_title='Image Generation', page_icon=':money_with_wings:', layout='wide')

st.title('Image Generation')

# Initialize session state
if 'dynamic_prompt' not in st.session_state:
    st.session_state['dynamic_prompt'] = ""
if 'image_generated' not in st.session_state:
    st.session_state['image_generated'] = False
if 'image_url' not in st.session_state:
    st.session_state['image_url'] = ""

# Main content
with st.form(key='my_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 order input
    pizza_names = ["Crispy Chicken", "Chicken Pesto", ...]  # truncated for brevity
    name = st.selectbox("Name of Pizza", pizza_names)
    size_pizza = st.selectbox("Size of Pizza", ["Small", "Medium", "Large", "Extra Large"])
    toppings_options = ["Arugula", "Bacon", ...]  # truncated for brevity
    toppings = st.multiselect("Toppings", toppings_options, default=["Mushroom", "Tomatoes", "Onion"])
    prompt = st.text_input(label='Enter additional Descriptions ')
    bg_color_options = ['White', 'Black', ...]  # truncated for brevity
    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"])
    
    generate_prompt_button = st.form_submit_button(label='Generate Prompt')

# Action when the generate prompt button is clicked
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 following toppings on it  {', '.join(toppings)}. I want to generate image for the given description in {orientation_checkbox} with back ground color as {bg_color} "
    st.write(st.session_state['dynamic_prompt'])
    st.session_state['image_url'] = "url_of_your_generated_image" # Here, integrate with your custom model
    st.session_state['image_generated'] = True

# Display the button to generate the image only if the prompt has been generated
if st.session_state['dynamic_prompt']:
    with st.form(key='image_form'):
        generate_image_button = st.form_submit_button(label='Generate Image')

    if generate_image_button:
        # Display the image
        if st.session_state['image_url']:
            st.image(st.session_state['image_url'], caption="Generated image", use_column_width=True)

        # 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)  # Saving the image data
                st.download_button(
                    label="Download image",
                    data=tmp_file,
                    file_name="generated_image.png",
                    mime="image/png"
                )

        combined_info = f"""
        Image Generation Details:
        - Text Prompt: {st.session_state['dynamic_prompt']}
        - Image Size: {size}
        - Number of Images: {num_images}
        """
        st.write(combined_info)