Upload folder using huggingface_hub

app.py

@@ -1,6 +1,7 @@
 import os
 import requests
 import time
+import functools
 import threading
 import uuid
 import base64
@@ -18,78 +19,142 @@ API_KEY = os.getenv("WAVESPEED_API_KEY")
 if not API_KEY:
     raise ValueError("WAVESPEED_API_KEY is not set in environment variables")
 
+
 MODEL_URL = "TostAI/nsfw-text-detection-large"
-CLASS_NAMES = {0: "✅ SAFE", 1: "⚠️ QUESTIONABLE", 2: "🚫 UNSAFE"}
+TITLE = "🖼️🔍 Image Prompt Safety Classifier 🛡️"
+DESCRIPTION = "✨ Enter an image generation prompt to classify its safety level! ✨"
 
-try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
-    model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL)
-except Exception as e:
-    raise RuntimeError(f"Failed to load safety model: {str(e)}")
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL)
 
+# Define class names with emojis and detailed descriptions
+CLASS_NAMES = {
+    0: "✅ SAFE - This prompt is appropriate and harmless.",
+    1: "⚠️ QUESTIONABLE - This prompt may require further review.",
+    2: "🚫 UNSAFE - This prompt is likely to generate inappropriate content."
+}
 
-class SessionManager:
+
+@functools.lru_cache(maxsize=128)
+def classify_text(text):
+    inputs = tokenizer(text,
+                       return_tensors="pt",
+                       truncation=True,
+                       padding=True,
+                       max_length=1024)
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    logits = outputs.logits
+    predicted_class = torch.argmax(logits, dim=1).item()
+
+    return predicted_class, CLASS_NAMES[predicted_class]
+
+
+class ClientManager:
     _instances = {}
     _lock = threading.Lock()
 
     @classmethod
-    def get_session(cls, session_id):
+    def get_manager(cls, client_id=None):
+        if not client_id:
+            client_id = str(uuid.uuid4())
+
         with cls._lock:
-            if session_id not in cls._instances:
-                cls._instances[session_id] = {
-                    'count': 0,
-                    'history': [],
-                    'last_active': time.time()
-                }
-            return cls._instances[session_id]
+            if client_id not in cls._instances:
+                cls._instances[client_id] = ClientGenerationManager()
+            return cls._instances[client_id]
 
     @classmethod
-    def cleanup_sessions(cls):
+    def cleanup_old_clients(cls, max_age=3600):  # 1 hour default
+        current_time = time.time()
         with cls._lock:
-            now = time.time()
-            expired = [
-                k for k, v in cls._instances.items()
-                if now - v['last_active'] > 3600
-            ]
-            for k in expired:
-                del cls._instances[k]
+            to_remove = []
+            for client_id, manager in cls._instances.items():
+                if (hasattr(manager, "last_activity")
+                        and current_time - manager.last_activity > max_age):
+                    to_remove.append(client_id)
+
+            for client_id in to_remove:
+                del cls._instances[client_id]
 
 
-class RateLimiter:
+class ClientGenerationManager:
 
     def __init__(self):
-        self.clients = {}
         self.lock = threading.Lock()
+        self.last_activity = time.time()
+        self.request_timestamps = []  # Track timestamps of requests
 
-    def check(self, client_id):
+    def update_activity(self):
         with self.lock:
-            now = time.time()
-            if client_id not in self.clients:
-                self.clients[client_id] = {'count': 1, 'reset': now + 3600}
-                return True
-            if now > self.clients[client_id]['reset']:
-                self.clients[client_id] = {'count': 1, 'reset': now + 3600}
-                return True
-            if self.clients[client_id]['count'] >= 20:
-                return False
-            self.clients[client_id]['count'] += 1
-            return True
+            self.last_activity = time.time()
 
+    def add_request_timestamp(self):
+        with self.lock:
+            self.request_timestamps.append(time.time())
 
-session_manager = SessionManager()
-rate_limiter = RateLimiter()
+    def has_exceeded_limit(self, limit=20):
+        with self.lock:
+            current_time = time.time()
+            # Filter timestamps to only include those within the last hour
+            self.request_timestamps = [
+                ts for ts in self.request_timestamps
+                if current_time - ts <= 3600
+            ]
+            return len(self.request_timestamps) >= limit
 
 
-def create_error_image(message):
-    img = Image.new("RGB", (512, 512), "#ffdddd")
-    try:
-        font = ImageFont.truetype("arial.ttf", 24)
-    except:
-        font = ImageFont.load_default()
-    draw = ImageDraw.Draw(img)
-    text = f"Error: {message[:60]}..." if len(message) > 60 else message
-    draw.text((50, 200), text, fill="#ff0000", font=font)
-    return img
+class SessionManager:
+    _instances = {}
+    _lock = threading.Lock()
+
+    @classmethod
+    def get_manager(cls, session_id=None):
+        if session_id is None:
+            session_id = str(uuid.uuid4())
+
+        with cls._lock:
+            if session_id not in cls._instances:
+                cls._instances[session_id] = GenerationManager()
+            return session_id, cls._instances[session_id]
+
+    @classmethod
+    def cleanup_old_sessions(cls, max_age=3600):  # 1 hour default
+        current_time = time.time()
+        with cls._lock:
+            to_remove = []
+            for session_id, manager in cls._instances.items():
+                if (hasattr(manager, "last_activity")
+                        and current_time - manager.last_activity > max_age):
+                    to_remove.append(session_id)
+
+            for session_id in to_remove:
+                del cls._instances[session_id]
+
+
+class GenerationManager:
+
+    def __init__(self):
+        self.last_activity = time.time()
+        self.request_timestamps = []  # Track timestamps of requests
+
+    def update_activity(self):
+        self.last_activity = time.time()
+
+    def add_request_timestamp(self):
+        self.request_timestamps.append(time.time())
+
+    def has_exceeded_limit(self,
+                           limit=10):  # Default limit: 10 requests per hour
+        current_time = time.time()
+        # Filter timestamps to only include those within the last hour
+        self.request_timestamps = [
+            ts for ts in self.request_timestamps if current_time - ts <= 3600
+        ]
+        return len(self.request_timestamps) >= limit
 
 
 @torch.no_grad()
@@ -112,28 +177,48 @@ def decode_base64_to_image(base64_str):
     return Image.open(io.BytesIO(image_data))
 
 
-def generate_image(image_file,
-                   prompt,
-                   seed,
-                   session_id,
-                   enable_safety_checker=True):
+def generate_image(
+    image_file,
+    prompt,
+    seed,
+    session_id,
+    enable_safety_checker,
+    request: gr.Request,
+):
     try:
-        if enable_safety_checker:
-            safety_level = classify_prompt(prompt)
-            if safety_level != 0:
-                error_img = create_error_image(CLASS_NAMES[safety_level])
-                yield f"❌ Blocked: {CLASS_NAMES[safety_level]}", error_img, "", None
-                return
+        client_ip = request.client.host
+        x_forwarded_for = request.headers.get('x-forwarded-for')
+        if x_forwarded_for:
+            client_ip = x_forwarded_for
+        print(f"Client IP: {client_ip}")
+        client_generation_manager = ClientManager.get_manager(client_ip)
+        client_generation_manager.update_activity()
+        if client_generation_manager.has_exceeded_limit(limit=20):
+            error_message = "❌ Your network has exceeded the limit of 20 requests per hour. Please try again later."
+            yield error_message, None, "", None
+            return
 
-        if not rate_limiter.check(session_id):
-            error_img = create_error_image(
-                "Hourly limit exceeded (20 requests)")
-            yield "❌ Too many requests, please try again later", error_img, "", None
+        client_generation_manager.add_request_timestamp()
+        """Generate images with big status box during generation"""
+        # Get or create a session manager
+        session_id, manager = SessionManager.get_manager(session_id)
+        manager.update_activity()
+
+        # Check if the user has exceeded the request limit
+        if manager.has_exceeded_limit(
+                limit=10):  # Set the limit to 10 requests per hour
+            error_message = "❌ You have exceeded the limit of 10 requests per hour. Please try again later."
+            yield error_message, None, "", None
             return
 
-        session = session_manager.get_session(session_id)
-        session['last_active'] = time.time()
-        session['count'] += 1
+        # Add the current request timestamp
+        manager.add_request_timestamp()
+
+        if not prompt or prompt.strip() == "":
+            # Handle empty prompt case
+            error_message = "⚠️ Please enter a prompt first"
+            yield error_message, None, "", None
+            return
 
         error_messages = []
         if not image_file:
@@ -143,16 +228,27 @@ def generate_image(image_file,
         if not prompt.strip():
             error_messages.append("Prompt cannot be empty")
         if error_messages:
-            error_img = create_error_image(" | ".join(error_messages))
-            yield "❌ Input validation failed", error_img, "", None
+            error_message = "❌ Input validation failed: " + ", ".join(
+                error_messages)
+            yield error_message, None, "", None
             return
 
+        # Check if the prompt is safe
+        classification, message = classify_text(prompt)
+        if classification == 2:  # UNSAFE
+            yield "❌ NSFW prompt detected", None, "", None
+            return
+
+        # Status message
+        status_message = f"🔄 PROCESSING: '{prompt}'"
+        yield status_message, None, "", None
+
         try:
             base64_image = image_to_base64(image_file)
             input_image = decode_base64_to_image(base64_image)
         except Exception as e:
-            error_img = create_error_image(f"File processing failed: {str(e)}")
-            yield "❌ File processing failed", error_img, "", None
+            error_message = f"❌ File processing failed: {str(e)}"
+            yield error_message, None, "", None
             return
 
         headers = {
@@ -178,7 +274,7 @@ def generate_image(image_file,
         start_time = time.time()
 
         for _ in range(60):
-            time.sleep(1)
+            time.sleep(1.0)
             resp = requests.get(result_url, headers=headers)
             resp.raise_for_status()
 
@@ -188,25 +284,28 @@ def generate_image(image_file,
             if status == "completed":
                 elapsed = time.time() - start_time
                 output_url = data["outputs"][0]
-                session["history"].append(output_url)
-                yield f"🎉 Generation successful! Time taken {elapsed:.1f}s", output_url, output_url, update_recent_gallery(prompt, input_image, output_url)
+                yield f"🎉 Generation successful! Time taken {elapsed:.1f}s", output_url, output_url, update_recent_gallery(
+                    prompt, input_image, output_url)
                 return
             elif status == "failed":
                 raise Exception(data.get("error", "Unknown error"))
             else:
-                yield f"⏳ Current status: {status.capitalize()}...", None, None, None
+                error_message = f"⏳ Current status: {status.capitalize()}..."
+                yield error_message, None, "", None
 
         raise Exception("Generation timed out")
 
     except Exception as e:
-        error_img = create_error_image(str(e))
-        yield f"❌ Generation failed: {str(e)}", error_img, "", None
+        error_message = f"❌ Generation failed: {str(e)}"
+        yield error_message, None, "", None
 
 
+# Schedule periodic cleanup of old sessions
 def cleanup_task():
-    while True:
-        session_manager.cleanup_sessions()
-        time.sleep(3600)
+    SessionManager.cleanup_old_sessions()
+    ClientManager.cleanup_old_clients()
+    # Schedule the next cleanup
+    threading.Timer(3600, cleanup_task).start()  # Run every hour
 
 
 # Store recent generations
@@ -238,14 +337,16 @@ with gr.Blocks(theme=gr.themes.Soft(),
 
     with gr.Row():
         with gr.Column(scale=1):
-            prompt = gr.Textbox(label="Prompt",
-                                placeholder="Please enter your prompt...",
-                                lines=3)
             image_file = gr.Image(label="Upload Image",
                                   type="filepath",
-                                  sources=["upload"],
+                                  sources=["upload", "clipboard"],
                                   interactive=True,
-                                  image_mode="RGB")
+                                  image_mode="RGB",
+                                  value="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-original.png")
+            prompt = gr.Textbox(label="Prompt",
+                                placeholder="Please enter your prompt...",
+                                lines=3,
+                                value="Convert the image into Claymation style.")
             seed = gr.Number(label="seed",
                              value=-1,
                              minimum=-1,
@@ -256,8 +357,8 @@ with gr.Blocks(theme=gr.themes.Soft(),
                                         value=True,
                                         interactive=False)
         with gr.Column(scale=1):
-            status = gr.Textbox(label="Status", elem_classes=["status-box"])
             output_image = gr.Image(label="Generated Result")
+            status = gr.Textbox(label="Status", elem_classes=["status-box"])
             output_url = gr.Textbox(label="Image URL",
                                     interactive=True,
                                     visible=False)
@@ -266,15 +367,15 @@ with gr.Blocks(theme=gr.themes.Soft(),
         examples=[
             [
                 "Convert the image into Claymation style.",
-                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
+                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/para-attn/flux-original.png"
             ],
             [
                 "Convert the image into Ghibli style.",
-                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flux_ip_adapter_input.jpg"
+                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
             ],
             [
-                "Add sunglasses to the face of the girl.",
-                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_mask_girl2.png"
+                "Add sunglasses to the face of the statue.",
+                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flux_ip_adapter_input.jpg"
             ],
             # [
             #     'Convert the image into an ink sketch style.',
@@ -319,11 +420,16 @@ with gr.Blocks(theme=gr.themes.Soft(),
         inputs=[image_file, prompt, seed, session_id, enable_safety],
         outputs=[status, output_image, output_url, recent_gallery],
         api_name=False,
+        max_batch_size=10,
+        concurrency_limit=20,
+        concurrency_id="generation",
     )
 
 if __name__ == "__main__":
-    threading.Thread(target=cleanup_task, daemon=True).start()
-    app.queue(max_size=8).launch(
+    # Start the cleanup task
+    cleanup_task()
+    app.queue(max_size=20).launch(
         server_name="0.0.0.0",
+        max_threads=10,
         share=False,
     )