app.py

import gradio as gr
from flask import Flask, request
import os
import time
from collections import defaultdict

# Environment variables for authentication and model
token = os.environ["TOKEN"]
model = os.environ["MODEL"]

# Create Flask app
app = Flask(__name__)

# Rate-limiting configuration
RATE_LIMIT_WINDOW = 60  # Time window in seconds
RATE_LIMIT_REQUESTS = 5  # Max requests per user per window

# Dictionary to store request logs
request_log = defaultdict(list)

# Function to apply rate limiting
def rate_limiter():
    user_ip = request.remote_addr  # Retrieve the user's IP address
    current_time = time.time()

    # Clean up old requests outside the rate limit window
    request_log[user_ip] = [
        timestamp for timestamp in request_log[user_ip]
        if current_time - timestamp < RATE_LIMIT_WINDOW
    ]

    # Check if the user exceeded the allowed requests
    if len(request_log[user_ip]) >= RATE_LIMIT_REQUESTS:
        return "❌ Rate limit exceeded. Please wait and try again later."

    # Log the current request timestamp
    request_log[user_ip].append(current_time)
    return None  # Indicate no rate limit violation

# Load the Gradio model from Hugging Face Spaces
demo = gr.load(model, src="spaces", token=token)

# Route to serve the Gradio interface with rate limiting
@app.route("/", methods=["GET", "POST"])
def gradio_interface():
    # Apply rate limiting
    rate_limit_message = rate_limiter()
    if rate_limit_message:
        return rate_limit_message  # Return rate limit error if exceeded

    # Launch Gradio app
    return demo.launch(prevent_thread_lock=True)

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)  # Run the Flask app