Update app.py

app.py

@@ -9,6 +9,7 @@ from PIL import Image
 import tempfile
 import mimetypes
 import uuid
+import io
 
 # Load Firebase credentials from Hugging Face Secrets
 firebase_creds = os.getenv("FIREBASE_CREDENTIALS")
@@ -35,15 +36,12 @@ 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
-if "load_more_pressed" not in st.session_state:
-    st.session_state.load_more_pressed = False # Initialize load_more_pressed
-    
+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
+
+
 TOKEN = os.getenv("TOKEN0")
 API_URL = os.getenv("API_URL")
 token_id = 0
@@ -126,8 +124,6 @@ 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}")
@@ -137,8 +133,6 @@ 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
@@ -151,8 +145,6 @@ def get_image_from_url(url):
         response.raise_for_status()
         image = Image.open(BytesIO(response.content))
         return image, url # Return the image and the URL
-
-
     except requests.exceptions.RequestException as e:
         return f"Error fetching image: {e}", None
     except Exception as e:
@@ -217,7 +209,7 @@ def download_image(image_url):
          return None
 
 # Function to store image and related data in Firebase
-def store_image_data_in_db(user_id, prompt, aspect_ratio, realism, image_url):
+def store_image_data_in_db(user_id, prompt, aspect_ratio, realism, image_url, thumbnail_url):
     try:
         ref = db.reference(f'users/{user_id}/images')
         new_image_ref = ref.push()
@@ -226,6 +218,7 @@ def store_image_data_in_db(user_id, prompt, aspect_ratio, realism, image_url):
             'aspect_ratio': aspect_ratio,
             'realism': realism,
             'image_url': image_url,
+            'thumbnail_url' : thumbnail_url,
             'timestamp': {'.sv': 'timestamp'}
         })
         st.success("Image data saved successfully!")
@@ -233,11 +226,14 @@ def store_image_data_in_db(user_id, prompt, aspect_ratio, realism, image_url):
         st.error(f"Failed to save image data: {e}")
 
 #Function to upload image to cloud storage
-def upload_image_to_storage(image, user_id):
+def upload_image_to_storage(image, user_id, is_thumbnail = False):
         try:
             bucket = storage.bucket()
             image_id = str(uuid.uuid4())
-            file_path = f"user_images/{user_id}/{image_id}.png"
+            if is_thumbnail:
+                 file_path = f"user_images/{user_id}/thumbnails/{image_id}.png" # path for thumbnail
+            else:
+                 file_path = f"user_images/{user_id}/{image_id}.png" # path for high resolution images
             blob = bucket.blob(file_path)
 
             # Convert PIL Image to BytesIO object
@@ -271,6 +267,19 @@ def load_image_data(user_id, start_index, batch_size):
         st.error(f"Failed to fetch image data from database: {e}")
         return []
 
+# Function to create low resolution thumbnail
+def create_thumbnail(image, thumbnail_size = (150,150)):
+    try:
+        img_byte_arr = BytesIO()
+        image.thumbnail(thumbnail_size)
+        image.save(img_byte_arr, format='PNG')
+        img_byte_arr = img_byte_arr.getvalue()
+        thumbnail = Image.open(io.BytesIO(img_byte_arr)) # convert byte to PIL image
+        return thumbnail
+    except Exception as e:
+      st.error(f"Failed to create thumbnail: {e}")
+      return None
+
 # Registration form
 def registration_form():
     with st.form("Registration"):
@@ -288,7 +297,6 @@ def login_form():
         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.")
@@ -308,15 +316,28 @@ def main_app():
                 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)
+                    #create thumbnail
+                    thumbnail = create_thumbnail(image)
+
+                    if thumbnail:
+                      # Upload thumbnail to cloud storage and store url
+                      thumbnail_url = upload_image_to_storage(thumbnail, st.session_state.current_user, is_thumbnail=True)
+
+                      if thumbnail_url:
+                          # Upload image to cloud storage and store url
+                          cloud_storage_url = upload_image_to_storage(image, st.session_state.current_user, is_thumbnail=False)
+
+                          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, thumbnail_url)
+                              st.success("Image stored to database successfully!")
+                      else:
+                         st.error("Failed to upload thumbnail to cloud storage.")
+                    else:
+                        st.error("Failed to create thumbnail")
 
-                    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!")
+                    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:
@@ -327,7 +348,6 @@ 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
     if "current_window_start" not in st.session_state:
       st.session_state.current_window_start = 0
@@ -363,22 +383,24 @@ def main_app():
 
         cols = st.columns(num_images_to_display)
         for i, image_data in enumerate(images_for_window):
-            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("---")
-
+          with cols[i]:
+             if image_data.get('thumbnail_url'):
+                  if st.button(image_data.get('prompt') ,key=f"thumbnail_{i}"):
+                    st.image(image_data['image_url'], width = 200)
+                  st.image(image_data['thumbnail_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.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!")
 
     # 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:
-    main_app()
-else:
-    registration_form()
-    login_form()
+        pass