Update app.py

app.py

@@ -285,8 +285,6 @@ def login_form():
         login_identifier = st.text_input("Email or Username", key="login_identifier")
         password = st.text_input("Password", type="password", key="login_password")
         submit_button = st.form_submit_button("Login", on_click=login_callback)
-
-# Main app screen (after login)
 def main_app():
     st.subheader(f"Welcome, {st.session_state.display_name}!")
     st.write("Enter a prompt below to generate an image.")
@@ -325,33 +323,102 @@ def main_app():
             st.warning("Please enter a prompt to generate an image.")
 
     st.header("Your Generated Images")
-    # Load initial batch of images
-    if st.session_state.images_data == []:
+
+    # Load initial batch of images (or reset if needed)
+    if st.session_state.images_data == [] or st.session_state.load_more_pressed == True:
         st.session_state.images_data = load_image_data(st.session_state.current_user, st.session_state.start_index, st.session_state.batch_size)
+        st.session_state.load_more_pressed = False
 
+    # Setup columns for horizontal slider layout
     if st.session_state.images_data:
-        for image_data in st.session_state.images_data:
-            col1, col2 = st.columns([1,2])
-            with col1:
-              st.image(image_data['image_url'], width = 150)
-            with col2:
-              st.write(f"**Prompt:** {image_data['prompt']}")
-              st.write(f"**Aspect Ratio:** {image_data['aspect_ratio']}")
-              st.write(f"**Realism:** {image_data['realism']}")
-              st.markdown("---")
+        num_images = len(st.session_state.images_data)
+        cols = st.columns(num_images) # dynamically create columns according to data
+        for i, image_data in enumerate(st.session_state.images_data):
+           with cols[i]:
+               st.image(image_data['image_url'], width = 150)
+               st.write(f"**Prompt:** {image_data['prompt']}")
+               st.write(f"**Aspect Ratio:** {image_data['aspect_ratio']}")
+               st.write(f"**Realism:** {image_data['realism']}")
+               st.markdown("---")
+    else:
+      st.write("No image generated yet!")
 
     # Load more button
     if st.button("Load More Images"):
         st.session_state.start_index += st.session_state.batch_size
-        new_images = load_image_data(st.session_state.current_user, st.session_state.start_index, st.session_state.batch_size)
-        if new_images:
-          st.session_state.images_data.extend(new_images)
-        else:
-          st.write("No more images available")
+        st.session_state.load_more_pressed = True
+        # No more loading logic here as we are now loading all the images together using `st.session_state.load_more_pressed` flag
+
 
     # Logout button
     if st.button("Logout", on_click=logout_callback):
         pass
+# Main app screen (after login)
+# def main_app():
+#     st.subheader(f"Welcome, {st.session_state.display_name}!")
+#     st.write("Enter a prompt below to generate an image.")
+
+#     # Input fields
+#     prompt = st.text_input("Prompt", key="image_prompt", placeholder="Describe the image you want to generate")
+#     aspect_ratio = st.radio(
+#         "Aspect Ratio", 
+#         options=["1:1", "3:4", "4:3", "9:16", "16:9", "9:21", "21:9"], 
+#         index=5
+#     )
+#     realism = st.checkbox("Realism", value=False)
+
+#     if st.button("Generate Image"):
+#         if prompt:
+#             with st.spinner("Generating Image..."):
+#                 image_result = generate_image(prompt, aspect_ratio, realism)
+#                 if isinstance(image_result, tuple) and len(image_result) == 2 and image_result[0] is not None:
+#                     image, image_url = image_result
+#                     st.image(image, caption="Generated Image", use_column_width=True)
+
+#                      # Upload image to cloud storage and store url
+#                     cloud_storage_url = upload_image_to_storage(image, st.session_state.current_user)
+
+#                     if cloud_storage_url:
+#                         # Store image data in database
+#                         store_image_data_in_db(st.session_state.current_user, prompt, aspect_ratio, realism, cloud_storage_url)
+#                         st.success("Image stored to database successfully!")
+
+#                     download_path = download_image(image_url)
+#                     if download_path:
+#                         st.download_button(label="Download Image", data = open(download_path, "rb"), file_name = f"image.png")
+#                 else:
+#                     st.error(f"Image generation failed: {image_result}")
+#         else:
+#             st.warning("Please enter a prompt to generate an image.")
+
+#     st.header("Your Generated Images")
+#     # Load initial batch of images
+#     if st.session_state.images_data == []:
+#         st.session_state.images_data = load_image_data(st.session_state.current_user, st.session_state.start_index, st.session_state.batch_size)
+
+#     if st.session_state.images_data:
+#         for image_data in st.session_state.images_data:
+#             col1, col2 = st.columns([1,2])
+#             with col1:
+#               st.image(image_data['image_url'], width = 150)
+#             with col2:
+#               st.write(f"**Prompt:** {image_data['prompt']}")
+#               st.write(f"**Aspect Ratio:** {image_data['aspect_ratio']}")
+#               st.write(f"**Realism:** {image_data['realism']}")
+#               st.markdown("---")
+
+#     # Load more button
+#     if st.button("Load More Images"):
+#         st.session_state.start_index += st.session_state.batch_size
+#         new_images = load_image_data(st.session_state.current_user, st.session_state.start_index, st.session_state.batch_size)
+#         if new_images:
+#           st.session_state.images_data.extend(new_images)
+#         else:
+#           st.write("No more images available")
+
+#     # Logout button
+#     if st.button("Logout", on_click=logout_callback):
+#         pass
 
 # Render the appropriate screen based on login status
 if st.session_state.logged_in: