Update app.py

app.py

@@ -40,6 +40,8 @@ if "window_size" not in st.session_state:
     st.session_state.window_size = 5
 if "current_window_start" not in st.session_state:
     st.session_state.current_window_start = 0
+if "selected_image" not in st.session_state:
+   st.session_state.selected_image = None
 
 
 TOKEN = os.getenv("TOKEN0")
@@ -349,13 +351,16 @@ def main_app():
             st.warning("Please enter a prompt to generate an image.")
 
     st.header("Your Generated Images")
-    # Initialize the current window, if it doesn't exist in session state
+     # Initialize the current window, if it doesn't exist in session state
     if "current_window_start" not in st.session_state:
       st.session_state.current_window_start = 0
 
     if "window_size" not in st.session_state:
       st.session_state.window_size = 5 # The number of images to display at a time
-
+    
+    if "selected_image" not in st.session_state:
+       st.session_state.selected_image = None
+    
     # Create left and right arrow buttons
     col_left, col_center, col_right = st.columns([1,8,1])
 
@@ -385,28 +390,34 @@ 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') 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']}")
-                 st.write(f"**Aspect Ratio:** {image_data['aspect_ratio']}")
-                 st.write(f"**Realism:** {image_data['realism']}")
-                 st.markdown("---")
+                if image_data.get('thumbnail_url') and image_data.get('image_url'):
+                  if st.button("More", key = f"more_{i}"):
+                      st.session_state.selected_image = image_data
+                      
+                  st.image(image_data['thumbnail_url'], width = 150)
+                else:
+                    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!")
 
+
+    # Display modal if an image is selected
+    if st.session_state.selected_image:
+       with st.container(border = True):
+           st.image(st.session_state.selected_image['image_url'], use_column_width=True)
+           st.write(f"**Prompt:** {st.session_state.selected_image['prompt']}")
+           st.write(f"**Aspect Ratio:** {st.session_state.selected_image['aspect_ratio']}")
+           st.write(f"**Realism:** {st.session_state.selected_image['realism']}")
+           download_path = download_image(st.session_state.selected_image['image_url'])
+           if download_path:
+              st.download_button(label="Download Image", data = open(download_path, "rb"), file_name = f"image.png")
+       if st.button("Close"):
+          st.session_state.selected_image = None # close the modal when "close" is clicked
+
     # Logout button
     if st.button("Logout", on_click=logout_callback):
         pass