Track the generation time in the block state

app.py

@@ -1,11 +1,8 @@
 import gradio as gr
 import time
-from collections import defaultdict
 from dotenv import load_dotenv
 import os
 from openai import OpenAI
-import random
-import uuid
 import secrets
 
 load_dotenv()
@@ -20,8 +17,6 @@ if APP_PASSWORD is None:
 
 GENERATION_TIMEOUT = float(os.getenv("GENERATION_TIMEOUT", 60))
 
-last_generation_time = defaultdict(float)
-
 def enhance_prompt(original_prompt):
     try:
         response = client.chat.completions.create(
@@ -38,20 +33,18 @@ def enhance_prompt(original_prompt):
         print("Ein Fehler ist beim Verbessern des Prompts aufgetreten:", e)
         return original_prompt
 
-def generate_image(prompt, use_magic_prompt, style, password, user_id):
-    global last_generation_time
-
+def generate_image(prompt, use_magic_prompt, style, password, user_id, last_generation_time):
     print(f"\nGenerierungsversuch von User {user_id} mit Prompt:\n{prompt}")
     
     if password != APP_PASSWORD:
         print("Generierung abgebrochen: Falsches Passwort.")
-        return "Falsches Passwort. Bitte versuche es erneut.", None
+        return "Falsches Passwort. Bitte versuche es erneut.", None,  last_generation_time
 
     current_time = time.time()
-    if current_time - last_generation_time[user_id] < GENERATION_TIMEOUT:
-        remaining_time = int(GENERATION_TIMEOUT - (current_time - last_generation_time[user_id]))
+    if current_time - last_generation_time < GENERATION_TIMEOUT:
+        remaining_time = max(0, int(GENERATION_TIMEOUT - (current_time - last_generation_time)))
         print("Generierung abgebrochen: Cooldown noch nicht abgelaufen")
-        return f"Bitte warte {remaining_time} Sekunden zwischen den Bildgenerierungen.", None
+        return f"Bitte warte noch {remaining_time} Sekunden bis zur nächsten Bildgenerierungen.\n\nNutze die Wartezeit, um deinen Prompt zu verfeinern. &#128521;", None,  last_generation_time
 
     if use_magic_prompt:
         prompt = enhance_prompt(prompt)
@@ -70,15 +63,15 @@ def generate_image(prompt, use_magic_prompt, style, password, user_id):
         image_url = response.data[0].url
         
         # Update the last generation time for this user
-        last_generation_time[user_id] = current_time
+        last_generation_time = current_time
         
-        return prompt, image_url
+        return prompt, image_url, last_generation_time
     except Exception as e:
         error_message = str(e)
         if "content_policy_violation" in error_message:
             return "Leider können wir deinen Prompt nicht verarbeiten, da er evtl. gegen die Content Policy von OpenAI verstösst. Wenn du denkst, dass das ein Fehler ist, versuche den Prompt so umzuformulieren, dass er jugendfreundlich ist.", None
         print("Ein Fehler ist aufgetreten:", e)
-        return f"Ein Fehler ist beim Generieren des Bildes aufgetreten: {error_message}", None
+        return f"Ein Fehler ist beim Generieren des Bildes aufgetreten: {error_message}", None, last_generation_time
 
 examples = [
     "Astronaut im Dschungel, kalte Farbpalette, gedämpfte Farben, detailliert, 8k",
@@ -91,9 +84,9 @@ examples = [
     "Ein mittelalterlicher Ritter, der gegen einen Drachen kämpft"
 ]
 
-def update_button_state(user_id):
+def update_button_state(last_generation_time):
     current_time = time.time()
-    remaining_time = max(0, int(GENERATION_TIMEOUT - (current_time - last_generation_time[user_id])))
+    remaining_time = max(0, int(GENERATION_TIMEOUT - (current_time - last_generation_time)))
     if remaining_time > 0:
         return gr.update(value=f"Bitte warte {remaining_time} Sekunden", interactive=False)
     else:
@@ -106,11 +99,12 @@ def generate_user_id():
     return secrets.token_urlsafe(16)
 
 def initialize_session():
-        new_user_id = generate_user_id()
-        return new_user_id, update_welcome_message(new_user_id)
+    new_user_id = generate_user_id()
+    return new_user_id, update_welcome_message(new_user_id), 0
 
 with gr.Blocks() as demo:
     user_id = gr.State(generate_user_id)
+    last_generation_time = gr.State(0)
 
     gr.Markdown("# codora DALL-E 3 Bildgenerator")
     welcome_message = gr.Markdown("Willkommen!")
@@ -140,23 +134,17 @@ with gr.Blocks() as demo:
 
     generate_button.click(
         fn=generate_image,
-        inputs=[prompt, use_magic_prompt, style, password, user_id],
-        outputs=[final_prompt, generated_image]
+        inputs=[prompt, use_magic_prompt, style, password, user_id, last_generation_time],
+        outputs=[final_prompt, generated_image, last_generation_time]
     )
 
-    demo.load(fn=initialize_session, outputs=[user_id, welcome_message])
+    demo.load(fn=initialize_session, outputs=[user_id, welcome_message, last_generation_time])
 
-    """demo.load(
-        fn=update_welcome_message,
-        inputs=[user_id],
-        outputs=[welcome_message]
-    )
-"""
     demo.load(
         fn=update_button_state,
-        inputs=[user_id],
+        inputs=[last_generation_time],
         outputs=[generate_button],
         every=1
     )
 
-demo.launch()
+demo.launch()