ready for demo

app.py

@@ -2,53 +2,91 @@ import streamlit as st
 import requests
 
 API_URL = 'https://pic-gai.up.railway.app'
-API_URL = 'http://localhost:8000'
+# API_URL = 'http://localhost:8000'
 
-st.title('Design Generator UI')
+st.title('CollageAI')
 
 # Input field for user prompt
-user_prompt = st.text_input('Enter a prompt describing the scenario:')
+user_prompt = st.text_area(
+    "Describe the design you'd like to create:",
+    placeholder="For our anniversary, I want to write a card to my partner to celebrate our love and share all the things I adore about them."
+)
+
+# pick number of photos
+photos_count = st.number_input(
+    'How many photos would you like to use?',
+    min_value=0,
+    max_value=10,
+)
 
 # Submit buttons for templates and stickers
-st.write('Select an action:')
-generate_templates_button = st.button('Generate Templates')
-generate_stickers_button = st.button('Generate Stickers')
+generate_button = st.button('Generate')
+
+def gallery(column, images):
+    groups = []
+    for i in range(0, len(images), column):
+        groups.append(images[i:i+column])
 
-# Containers for templates and stickers
-templates_container = st.container()
-stickers_container = st.container()
+    for group in groups:
+        cols = st.columns(column)
+        for i, image in enumerate(group):
+            cols[i].image(image)
 
-if generate_templates_button or generate_stickers_button:
+if generate_button:
     if user_prompt:
-        if generate_templates_button:
+        # Prepare the params with the user prompt
+        params = {
+            'prompt': user_prompt,
+            'photos_count': photos_count
+        }
+
+        # remove empty params
+        params = {k: v for k, v in params.items() if v is not None}
+
+        print(params)
+        # Templates
+        with st.container():
             # Define the FastAPI server URL for templates
             url = f"{API_URL}/api/templates"
 
-        if generate_stickers_button:
-            # Define the FastAPI server URL for stickers
+            with st.spinner('Generating templates...'):
+                # Make a request to the FastAPI server
+                response = requests.get(url, params=params)
+
+                # Display the response in the appropriate output container
+                if response.status_code == 200:
+                    templates = response.json().get('result', [])
+                    image_urls = [template.get('image_medium') for template in templates]
+                    
+                    if image_urls:
+                        st.divider()
+                        st.subheader('Generated templates:')
+                        gallery(4, image_urls[:8])
+                    else:
+                        st.warning('No images were generated. Please try again with a different prompt.')
+                else:
+                    st.error(f"Error: {response.status_code}")
+
+        with st.container():
+            # Define the FastAPI server URL for templates
             url = f"{API_URL}/api/stickers"
 
-        # Prepare the payload with the user prompt
-        payload = {'user_prompt': user_prompt}
-
-        # Make a request to the FastAPI server
-        response = requests.get(url, json=payload)
-
-        # Display the response in the appropriate output container
-        if response.status_code == 200:
-            result = response.json()
-            image_urls = result.get('generated_images', [])
-            
-            if image_urls:
-                output_text = 'Generated templates:' if generate_templates_button else 'Generated stickers:'
-                output_container = templates_container if generate_templates_button else stickers_container
-                output_container.write(output_text)
-                
-                for image_url in image_urls:
-                    output_container.image(image_url)
-            else:
-                st.warning('No images were generated. Please try again with a different prompt.')
-        else:
-            st.error(f"Error: {response.status_code}")
+            with st.spinner('Generating stickers...'):
+                # Make a request to the FastAPI server
+                response = requests.get(url, params=params)
+
+                # Display the response in the appropriate output container
+                if response.status_code == 200:
+                    stickers = response.json().get('result', [])
+                    image_urls = [sticker.get('image_url') for sticker in stickers]
+                    
+                    if image_urls:
+                        st.divider()
+                        st.subheader('Generated stickers:')
+                        gallery(4, image_urls[:8])
+                    else:
+                        st.warning('No images were generated. Please try again with a different prompt.')
+                else:
+                    st.error(f"Error: {response.status_code}")
     else:
         st.warning('Please enter a prompt before submitting.')