Update app.py

app.py

@@ -12,6 +12,11 @@ from datetime import datetime
 import pandas as pd
 import json
 import re
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
@@ -60,6 +65,7 @@ def save_image(img, prompt):
     image_metadata = pd.concat([image_metadata, new_row], ignore_index=True)
     likes_cache[filename] = {'likes': 0, 'dislikes': 0, 'hearts': 0}
     save_likes_cache(likes_cache)
+    logger.info(f"Saved new image: {filename}")
     return filename
 
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
@@ -90,6 +96,7 @@ def delete_all_images():
     image_metadata = pd.DataFrame(columns=['Filename', 'Prompt', 'Likes', 'Dislikes', 'Hearts', 'Created'])
     likes_cache = {}
     save_likes_cache(likes_cache)
+    logger.info("All images deleted")
     return get_image_gallery(), image_metadata.values.tolist()
 
 def delete_image(filename):
@@ -100,13 +107,21 @@ def delete_image(filename):
         if filename in likes_cache:
             del likes_cache[filename]
             save_likes_cache(likes_cache)
+        logger.info(f"Deleted image: {filename}")
     return get_image_gallery(), image_metadata.values.tolist()
 
 def vote(filename, vote_type):
-    global likes_cache
+    global likes_cache, image_metadata
+    logger.info(f"Voting on {filename} with {vote_type}")
     if filename in likes_cache:
         likes_cache[filename][vote_type.lower()] += 1
         save_likes_cache(likes_cache)
+        # Update image_metadata as well
+        if filename in image_metadata['Filename'].values:
+            image_metadata.loc[image_metadata['Filename'] == filename, vote_type] += 1
+        logger.info(f"Updated vote count for {filename}: {likes_cache[filename]}")
+    else:
+        logger.warning(f"File {filename} not found in likes_cache")
     return get_image_gallery(), image_metadata.values.tolist()
 
 def get_random_style():
@@ -296,19 +311,19 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     )
 
     like_button.click(
-        fn=lambda x: vote(x, 'likes'),
+        fn=lambda x: vote(x, 'Likes') if x else (get_image_gallery(), image_metadata.values.tolist()),
         inputs=[selected_image],
         outputs=[image_gallery, metadata_df],
     )
 
     dislike_button.click(
-        fn=lambda x: vote(x, 'dislikes'),
+        fn=lambda x: vote(x, 'Dislikes') if x else (get_image_gallery(), image_metadata.values.tolist()),
         inputs=[selected_image],
         outputs=[image_gallery, metadata_df],
     )
 
     heart_button.click(
-        fn=lambda x: vote(x, 'hearts'),
+        fn=lambda x: vote(x, 'Hearts') if x else (get_image_gallery(), image_metadata.values.tolist()),
         inputs=[selected_image],
         outputs=[image_gallery, metadata_df],
     )
@@ -341,9 +356,4 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         ],
         outputs=[result, seed, gr.HTML(visible=False), image_gallery, metadata_df],
         api_name="run",
-    )
-
-    demo.load(fn=update_gallery_and_metadata, outputs=[image_gallery, metadata_df])
-
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch(share=True, debug=False)
+    )