Update app.py

app.py

@@ -8,6 +8,79 @@ from PIL import Image
 from io import BytesIO
 import tempfile
 import ffmpeg
+import json
+from datetime import datetime, date
+from pathlib import Path
+
+# --- Database Setup ---
+DATA_DIR = Path("/data")
+DATA_DIR.mkdir(exist_ok=True)
+USAGE_DB_PATH = DATA_DIR / "usage_limits.json"
+
+def load_usage_db() -> dict:
+    """Load the usage database from disk."""
+    if USAGE_DB_PATH.exists():
+        try:
+            with open(USAGE_DB_PATH, 'r') as f:
+                return json.load(f)
+        except Exception as e:
+            print(f"Error loading usage database: {e}")
+            return {}
+    return {}
+
+def save_usage_db(db: dict):
+    """Save the usage database to disk."""
+    try:
+        with open(USAGE_DB_PATH, 'w') as f:
+            json.dump(db, f, indent=2)
+    except Exception as e:
+        print(f"Error saving usage database: {e}")
+
+def check_and_update_usage(username: str) -> bool:
+    """
+    Check if user has reached daily limit and update usage.
+    Returns True if user can generate, False if limit reached.
+    """
+    db = load_usage_db()
+    today = str(date.today())
+    
+    # Initialize user record if not exists
+    if username not in db:
+        db[username] = {"date": today, "count": 0}
+    
+    user_data = db[username]
+    
+    # Reset count if it's a new day
+    if user_data["date"] != today:
+        user_data["date"] = today
+        user_data["count"] = 0
+    
+    # Check if limit reached
+    if user_data["count"] >= 80:
+        return False
+    
+    # Increment count
+    user_data["count"] += 1
+    db[username] = user_data
+    save_usage_db(db)
+    
+    return True
+
+def get_remaining_generations(username: str) -> int:
+    """Get the number of remaining generations for today."""
+    db = load_usage_db()
+    today = str(date.today())
+    
+    if username not in db:
+        return 80
+    
+    user_data = db[username]
+    
+    # Reset if it's a new day
+    if user_data["date"] != today:
+        return 80
+    
+    return max(0, 80 - user_data["count"])
 
 # --- Google Gemini API Configuration ---
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
@@ -38,6 +111,27 @@ def verify_pro_status(token: Optional[Union[gr.OAuthToken, str]]) -> bool:
         print(f"Could not verify user's PRO/Enterprise status: {e}")
         return False
 
+def get_username(token: Optional[Union[gr.OAuthToken, str]]) -> Optional[str]:
+    """Get the username from the token."""
+    if not token:
+        return None
+    
+    if isinstance(token, gr.OAuthToken):
+        token_str = token.token
+    elif isinstance(token, str):
+        token_str = token
+    else:
+        return None
+    
+    try:
+        user_info = whoami(token=token_str)
+        username = user_info.get("name", None)
+        print(f"Username: {username}")
+        return username
+    except Exception as e:
+        print(f"Could not get username: {e}")
+        return None
+
 def _extract_image_data_from_response(response) -> Optional[bytes]:
     """Helper to extract image data from the model's response."""
     if hasattr(response, 'candidates') and response.candidates:
@@ -118,7 +212,17 @@ def _generate_video_segment(input_image_path: str, output_image_path: str, promp
     return result[0]["video"]
 
 def unified_image_generator(prompt: str, images: Optional[List[str]], previous_video_path: Optional[str], last_frame_path: Optional[str], manual_token: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
-    if not (verify_pro_status(oauth_token) or verify_pro_status(manual_token)): raise gr.Error("Access Denied.")
+    if not (verify_pro_status(oauth_token) or verify_pro_status(manual_token)): 
+        raise gr.Error("Access Denied.")
+    
+    # Check rate limit
+    username = get_username(oauth_token) or get_username(manual_token)
+    if not username:
+        raise gr.Error("Could not identify user.")
+    
+    if not check_and_update_usage(username):
+        raise gr.Error("This demo is made for interactive generations, not automated workflows. You have generated 80 images today, come back tomorrow for more.")
+    
     try:
         contents = [Image.open(image_path[0]) for image_path in images] if images else []
         contents.append(prompt)
@@ -135,7 +239,7 @@ def unified_image_generator(prompt: str, images: Optional[List[str]], previous_v
             # The crucial check for continuity
             if images[0][0] == last_frame_path:
                 can_extend_video = True
-                
+        
         return (output_path, gr.update(visible=can_create_video), gr.update(visible=can_extend_video), gr.update(visible=False))
     except Exception as e:
         raise gr.Error(f"Image generation failed: {e}. Rephrase your prompt to make image generation explicit and try again")
@@ -199,6 +303,7 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
                 generate_button = gr.Button("Generate", variant="primary")
             with gr.Column(scale=1):
                 output_image = gr.Image(label="Output", interactive=False, elem_id="output", type="filepath")
+                status_message = gr.Markdown("", visible=True)
                 use_image_button = gr.Button("♻️ Use this Image for Next Edit", variant="primary")
                 with gr.Row():
                     create_video_button = gr.Button("Create video between the two images 🎥", variant="secondary", visible=False)
@@ -215,15 +320,15 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
         triggers=[generate_button.click, prompt_input.submit],
         fn=unified_image_generator,
         inputs=[prompt_input, image_input_gallery, previous_video_state, last_frame_of_video_state, manual_token],
-        outputs=[output_image, create_video_button, extend_video_button, video_group]
+        outputs=[output_image, create_video_button, extend_video_button, video_group, status_message]
     )
     use_image_button.click(
         fn=lambda img: (
             [img] if img else None, None, gr.update(visible=False),
-            gr.update(visible=False), gr.update(visible=False)
+            gr.update(visible=False), gr.update(visible=False), ""
         ),
         inputs=[output_image],
-        outputs=[image_input_gallery, output_image, create_video_button, extend_video_button, video_group]
+        outputs=[image_input_gallery, output_image, create_video_button, extend_video_button, video_group, status_message]
     )
     create_video_button.click(
         fn=lambda: gr.update(visible=True), outputs=[video_group]
@@ -241,8 +346,12 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
     )
 
     def control_access(profile: Optional[gr.OAuthProfile] = None, oauth_token: Optional[gr.OAuthToken] = None):
-        if not profile: return gr.update(visible=False), gr.update(visible=False)
-        if verify_pro_status(oauth_token): return gr.update(visible=True), gr.update(visible=False)
+        if not profile: return gr.update(visible=False), gr.update(visible=False), ""
+        if verify_pro_status(oauth_token):
+            username = get_username(oauth_token)
+            remaining = get_remaining_generations(username) if username else 80
+            status = f"Welcome! You have {remaining} generations remaining today."
+            return gr.update(visible=True), gr.update(visible=False), status
         else:
             message = (
                 "## ✨ Exclusive Access for PRO Users\n\n"
@@ -250,8 +359,8 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
                 "To unlock this and many other cool stuff, please consider upgrading your account.\n\n"
                 "### [**Become a PRO Today!**](http://huggingface.co/subscribe/pro?source=nana_banana)"
             )
-            return gr.update(visible=False), gr.update(visible=True, value=message)
-    demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])
+            return gr.update(visible=False), gr.update(visible=True, value=message), ""
+    demo.load(control_access, inputs=None, outputs=[main_interface, pro_message, status_message])
 
 if __name__ == "__main__":
     demo.queue(max_size=None, default_concurrency_limit=None).launch(show_error=True)