From 13b9a14c16a9bb44e49660d803c4ba97f2c6158c Mon Sep 17 00:00:00 2001
From: sheepgreen <58287293+slippersheepig@users.noreply.github.com>
Date: Sun, 21 Jan 2024 10:56:22 +0800
Subject: [PATCH] Update main.py

---
 main.py | 28 ++++++++++++++++++++++++----
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/main.py b/main.py
index 22bf70c..440100b 100644
--- a/main.py
+++ b/main.py
@@ -5,6 +5,7 @@
 import io
 from decouple import config
 from queue import Queue
+import time
 
 # 常量
 BOT_TOKEN = config('BOT_TOKEN')
@@ -18,6 +19,9 @@
 # 队列
 queue = Queue()
 
+# 用户请求时间字典
+user_last_request_time = {}
+
 # 请求到stablediffusion
 headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}
 
@@ -45,6 +49,19 @@ def generate_image(message, user, prompt):
     finally:
         queue.get()  # 从队列中移除用户
 
+# 增加频率限制函数
+def check_request_limit(user):
+    current_time = time.time()
+    last_request_time = user_last_request_time.get(user, 0)
+    time_since_last_request = current_time - last_request_time
+    # 设置请求时间间隔为60秒
+    request_interval = 60
+    if time_since_last_request < request_interval:
+        return False
+    else:
+        user_last_request_time[user] = current_time
+        return True
+
 @bot.message_handler(commands=['start'])
 def send_welcome(message):
     bot.reply_to(message, "嗨！我是一个用于生成stablediffusion图像的机器人，请输入 /help 以获取详细信息")
@@ -59,9 +76,12 @@ def stablediffusion_command(message):
     if user not in queue.queue:
         prompt = message.text.replace("/sd", "").strip()
         if prompt:
-            queue.put(user)
-            bot.reply_to(message, f'请稍等... \n您在队列中的位置: {queue.qsize()}')
-            generate_image(message, user, prompt)
+            if check_request_limit(user):
+                queue.put(user)
+                bot.reply_to(message, f'请稍等... \n您在队列中的位置: {queue.qsize()}')
+                generate_image(message, user, prompt)
+            else:
+                bot.reply_to(message, f"请求太频繁，请等待 {request_interval} 秒后再试.")
         else:
             bot.reply_to(message, '请求不能为空.')
     else:
@@ -77,6 +97,6 @@ def queue_command(message):
 
 def main():
     bot.polling()
-
+    
 if __name__ == "__main__":
     main()