add more

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 from datasets import load_dataset
 from PIL import Image
 
+from set import ExpiringMap
 # from model import get_sd_small, get_sd_tiny, get_sd_every
 from trans_google import google_translator
 import replicate
@@ -52,7 +53,6 @@ re_sampler = [
 rand = random.Random()
 translator = google_translator()
 
-
 # tiny_pipe = get_sd_tiny()
 # small_pipe = get_sd_small()
 # every_pipe = get_sd_every()
@@ -66,8 +66,29 @@ translator = google_translator()
 #     else:
 #         return every_pipe
 
+time_client_map = ExpiringMap()
+count_client_map = ExpiringMap()
+
+
+def infer(prompt: str, negative: str, width: int, height: int, sampler: str,
+          steps: int, seed: int, scale, request: gr.Request):
+    client_ip = request.client.host
+    if client_ip != '127.0.0.1' and client_ip != 'localhost' and client_ip != '0.0.0.0':
+        if time_client_map.get(client_ip):
+            return "Too many requests from this IP, please try again later."
+        else:
+            time_client_map.put(client_ip, 1, 10)  # 添加一个过期时间为 10 秒的项
+
+    count = count_client_map.get(client_ip)
+    if count is None:
+        count = 0
+
+    count += 1
+    if count > 5:
+        return "Too many requests from this IP, please try again later."
+    else:
+        count_client_map.put(client_ip, count, 24 * 60 * 60)  # 添加一个过期时间为 24 小时的项
 
-def infer(prompt: str, negative: str, width: int, height: int, sampler: str, steps: int, seed: int, scale):
     global is_gpu_busy
 
     if seed == 0:
@@ -350,7 +371,7 @@ with block:
                    Also Change params to have a try <br />
                    more size may cost more time. <br />
                    It's just a simplified demo, you can use more advanced features optimize image quality <br />"""
-    tutorial_link = "https://docs.cworld.ai/docs/cworld-ai/%E5%BF%AB%E9%80%9F%E5%BC%80%E5%A7%8B/stable-diffusion/"
+    tutorial_link = "https://docs.cworld.ai/docs/cworld-ai/quick-start-stable-diffusion"
 
     gr.HTML(
         f"""

set.py

@@ -0,0 +1,22 @@
+from datetime import datetime, timedelta
+
+
+class ExpiringMap:
+    def __init__(self):
+        self.data = defaultdict(dict)
+
+    def put(self, key, value, expire_time):
+        current_time = datetime.now()
+        expire_timestamp = current_time + timedelta(seconds=expire_time)
+        self.data[key] = {'value': value, 'expire_time': expire_timestamp}
+
+    def get(self, key):
+        current_time = datetime.now()
+        if key in self.data and self.data[key]['expire_time'] > current_time:
+            return self.data[key]['value']
+        else:
+            return None
+
+    def remove(self, key):
+        if key in self.data:
+            del self.data[key]