Update app.py

app.py

@@ -332,16 +332,17 @@ def main_app():
                               # Store image data in database
                               store_image_data_in_db(st.session_state.current_user, prompt, aspect_ratio, realism, cloud_storage_url, thumbnail_url)
                               st.success("Image stored to database successfully!")
+                              st.image(image, caption="Generated Image", use_column_width=True) # Display high res image on top
+                              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("Failed to upload image to cloud storage.")
                       else:
                          st.error("Failed to upload thumbnail to cloud storage.")
                     else:
                         st.error("Failed to create thumbnail")
 
-                    st.image(image, caption="Generated Image", use_column_width=True) # Display high res image on top
-
-                    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:
@@ -384,14 +385,19 @@ def main_app():
         cols = st.columns(num_images_to_display)
         for i, image_data in enumerate(images_for_window):
             with cols[i]:
-              if image_data.get('thumbnail_url'):
-                 with st.expander("More", expanded=False):
-                      st.image(image_data['image_url'], width = 200)
-                      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("---")
-                 st.image(image_data['thumbnail_url'], width = 150)
+              if image_data.get('thumbnail_url') and image_data.get('image_url'):
+                with st.container():
+                      if st.button("More", key = f"more_{i}"):
+                        with st.container(border = True):
+                          st.image(image_data['image_url'], use_column_width=True)
+                          st.write(f"**Prompt:** {image_data['prompt']}")
+                          st.write(f"**Aspect Ratio:** {image_data['aspect_ratio']}")
+                          st.write(f"**Realism:** {image_data['realism']}")
+                          download_path = download_image(image_data['image_url'])
+                          if download_path:
+                            st.download_button(label="Download Image", data = open(download_path, "rb"), file_name = f"image.png")
+                      st.image(image_data['thumbnail_url'], width = 150) # display thumbnail
+
               else:
                  st.image(image_data['image_url'], width = 150)
                  st.write(f"**Prompt:** {image_data['prompt']}")