Update app.py

app.py

@@ -35,12 +35,18 @@ if "current_user" not in st.session_state:
     st.session_state.current_user = None
 if "display_name" not in st.session_state:
     st.session_state.display_name = None
+if "images_data" not in st.session_state:
+     st.session_state.images_data = []
+if "start_index" not in st.session_state:
+    st.session_state.start_index = 0
+if "batch_size" not in st.session_state:
+     st.session_state.batch_size = 6
 
 TOKEN = os.getenv("TOKEN0")
 API_URL = os.getenv("API_URL")
 token_id = 0
 tokens_tried = 0
-no_of_accounts = 7
+no_of_accounts = 11
 model_id = os.getenv("MODEL_ID")
 
 def send_verification_email(id_token):
@@ -118,6 +124,8 @@ def login_callback():
         st.session_state.logged_in = True
         st.session_state.current_user = user.uid
         st.session_state.display_name = user.display_name # Store the display name
+        st.session_state.images_data = []  # Reset images data on login
+        st.session_state.start_index = 0
         st.success("Logged in successfully!")
     except Exception as e:
         st.error(f"Login failed: {e}")
@@ -127,6 +135,8 @@ def logout_callback():
     st.session_state.logged_in = False
     st.session_state.current_user = None
     st.session_state.display_name = None
+    st.session_state.images_data = []  # Clear images data on logout
+    st.session_state.start_index = 0
     st.info("Logged out successfully!")
 
 # Function to get image from url
@@ -241,6 +251,24 @@ def upload_image_to_storage(image, user_id):
             st.error(f"Failed to upload image to cloud storage: {e}")
             return None
 
+#Function to load image data from the database
+def load_image_data(user_id, start_index, batch_size):
+    try:
+        ref = db.reference(f'users/{user_id}/images')
+        snapshot = ref.order_by_child('timestamp').limit_to_last(start_index + batch_size).get()
+        if snapshot:
+            image_list = list(snapshot.items())
+            image_list.reverse()  # Reverse to show latest first
+            new_images = []
+            for key, val in image_list[start_index:]:
+              new_images.append(val)
+            return new_images
+        else:
+           return []
+    except Exception as e:
+        st.error(f"Failed to fetch image data from database: {e}")
+        return []
+
 # Registration form
 def registration_form():
     with st.form("Registration"):
@@ -296,6 +324,31 @@ def main_app():
         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