init

app.py

@@ -3,11 +3,59 @@ import threading
 import os
 import shutil
 import tempfile
+import time
+from collections import defaultdict, deque
 from util import process_image_edit, get_country_info_safe, get_location_info_safe, contains_chinese
 from nfsw import NSFWDetector
 
+# 配置参数
+NSFW_TIME_WINDOW = 5  # 时间窗口：5分钟
+NSFW_LIMIT = 6        # 限制次数：6次
+
 IP_Dict = {}
 NSFW_Dict = {}  # 记录每个IP的NSFW违规次数
+NSFW_Time_Dict = defaultdict(deque)  # 记录每个IP的NSFW检测时间戳
+
+def check_nsfw_rate_limit(client_ip):
+    """
+    检查IP的NSFW检测频率限制
+    
+    Args:
+        client_ip (str): 客户端IP地址
+    
+    Returns:
+        tuple: (是否超过限制, 剩余等待时间)
+    """
+    current_time = time.time()
+    time_window_start = current_time - (NSFW_TIME_WINDOW * 60)  # 5分钟前的时间戳
+    
+    # 清理过期的时间戳
+    while NSFW_Time_Dict[client_ip] and NSFW_Time_Dict[client_ip][0] < time_window_start:
+        NSFW_Time_Dict[client_ip].popleft()
+    
+    # 检查是否超过限制
+    if len(NSFW_Time_Dict[client_ip]) >= NSFW_LIMIT:
+        # 计算需要等待的时间
+        oldest_timestamp = NSFW_Time_Dict[client_ip][0]
+        wait_time = oldest_timestamp + (NSFW_TIME_WINDOW * 60) - current_time
+        return True, max(0, wait_time)
+    
+    return False, 0
+
+def record_nsfw_detection(client_ip):
+    """
+    记录IP的NSFW检测时间
+    
+    Args:
+        client_ip (str): 客户端IP地址
+    """
+    current_time = time.time()
+    NSFW_Time_Dict[client_ip].append(current_time)
+    
+    # 记录到NSFW_Dict中（兼容现有逻辑）
+    if client_ip not in NSFW_Dict:
+        NSFW_Dict[client_ip] = 0
+    NSFW_Dict[client_ip] += 1
 
 # 初始化NSFW检测器（从Hugging Face下载）
 try:
@@ -78,10 +126,17 @@ def edit_image_interface(input_image, prompt, request: gr.Request, progress=gr.P
             print(f"🔍 NSFW检测结果: {nsfw_result} - IP: {client_ip}({country_info})")
             
             if nsfw_result.lower() == "nsfw":
-                # 记录NSFW违规次数
-                if client_ip not in NSFW_Dict:
-                    NSFW_Dict[client_ip] = 0
-                NSFW_Dict[client_ip] += 1
+                # 检查NSFW频率限制
+                is_rate_limited, wait_time = check_nsfw_rate_limit(client_ip)
+                
+                if is_rate_limited:
+                    # 超过频率限制，显示等待提示
+                    wait_minutes = int(wait_time / 60) + 1  # 向上取整到分钟
+                    print(f"⚠️ NSFW频率限制 - IP: {client_ip}({country_info}), 需要等待 {wait_minutes} 分钟")
+                    return None, f"❌ 请等待 {wait_minutes} 分钟后继续生成"
+                
+                # 记录NSFW检测
+                record_nsfw_detection(client_ip)
                 
                 # 检查是否需要特殊处理（中文prompt或中国IP）
                 has_chinese = contains_chinese(prompt.strip())