init

app.py

@@ -19,10 +19,91 @@ PHASE_2_WINDOW = 10     # 30-40 tries: 6 minutes
 PHASE_3_WINDOW = 15    # 40-60 tries: 10 minutes
 MAX_IMAGES_PER_WINDOW = 2  # Max images per time window
 
+# Restricted countries configuration
+RESTRICTED_COUNTRIES = ["印度", "巴基斯坦", "俄罗斯", "中国"]
+RESTRICTED_COUNTRY_LIMIT = 1  # Max usage for restricted countries
+
 IP_Dict = {}
 # IP generation statistics and time window tracking
 IP_Generation_Count = {}  # Record total generation count for each IP
 IP_Rate_Limit_Track = {}  # Record generation count and timestamp in current time window for each IP
+IP_Active_Tasks = {}  # Track active tasks for each IP (to prevent concurrent submissions)
+IP_Country_Cache = {}  # Cache IP country information to avoid repeated queries
+
+def query_ip_country(client_ip):
+    """
+    Query IP address geo information with robust error handling
+    
+    Returns:
+        dict: {"country": str, "region": str, "city": str}
+    """
+    # Check cache first
+    if client_ip in IP_Country_Cache:
+        print(f"Using cached IP data for {client_ip}")
+        return IP_Country_Cache[client_ip]
+    
+    # Validate IP address
+    if not client_ip or client_ip in ["127.0.0.1", "localhost", "::1"]:
+        print(f"Invalid or local IP address: {client_ip}, using default")
+        default_geo = {"country": "Unknown", "region": "Unknown", "city": "Unknown"}
+        IP_Country_Cache[client_ip] = default_geo
+        return default_geo
+    
+    # Query API with robust error handling
+    print(f"Querying IP geolocation for {client_ip}...")
+    
+    try:
+        import requests
+        from requests.exceptions import Timeout, ConnectionError, RequestException
+        
+        api_url = f"https://api.vore.top/api/IPdata?ip={client_ip}"
+        
+        # Make request with 3 second timeout
+        response = requests.get(api_url, timeout=3)
+        
+        if response.status_code == 200:
+            data = response.json()
+            if data.get("code") == 200 and "ipdata" in data:
+                ipdata = data["ipdata"]
+                geo_info = {
+                    "country": ipdata.get("info1", "Unknown"),
+                    "region": ipdata.get("info2", "Unknown"),
+                    "city": ipdata.get("info3", "Unknown")
+                }
+                IP_Country_Cache[client_ip] = geo_info
+                print(f"Successfully detected location for {client_ip}: {geo_info['country']}")
+                return geo_info
+        
+    except Exception as e:
+        print(f"Error querying IP {client_ip}: {e}, using default")
+    
+    # All failures lead here
+    default_geo = {"country": "Unknown", "region": "Unknown", "city": "Unknown"}
+    IP_Country_Cache[client_ip] = default_geo
+    return default_geo
+
+def is_restricted_country_ip(client_ip):
+    """
+    Check if IP is from a restricted country
+    
+    Returns:
+        bool: True if from restricted country
+    """
+    geo_info = query_ip_country(client_ip)
+    country = geo_info["country"]
+    return country in RESTRICTED_COUNTRIES
+
+def get_ip_max_limit(client_ip):
+    """
+    Get max usage limit for IP based on country
+    
+    Returns:
+        int: Max usage limit
+    """
+    if is_restricted_country_ip(client_ip):
+        return RESTRICTED_COUNTRY_LIMIT
+    else:
+        return RATE_LIMIT_60
 
 def get_ip_generation_count(client_ip):
     """
@@ -49,16 +130,25 @@ def get_ip_phase(client_ip):
         str: 'free', 'rate_limit_1', 'rate_limit_2', 'rate_limit_3', 'blocked'
     """
     count = get_ip_generation_count(client_ip)
+    max_limit = get_ip_max_limit(client_ip)
     
-    if count < FREE_TRY_N:           # 0-19 tries
+    # For restricted countries, check if they've reached their limit (1 task)
+    if is_restricted_country_ip(client_ip):
+        if count >= max_limit:
+            return 'blocked'
+        else:
+            return 'free'  # Only 1 free task for restricted countries
+    
+    # For normal countries, use standard limits
+    if count < FREE_TRY_N:           # 0-2 tries
         return 'free'
-    elif count < SLOW_TRY_N:         # 20-29 tries
-        return 'rate_limit_1'        # NSFW blur + 3 minutes 2 images
-    elif count < SLOW2_TRY_N:                 # 30-39 tries
-        return 'rate_limit_2'        # NSFW blur + 6 minutes 2 images
-    elif count < RATE_LIMIT_60:      # 40-59 tries
-        return 'rate_limit_3'        # NSFW blur + 10 minutes 2 images
-    else:                            # 60+ tries
+    elif count < SLOW_TRY_N:         # 3-4 tries
+        return 'rate_limit_1'        # NSFW blur + 5 minutes 2 images
+    elif count < SLOW2_TRY_N:        # 5-7 tries
+        return 'rate_limit_2'        # NSFW blur + 10 minutes 2 images
+    elif count < max_limit:          # 8-9 tries
+        return 'rate_limit_3'        # NSFW blur + 15 minutes 2 images
+    else:                            # 10+ tries
         return 'blocked'             # Generation blocked
 
 def check_rate_limit_for_phase(client_ip, phase):
@@ -113,6 +203,31 @@ def check_rate_limit_for_phase(client_ip, phase):
     
     return False, 0, track_data['count']
 
+def has_active_task(client_ip):
+    """
+    Check if IP has an active task in progress
+    
+    Returns:
+        bool: True if IP has an active task
+    """
+    return client_ip in IP_Active_Tasks and IP_Active_Tasks[client_ip]
+
+def set_active_task(client_ip, active=True):
+    """
+    Set or clear active task status for IP
+    
+    Args:
+        client_ip (str): Client IP address
+        active (bool): True to mark as active, False to clear
+    """
+    if active:
+        IP_Active_Tasks[client_ip] = True
+        print(f"🔒 Task locked for IP: {client_ip}")
+    else:
+        if client_ip in IP_Active_Tasks:
+            del IP_Active_Tasks[client_ip]
+        print(f"🔓 Task unlocked for IP: {client_ip}")
+
 def record_generation_attempt(client_ip, phase):
     """
     Record generation attempt
@@ -211,11 +326,38 @@ def edit_image_interface(input_image1, input_image2, input_image3, prompt, aspec
         print(f"⚠️ Request preprocessing error: {e}")
         return None, "❌ Request processing error", gr.update(visible=False)
     
+    # Check if user has an active task in progress
+    if has_active_task(client_ip):
+        concurrent_button_html = f"""
+        <div style='display: flex; justify-content: center; gap: 15px; margin: 10px 0 5px 0; padding: 0px;'>
+            <a href='https://omnicreator.net/multi-image-edit#generator' target='_blank' style='
+                display: inline-flex; 
+                align-items: center;
+                justify-content: center;
+                padding: 16px 32px; 
+                background: linear-gradient(135deg, #f39c12 0%, #e67e22 100%);
+                color: white; 
+                text-decoration: none; 
+                border-radius: 12px; 
+                font-weight: 600;
+                font-size: 16px;
+                text-align: center;
+                min-width: 200px;
+                box-shadow: 0 4px 15px rgba(243, 156, 18, 0.4);
+                transition: all 0.3s ease;
+                border: none;
+            '>⏰ Process Multiple Tasks Simultaneously</a>
+        </div>
+        """
+        return None, "❌ You have a task in progress. Please wait for it to complete before submitting a new task. Or visit https://omnicreator.net/multi-image-edit#generator to process multiple tasks simultaneously.", gr.update(value=concurrent_button_html, visible=True)
+    
     # Get user current phase
     current_phase = get_ip_phase(client_ip)
     current_count = get_ip_generation_count(client_ip)
+    geo_info = IP_Country_Cache.get(client_ip, {"country": "Unknown", "region": "Unknown", "city": "Unknown"})
+    is_restricted = is_restricted_country_ip(client_ip)
     
-    print(f"📊 User phase info - IP: {client_ip}, current phase: {current_phase}, generation count: {current_count}")
+    print(f"📊 User phase info - IP: {client_ip}, Location: {geo_info['country']}/{geo_info['region']}/{geo_info['city']}, Phase: {current_phase}, Count: {current_count}, Restricted: {is_restricted}, Active task: {has_active_task(client_ip)}")
     
     # Check if completely blocked
     if current_phase == 'blocked':
@@ -321,10 +463,16 @@ def edit_image_interface(input_image1, input_image2, input_image3, prompt, aspec
         record_generation_attempt(client_ip, current_phase)
         updated_count = get_ip_generation_count(client_ip)
         
+        # Mark this IP as having an active task
+        set_active_task(client_ip, active=True)
+        
         print(f"✅ Multi-image processing started - IP: {client_ip}, phase: {current_phase}, total count: {updated_count}, images: {len(valid_images)}, size: {width}x{height}, prompt: {prompt.strip()}", flush=True)
         
+        # Determine priority (first few tasks for each IP get priority)
+        task_priority = 1 if current_count < 2 else 0
+        
         # Call multi-image editing processing function
-        result_url, message, task_uuid = process_multi_image_edit(valid_images, prompt.strip(), width, height, progress_callback)
+        result_url, message, task_uuid = process_multi_image_edit(valid_images, prompt.strip(), width, height, progress_callback, priority=task_priority, client_ip=client_ip)
         
         if result_url:
             print(f"✅ Processing completed successfully - IP: {client_ip}, result_url: {result_url}, task_uuid: {task_uuid}", flush=True)
@@ -453,6 +601,9 @@ def edit_image_interface(input_image1, input_image2, input_image3, prompt, aspec
     except Exception as e:
         print(f"❌ Processing exception - IP: {client_ip}, error: {str(e)}")
         return None, f"❌ Error occurred during processing: {str(e)}", gr.update(visible=False)
+    finally:
+        # Always clear active task status when done (success or failure)
+        set_active_task(client_ip, active=False)
 
 def local_edit_interface(image_dict, prompt, request: gr.Request, progress=gr.Progress()):
     """
@@ -491,11 +642,38 @@ def local_edit_interface(image_dict, prompt, request: gr.Request, progress=gr.Pr
         print(f"⚠️ Local edit request preprocessing error: {e}")
         return None, "❌ Request processing error", gr.update(visible=False)
     
+    # Check if user has an active task in progress
+    if has_active_task(client_ip):
+        concurrent_button_html = f"""
+        <div style='display: flex; justify-content: center; gap: 15px; margin: 10px 0 5px 0; padding: 0px;'>
+            <a href='https://omnicreator.net/multi-image-edit#generator' target='_blank' style='
+                display: inline-flex; 
+                align-items: center;
+                justify-content: center;
+                padding: 16px 32px; 
+                background: linear-gradient(135deg, #f39c12 0%, #e67e22 100%);
+                color: white; 
+                text-decoration: none; 
+                border-radius: 12px; 
+                font-weight: 600;
+                font-size: 16px;
+                text-align: center;
+                min-width: 200px;
+                box-shadow: 0 4px 15px rgba(243, 156, 18, 0.4);
+                transition: all 0.3s ease;
+                border: none;
+            '>⏰ Process Multiple Tasks Simultaneously</a>
+        </div>
+        """
+        return None, "❌ You have a task in progress. Please wait for it to complete before submitting a new task. Or visit https://omnicreator.net/multi-image-edit#generator to process multiple tasks simultaneously.", gr.update(value=concurrent_button_html, visible=True)
+    
     # Get user current phase
     current_phase = get_ip_phase(client_ip)
     current_count = get_ip_generation_count(client_ip)
+    geo_info = IP_Country_Cache.get(client_ip, {"country": "Unknown", "region": "Unknown", "city": "Unknown"})
+    is_restricted = is_restricted_country_ip(client_ip)
     
-    print(f"📊 Local edit user phase info - IP: {client_ip}, current phase: {current_phase}, generation count: {current_count}")
+    print(f"📊 Local edit user phase info - IP: {client_ip}, Location: {geo_info['country']}/{geo_info['region']}/{geo_info['city']}, Phase: {current_phase}, Count: {current_count}, Restricted: {is_restricted}, Active task: {has_active_task(client_ip)}")
     
     # Check if completely blocked
     if current_phase == 'blocked':
@@ -588,10 +766,16 @@ def local_edit_interface(image_dict, prompt, request: gr.Request, progress=gr.Pr
         record_generation_attempt(client_ip, current_phase)
         updated_count = get_ip_generation_count(client_ip)
         
+        # Mark this IP as having an active task
+        set_active_task(client_ip, active=True)
+        
         print(f"✅ Local editing started - IP: {client_ip}, phase: {current_phase}, total count: {updated_count}, prompt: {prompt.strip()}", flush=True)
         
+        # Determine priority (first few tasks for each IP get priority)
+        task_priority = 1 if current_count < 2 else 0
+        
         # Call local image editing processing function
-        result_url, message, task_uuid = process_local_image_edit(base_image, layers, prompt.strip(), progress_callback)
+        result_url, message, task_uuid = process_local_image_edit(base_image, layers, prompt.strip(), progress_callback, priority=task_priority, client_ip=client_ip)
         
         if result_url:
             print(f"✅ Local editing completed successfully - IP: {client_ip}, result_url: {result_url}, task_uuid: {task_uuid}", flush=True)
@@ -720,6 +904,9 @@ def local_edit_interface(image_dict, prompt, request: gr.Request, progress=gr.Pr
     except Exception as e:
         print(f"❌ Local editing exception - IP: {client_ip}, error: {str(e)}")
         return None, f"❌ Error occurred during processing: {str(e)}", gr.update(visible=False)
+    finally:
+        # Always clear active task status when done (success or failure)
+        set_active_task(client_ip, active=False)
 
 # Create Gradio interface
 def create_app():
@@ -1380,6 +1567,33 @@ def create_app():
         </p>
     </div>
     """, padding=False)
+        
+        # Load IP country information on page load
+        def on_load(request: gr.Request):
+            """
+            Load page and query IP country information
+            """
+            try:
+                # Extract client IP
+                client_ip = request.client.host
+                headers = dict(request.headers) if hasattr(request, 'headers') else {}
+                x_forwarded_for = headers.get('x-forwarded-for') or headers.get('X-Forwarded-For')
+                if x_forwarded_for:
+                    client_ip = x_forwarded_for.split(',')[0].strip()
+                
+                print(f"Loading page for IP: {client_ip}")
+                
+                # Query IP country information (cached for subsequent requests)
+                query_ip_country(client_ip)
+                
+            except Exception as e:
+                print(f"Error in on_load: {e}")
+        
+        app.load(
+            on_load,
+            inputs=None,
+            outputs=None
+        )
     
     return app
 

util.py

@@ -32,6 +32,8 @@ R2_ACCESS_KEY = OneKey[4]
 R2_SECRET_KEY = OneKey[5]
 R2_ENDPOINT = OneKey[6]
 
+# HuggingFace用户任务限制配置
+HF_TASK_LIMIT = 20
 
 
 samples = [
@@ -273,10 +275,14 @@ def upload_mask_image_r2(client_ip, time_id, mask_image):
 
 
 
-def submit_image_edit_task(user_image_url, prompt, task_type="80", mask_image_url="", user_image2_url="", user_image3_url="", width=0, height=0):
+def submit_image_edit_task(user_image_url, prompt, task_type="80", mask_image_url="", user_image2_url="", user_image3_url="", width=0, height=0, client_ip="", priority=0):
     """
     Submit image editing task with improved error handling using API v2
     Supports multi-image editing with user_image2, user_image3, width, height parameters
+    
+    Args:
+        client_ip (str): Client IP address for task tracking
+        priority (int): Task priority, 1 for high priority, 0 for normal priority
     """
     headers = {
         'Content-Type': 'application/json',
@@ -284,6 +290,7 @@ def submit_image_edit_task(user_image_url, prompt, task_type="80", mask_image_ur
     }
     
     data = {
+        "client_ip": client_ip,
         "user_image": user_image_url,
         "user_image2": user_image2_url,
         "user_image3": user_image3_url,
@@ -293,8 +300,9 @@ def submit_image_edit_task(user_image_url, prompt, task_type="80", mask_image_ur
         "height": height,
         "text": prompt,
         "user_uuid": APIKEY,
-        "priority": 0,
-        "secret_key": "219ngu"
+        "priority": priority,
+        "secret_key": "219ngu",
+        "hf_task_limit": HF_TASK_LIMIT  # 传递HF用户任务限制配置
     }
     
     retry_count = 0
@@ -407,7 +415,7 @@ def check_task_status(task_id):
     return 'error', None, f"Failed after {max_retries} retries"
 
 
-def process_multi_image_edit(img_inputs, prompt, width=0, height=0, progress_callback=None):
+def process_multi_image_edit(img_inputs, prompt, width=0, height=0, progress_callback=None, priority=0, client_ip="127.0.0.1"):
     """
     Complete process for multi-image editing
     
@@ -417,10 +425,11 @@ def process_multi_image_edit(img_inputs, prompt, width=0, height=0, progress_cal
         width: Output width (0 for auto)
         height: Output height (0 for auto)
         progress_callback: Progress callback function
+        priority (int): Task priority, 1 for high priority, 0 for normal priority
+        client_ip (str): Client IP address for task tracking
     """
     try:
-        # Generate client IP and timestamp
-        client_ip = "127.0.0.1"  # Default IP
+        # Use provided client IP or default
         time_id = int(time.time())
         
         # Validate input images
@@ -476,7 +485,9 @@ def process_multi_image_edit(img_inputs, prompt, width=0, height=0, progress_cal
             user_image2_url=user_image2_url,
             user_image3_url=user_image3_url,
             width=width,
-            height=height
+            height=height,
+            client_ip=client_ip,
+            priority=priority
         )
         if error:
             return None, error, None
@@ -540,7 +551,7 @@ def process_image_edit(img_input, prompt, progress_callback=None):
     return process_multi_image_edit([img_input], prompt, 0, 0, progress_callback)
 
 
-def process_local_image_edit(base_image, layers, prompt, progress_callback=None):
+def process_local_image_edit(base_image, layers, prompt, progress_callback=None, priority=0, client_ip="127.0.0.1"):
     """
     处理局部图片编辑的完整流程
     
@@ -549,10 +560,11 @@ def process_local_image_edit(base_image, layers, prompt, progress_callback=None)
         layers (list): ImageEditor的层数据
         prompt (str): 编辑指令
         progress_callback: 进度回调函数
+        priority (int): Task priority, 1 for high priority, 0 for normal priority
+        client_ip (str): Client IP address for task tracking
     """
     try:
-        # Generate client IP and timestamp
-        client_ip = "127.0.0.1"  # Default IP
+        # Use provided client IP or default
         time_id = int(time.time())
         
         if progress_callback:
@@ -600,7 +612,14 @@ def process_local_image_edit(base_image, layers, prompt, progress_callback=None)
             progress_callback("submitting local edit task...")
         
         # 提交局部图片编辑任务 (task_type=81)
-        task_id, error = submit_image_edit_task(uploaded_url, prompt, task_type="81", mask_image_url=mask_url)
+        task_id, error = submit_image_edit_task(
+            uploaded_url, 
+            prompt, 
+            task_type="81", 
+            mask_image_url=mask_url,
+            client_ip=client_ip,
+            priority=priority
+        )
         if error:
             return None, error, None