feat: implement rate limiting for video generation requests

app.py

@@ -16,6 +16,8 @@ from openai import OpenAI
 import base64
 from google.cloud import vision
 from google.oauth2 import service_account
+import time
+from collections import defaultdict, deque
 
 dotenv.load_dotenv()
 
@@ -25,6 +27,13 @@ SCRIPT_DIR = Path(__file__).parent
 MODAL_ENDPOINT = os.getenv('FAL_MODAL_ENDPOINT')
 MODAL_AUTH_TOKEN = os.getenv('MODAL_AUTH_TOKEN')
 
+# Rate limiting configuration
+RATE_LIMIT_GENERATIONS = int(os.getenv('RATE_LIMIT_GENERATIONS', '5'))  # Default 5 generations per hour
+RATE_LIMIT_WINDOW = int(os.getenv('RATE_LIMIT_WINDOW', '3600'))  # Default 1 hour in seconds
+
+# In-memory rate limiting storage (for production, consider Redis)
+user_generations = defaultdict(deque)
+
 loras = [
    {
       "image": "https://huggingface.co/Remade-AI/Crash-zoom-out/resolve/main/example_videos/1.gif",
@@ -576,11 +585,30 @@ def update_selection(evt: gr.SelectData):
   sentence = f"Selected LoRA: {selected_lora['title']}"
   return selected_lora['id'], sentence
 
-async def handle_generation(image_input, subject, selected_index, progress=gr.Progress(track_tqdm=True)):
+async def handle_generation(image_input, subject, selected_index, request: gr.Request, progress=gr.Progress(track_tqdm=True)):
     try:
         if selected_index is None:
             raise gr.Error("You must select a LoRA before proceeding.")
         
+        # Check rate limit first
+        user_identifier = get_user_identifier(request)
+        is_allowed, remaining, reset_time = check_rate_limit(user_identifier)
+        
+        if not is_allowed:
+            minutes = reset_time // 60
+            seconds = reset_time % 60
+            time_str = f"{minutes}m {seconds}s" if minutes > 0 else f"{seconds}s"
+            # Re-enable button on rate limit
+            yield None, None, gr.update(visible=False), gr.update(value="Generate", interactive=True)
+            raise gr.Error(f"Rate limit exceeded. Go to https://app.remade.ai for more generations and effects. Otherwise, you can generate {RATE_LIMIT_GENERATIONS} videos per hour. Try again in {time_str}.")
+        
+        # Record this generation attempt
+        record_generation(user_identifier)
+        
+        # Show remaining generations to user
+        if remaining > 0:
+            print(f"User {user_identifier} has {remaining} generations remaining this hour")
+        
         # First, moderate the prompt
         prompt_moderation = await moderate_prompt(subject)
         print(f"Prompt moderation result: {prompt_moderation}")
@@ -823,6 +851,30 @@ css = '''
 }
 '''
 
+def get_user_identifier(request: gr.Request) -> str:
+    """Get user identifier from request (IP address)"""
+    if request and hasattr(request, 'client') and hasattr(request.client, 'host'):
+        return request.client.host
+    return "unknown"
+
+def get_rate_limit_status(request: gr.Request) -> str:
+    """Get current rate limit status for display to user"""
+    try:
+        user_identifier = get_user_identifier(request)
+        is_allowed, remaining, reset_time = check_rate_limit(user_identifier)
+        
+        if remaining == 0 and reset_time > 0:
+            minutes = reset_time // 60
+            seconds = reset_time % 60
+            time_str = f"{minutes}m {seconds}s" if minutes > 0 else f"{seconds}s"
+            return f"⚠️ Rate limit reached. Try again in {time_str}"
+        elif remaining <= 2:
+            return f"⚡ {remaining} generations remaining this hour"
+        else:
+            return f"✅ {remaining} generations remaining this hour"
+    except:
+        return "✅ Ready to generate"
+
 with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate", text_size="lg")) as demo:
     selected_index = gr.State(None)
     current_generation_id = gr.State(None)
@@ -892,6 +944,9 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="
 
             subject = gr.Textbox(label="Describe your subject", placeholder="Cat toy")
             
+            # Rate limit status display
+            rate_limit_status = gr.Markdown("✅ Ready to generate", elem_id="rate_limit_status")
+            
             with gr.Row():
                 button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
                 audio_button = gr.Button("Add Audio 🔒", interactive=False)
@@ -997,6 +1052,13 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="
         inputs=None,
         outputs=None
     )
+    
+    # Update rate limit status on page load
+    demo.load(
+        fn=get_rate_limit_status,
+        inputs=None,
+        outputs=[rate_limit_status]
+    )
 
 def get_modal_auth_headers():
     """Get authentication headers for Modal API requests"""
@@ -1008,6 +1070,32 @@ def get_modal_auth_headers():
         'Content-Type': 'application/json'
     }
 
+def check_rate_limit(user_identifier: str) -> tuple[bool, int, int]:
+    """
+    Check if user has exceeded rate limit
+    Returns: (is_allowed, remaining_generations, reset_time_seconds)
+    """
+    current_time = time.time()
+    user_queue = user_generations[user_identifier]
+    
+    # Remove old entries outside the time window
+    while user_queue and current_time - user_queue[0] > RATE_LIMIT_WINDOW:
+        user_queue.popleft()
+    
+    # Check if user has exceeded limit
+    if len(user_queue) >= RATE_LIMIT_GENERATIONS:
+        # Calculate when the oldest entry will expire
+        reset_time = int(user_queue[0] + RATE_LIMIT_WINDOW - current_time)
+        return False, 0, reset_time
+    
+    remaining = RATE_LIMIT_GENERATIONS - len(user_queue)
+    return True, remaining, 0
+
+def record_generation(user_identifier: str):
+    """Record a new generation for the user"""
+    current_time = time.time()
+    user_generations[user_identifier].append(current_time)
+
 if __name__ == "__main__":
     demo.queue(default_concurrency_limit=20)
     demo.launch(ssr_mode=False, share=True)