Update app.py

app.py

@@ -43,7 +43,6 @@ if "current_window_start" not in st.session_state:
 if "selected_image" not in st.session_state:
    st.session_state.selected_image = None
 
-
 TOKEN = os.getenv("TOKEN0")
 API_URL = os.getenv("API_URL")
 token_id = 0
@@ -58,10 +57,10 @@ def send_verification_email(id_token):
         'requestType': 'VERIFY_EMAIL',
         'idToken': id_token
     }
-    
+
     response = requests.post(url, headers=headers, json=data)
     result = response.json()
-    
+
     if 'error' in result:
         return {'status': 'error', 'message': result['error']['message']}
     else:
@@ -75,10 +74,10 @@ def register_callback():
     try:
         # Step 1: Create a new user in Firebase
         user = auth.create_user(email=email, password=password)
-        
+
         # Step 2: Update the user profile with the display name
         auth.update_user(user.uid, display_name=display_name)
-        
+
         st.success("Registration successful! Sending verification email...")
 
         # Step 3: Sign in the user programmatically to get the id_token
@@ -150,7 +149,7 @@ def login_callback():
                       st.session_state.display_name = user.display_name # Store the display name
                       st.success("Logged in successfully!")
                       return
-                
+
                 raise Exception("User not found with provided credentials.")  # if not found, raise exception.
             except Exception as e:
                 st.error(f"Login failed: {e}") # if any error, display this message.
@@ -191,7 +190,7 @@ def generate_image(prompt, aspect_ratio, realism):
     global no_of_accounts
     global model_id
     payload = {
-        "id": model_id,  
+        "id": model_id,
         "inputs": [prompt, aspect_ratio, str(realism).lower()],
     }
     headers = {"Authorization": f"Bearer {TOKEN}"}
@@ -258,7 +257,6 @@ def store_image_data_in_db(user_id, prompt, aspect_ratio, realism, image_url, th
     except Exception as e:
         st.error(f"Failed to save image data: {e}")
 
-#Function to upload image to cloud storage
 #Function to upload image to cloud storage
 def upload_image_to_storage(image, user_id, is_thumbnail = False):
         try:
@@ -291,7 +289,7 @@ def load_image_data(user_id, start_index, batch_size):
         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)
@@ -339,8 +337,8 @@ def main_app():
     # Input fields
     prompt = st.text_input("Prompt", key="image_prompt", placeholder="Describe the image you want to generate")
     aspect_ratio = st.radio(
-        "Aspect Ratio", 
-        options=["1:1", "3:4", "4:3", "9:16", "16:9", "9:21", "21:9"], 
+        "Aspect Ratio",
+        options=["1:1", "3:4", "4:3", "9:16", "16:9", "9:21", "21:9"],
         index=5
     )
     realism = st.checkbox("Realism", value=False)
@@ -368,7 +366,7 @@ def main_app():
                                 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!")
                                 with st.container(border=True):
-                                    st.image(image, use_column_width=True)  # Display the (now thumbnailed) image here - consider fetching original for display
+                                    st.image(cloud_storage_url, use_column_width=True)  # Display high-res image from URL
                                     st.write(f"**Prompt:** {prompt}")
                                     st.write(f"**Aspect Ratio:** {aspect_ratio}")
                                     st.write(f"**Realism:** {realism}")
@@ -394,10 +392,10 @@ def main_app():
 
     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])
 
@@ -411,14 +409,14 @@ def main_app():
 
     # Dynamically load images for the window
     all_images = load_image_data(st.session_state.current_user, 0, 1000) # load all images
-    
+
     if all_images:
         num_images = len(all_images)
 
         # Calculate the range for images to display
         start_index = st.session_state.current_window_start
         end_index = min(start_index + st.session_state.window_size, num_images)
-        
+
         images_for_window = all_images[start_index:end_index]
 
         # Setup columns for horizontal slider layout
@@ -441,17 +439,16 @@ def main_app():
     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(cloud_storage_url, use_column_width=True)  # Display high-res image from URL
-            st.write(f"**Prompt:** {prompt}")
-            st.write(f"**Aspect Ratio:** {aspect_ratio}")
-            st.write(f"**Realism:** {realism}")
-            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", key=f"download_high_res_{uuid.uuid4()}")
+        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", key=f"download_overlay_{uuid.uuid4()}")
 
         if st.button("Close"):
           st.session_state.selected_image = None # close the modal when "close" is clicked