Update app.py

app.py

@@ -306,7 +306,7 @@ def main_app():
                     image, image_url = image_result
                     st.image(image, caption="Generated Image", use_column_width=True)
 
-                     # Upload image to cloud storage and store url
+                    # 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:
@@ -324,8 +324,11 @@ def main_app():
 
     st.header("Your Generated Images")
 
-    # Load initial batch of images (or reset if needed)
+     # Load initial batch of images
     if st.session_state.images_data == [] or st.session_state.load_more_pressed == True:
+        if st.session_state.load_more_pressed == True:
+           st.session_state.start_index += st.session_state.batch_size
+        # load all images if new login or load_more is pressed, also reset the flag
         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
 
@@ -343,11 +346,11 @@ def main_app():
     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
         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
+        # No loading here, we are now loading all the images when `load_more_pressed` flag is set to True
 
 
     # Logout button