util.py

import os
import sys
import cv2
import json
import random
import time
import datetime
import requests
import func_timeout
import numpy as np
import gradio as gr
import boto3
import tempfile
import io
import uuid
from botocore.client import Config
from PIL import Image


# TOKEN = os.environ['TOKEN']
# APIKEY = os.environ['APIKEY']
# UKAPIURL = os.environ['UKAPIURL']

OneKey = os.environ['OneKey'].strip()
OneKey = OneKey.split("#")
TOKEN = OneKey[0]
APIKEY = OneKey[1]
UKAPIURL = OneKey[2]
LLMKEY = OneKey[3]
R2_ACCESS_KEY = OneKey[4]
R2_SECRET_KEY = OneKey[5]
R2_ENDPOINT = OneKey[6]


# tmpFolder is no longer needed since we upload directly from memory
# tmpFolder = "tmp"
# os.makedirs(tmpFolder, exist_ok=True)


# Legacy function - no longer used since we upload directly from memory
# def upload_user_img(clientIp, timeId, img):
#     fileName = clientIp.replace(".", "")+str(timeId)+".jpg"
#     local_path = os.path.join(tmpFolder, fileName)
#     img = cv2.imread(img)
#     cv2.imwrite(os.path.join(tmpFolder, fileName), img)
# 
#     json_data = {
#         "token": TOKEN,  
#         "input1": fileName,
#         "input2": "",
#         "protocol": "",
#         "cloud": "ali"
#     }
#     
#     session = requests.session()
#     ret = requests.post(
#         f"{UKAPIURL}/upload",
#         headers={'Content-Type': 'application/json'},
#         json=json_data
#     )
#     
#     res = ""
#     if ret.status_code==200:
#         if 'upload1' in ret.json():
#             upload_url = ret.json()['upload1']
#             headers = {'Content-Type': 'image/jpeg'}
#             response = session.put(upload_url, data=open(local_path, 'rb').read(), headers=headers)
#             # print(response.status_code)
#             if response.status_code == 200:
#                 res = upload_url
#     if os.path.exists(local_path):
#         os.remove(local_path)
#     return res


class R2Api:

    def __init__(self, session=None):
        super().__init__()
        self.R2_BUCKET = "omni-creator"
        self.domain = "https://www.omnicreator.net/"
        self.R2_ACCESS_KEY = R2_ACCESS_KEY
        self.R2_SECRET_KEY = R2_SECRET_KEY
        self.R2_ENDPOINT = R2_ENDPOINT

        self.client = boto3.client(
            "s3",
            endpoint_url=self.R2_ENDPOINT,
            aws_access_key_id=self.R2_ACCESS_KEY,
            aws_secret_access_key=self.R2_SECRET_KEY,
            config=Config(signature_version="s3v4")
        )

        self.session = requests.Session() if session is None else session

    def upload_from_memory(self, image_data, filename, content_type='image/jpeg'):
        """
        Upload image data directly from memory to R2
        
        Args:
            image_data (bytes): Image data in bytes
            filename (str): Filename for the uploaded file
            content_type (str): MIME type of the image
            
        Returns:
            str: URL of the uploaded file
        """
        t1 = time.time()
        headers = {"Content-Type": content_type}
        
        cloud_path = f"ImageEdit/Uploads/{str(datetime.date.today())}/{filename}"
        url = self.client.generate_presigned_url(
            "put_object",
            Params={"Bucket": self.R2_BUCKET, "Key": cloud_path, "ContentType": content_type},
            ExpiresIn=604800
        )

        retry_count = 0
        while retry_count < 3:
            try:
                response = self.session.put(url, data=image_data, headers=headers, timeout=15)
                if response.status_code == 200:
                    break
                else:
                    print(f"⚠️ Upload failed with status code: {response.status_code}")
                    retry_count += 1
            except (requests.exceptions.Timeout, requests.exceptions.RequestException) as e:
                print(f"⚠️ Upload retry {retry_count + 1}/3 failed: {e}")
                retry_count += 1
                if retry_count == 3:
                    raise Exception(f'Failed to upload file to R2 after 3 retries! Last error: {str(e)}')
                time.sleep(1)  # 等待1秒后重试
                continue
        print("upload_from_memory time is ====>", time.time() - t1)
        return f"{self.domain}{cloud_path}"

def upload_user_img_r2(clientIp, timeId, pil_image):
    """
    Upload PIL Image directly to R2 without saving to local file
    
    Args:
        clientIp (str): Client IP address
        timeId (int): Timestamp
        pil_image (PIL.Image): PIL Image object
        
    Returns:
        str: Uploaded URL
    """
    # Generate unique filename using UUID to prevent file conflicts in concurrent environment
    unique_id = str(uuid.uuid4())
    fileName = f"user_img_{unique_id}_{timeId}.jpg"
    
    # Convert PIL Image to bytes
    img_buffer = io.BytesIO()
    if pil_image.mode != 'RGB':
        pil_image = pil_image.convert('RGB')
    pil_image.save(img_buffer, format='JPEG', quality=95)
    img_data = img_buffer.getvalue()
    
    # Upload directly from memory
    res = R2Api().upload_from_memory(img_data, fileName, 'image/jpeg')
    return res



def create_mask_from_layers(base_image, layers):
    """
    Create mask image from ImageEditor layers
    
    Args:
        base_image (PIL.Image): Original image
        layers (list): ImageEditor layer data
        
    Returns:
        PIL.Image: Black and white mask image
    """
    from PIL import Image, ImageDraw
    import numpy as np
    
    # Create blank mask with same size as original image
    mask = Image.new('L', base_image.size, 0)  # 'L' mode is grayscale, 0 is black
    
    if not layers:
        return mask
    
    # Iterate through all layers, set drawn areas to white
    for layer in layers:
        if layer is not None:
            # Convert layer to numpy array
            layer_array = np.array(layer)
            
            # Check layer format
            if len(layer_array.shape) == 3:  # RGB/RGBA format
                # If RGBA, check alpha channel
                if layer_array.shape[2] == 4:
                    # Use alpha channel as mask
                    alpha_channel = layer_array[:, :, 3]
                    # Set non-transparent areas (alpha > 0) to white
                    mask_array = np.where(alpha_channel > 0, 255, 0).astype(np.uint8)
                else:
                    # RGB format, check if not pure black (0,0,0)
                    # Assume drawn areas are non-black
                    non_black = np.any(layer_array > 0, axis=2)
                    mask_array = np.where(non_black, 255, 0).astype(np.uint8)
            elif len(layer_array.shape) == 2:  # Grayscale
                # Use grayscale values directly, set non-zero areas to white
                mask_array = np.where(layer_array > 0, 255, 0).astype(np.uint8)
            else:
                continue
            
            # Convert mask_array to PIL image and merge into total mask
            layer_mask = Image.fromarray(mask_array, mode='L')
            # Resize to match original image
            if layer_mask.size != base_image.size:
                layer_mask = layer_mask.resize(base_image.size, Image.LANCZOS)
            
            # Merge masks (use maximum value to ensure all drawn areas are included)
            mask_array_current = np.array(mask)
            layer_mask_array = np.array(layer_mask)
            combined_mask_array = np.maximum(mask_array_current, layer_mask_array)
            mask = Image.fromarray(combined_mask_array, mode='L')
    
    return mask


def upload_mask_image_r2(client_ip, time_id, mask_image):
    """
    Upload mask image to R2 directly from memory
    
    Args:
        client_ip (str): Client IP
        time_id (int): Timestamp
        mask_image (PIL.Image): Mask image
        
    Returns:
        str: Uploaded URL
    """
    # Generate unique filename using UUID to prevent file conflicts in concurrent environment
    unique_id = str(uuid.uuid4())
    file_name = f"mask_img_{unique_id}_{time_id}.png"
    
    try:
        # Convert mask image to bytes
        img_buffer = io.BytesIO()
        mask_image.save(img_buffer, format='PNG')
        img_data = img_buffer.getvalue()
        
        # Upload directly from memory
        res = R2Api().upload_from_memory(img_data, file_name, 'image/png')
        
        return res
    except Exception as e:
        print(f"Failed to upload mask image: {e}")
        return None



def submit_image_edit_task(user_image_url, prompt, task_type="80", mask_image_url="", reference_image_url=""):
    """
    Submit image editing task with improved error handling using API v2
    """
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {APIKEY}'
    }
    
    data = {
        "user_image": user_image_url,
        "user_mask": mask_image_url,
        "type": task_type,
        "text": prompt,
        "user_uuid": APIKEY,
        "priority": 0,
        "secret_key": "219ngu"
    }

    if reference_image_url:
        data["user_image2"] = reference_image_url
    
    retry_count = 0
    max_retries = 3
    
    while retry_count < max_retries:
        try:
            response = requests.post(
                f'{UKAPIURL}/public_image_edit_v2',
                headers=headers,
                json=data,
                timeout=30  # 增加超时时间
            )
            
            if response.status_code == 200:
                result = response.json()
                if result.get('code') == 0:
                    return result['data']['task_id'], None
                else:
                    return None, f"API Error: {result.get('message', 'Unknown error')}"
            elif response.status_code in [502, 503, 504]:  # 服务器错误，可以重试
                retry_count += 1
                if retry_count < max_retries:
                    print(f"⚠️ Server error {response.status_code}, retrying {retry_count}/{max_retries}")
                    time.sleep(2)  # 等待2秒后重试
                    continue
                else:
                    return None, f"HTTP Error after {max_retries} retries: {response.status_code}"
            else:
                return None, f"HTTP Error: {response.status_code}"
                
        except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
            retry_count += 1
            if retry_count < max_retries:
                print(f"⚠️ Network error, retrying {retry_count}/{max_retries}: {e}")
                time.sleep(2)
                continue
            else:
                return None, f"Network error after {max_retries} retries: {str(e)}"
        except Exception as e:
            return None, f"Request Exception: {str(e)}"
    
    return None, f"Failed after {max_retries} retries"


def check_task_status(task_id):
    """
    Query task status with improved error handling using API v2
    """
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {APIKEY}'
    }
    
    data = {
        "task_id": task_id
    }
    
    retry_count = 0
    max_retries = 2  # 状态查询重试次数少一些
    
    while retry_count < max_retries:
        try:
            response = requests.post(
                f'{UKAPIURL}/status_image_edit_v2',
                headers=headers,
                json=data,
                timeout=15  # 状态查询超时时间短一些
            )
            
            if response.status_code == 200:
                result = response.json()
                if result.get('code') == 0:
                    task_data = result['data']
                    status = task_data['status']
                    image_url = task_data.get('image_url')
                    
                    # Extract and log queue information for better user feedback
                    queue_info = task_data.get('queue_info', {})
                    if queue_info:
                        tasks_ahead = queue_info.get('tasks_ahead', 0)
                        current_priority = queue_info.get('current_priority', 0)
                        description = queue_info.get('description', '')
                        # print(f"📊 Queue Status - Tasks ahead: {tasks_ahead}, Priority: {current_priority}, Status: {status}")
                    
                    return status, image_url, task_data
                else:
                    return 'error', None, result.get('message', 'Unknown error')
            elif response.status_code in [502, 503, 504]:  # 服务器错误，可以重试
                retry_count += 1
                if retry_count < max_retries:
                    print(f"⚠️ Status check server error {response.status_code}, retrying {retry_count}/{max_retries}")
                    time.sleep(1)  # 状态查询重试间隔短一些
                    continue
                else:
                    return 'error', None, f"HTTP Error after {max_retries} retries: {response.status_code}"
            else:
                return 'error', None, f"HTTP Error: {response.status_code}"
                
        except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
            retry_count += 1
            if retry_count < max_retries:
                print(f"⚠️ Status check network error, retrying {retry_count}/{max_retries}: {e}")
                time.sleep(1)
                continue
            else:
                return 'error', None, f"Network error after {max_retries} retries: {str(e)}"
        except Exception as e:
            return 'error', None, f"Request Exception: {str(e)}"
    
    return 'error', None, f"Failed after {max_retries} retries"


def process_image_edit(img_input, prompt, reference_image=None, progress_callback=None):
    """
    Complete process for image editing
    
    Args:
        img_input: Can be file path (str) or PIL Image object
        prompt: Editing instructions
        progress_callback: Progress callback function
    """
    try:
        # Generate client IP and timestamp
        client_ip = "127.0.0.1"  # Default IP
        time_id = int(time.time())
        
        # Process input image - supports PIL Image and file path
        if hasattr(img_input, 'save'):  # PIL Image object
            pil_image = img_input
            print(f"💾 Using PIL Image directly from memory")
        else:
            # Load from file path
            pil_image = Image.open(img_input)
            print(f"📁 Loaded image from file: {img_input}")
        
        if progress_callback:
            progress_callback("uploading image...")
        
        # Upload user image directly from memory
        uploaded_url = upload_user_img_r2(client_ip, time_id, pil_image)
        if not uploaded_url:
            return None, None, "image upload failed", None
        
        # Extract actual image URL from upload URL
        if "?" in uploaded_url:
            uploaded_url = uploaded_url.split("?")[0]
        
        if progress_callback:
            progress_callback("submitting edit task...")
        
        reference_url = ""
        if reference_image is not None:
            try:
                if progress_callback:
                    progress_callback("uploading reference image...")

                if hasattr(reference_image, 'save'):
                    reference_pil = reference_image
                else:
                    reference_pil = Image.open(reference_image)

                reference_url = upload_user_img_r2(client_ip, time_id, reference_pil)
                if not reference_url:
                    return None, None, "reference image upload failed", None

                if "?" in reference_url:
                    reference_url = reference_url.split("?")[0]
            except Exception as e:
                return None, None, f"reference image processing failed: {str(e)}", None

        # Submit image editing task
        task_id, error = submit_image_edit_task(uploaded_url, prompt, reference_image_url=reference_url)
        if error:
            return None, None, error, None
        
        if progress_callback:
            progress_callback(f"task submitted, ID: {task_id}, processing...")
        
        # Wait for task completion
        max_attempts = 60  # Wait up to 10 minutes
        task_uuid = None
        for attempt in range(max_attempts):
            status, output_url, task_data = check_task_status(task_id)
            
            # Extract task_uuid from task_data
            if task_data and isinstance(task_data, dict):
                task_uuid = task_data.get('uuid', None)
            
            if status == 'completed':
                if output_url:
                    return uploaded_url, output_url, "image edit completed", task_uuid
                else:
                    return None, None, "Task completed but no result image returned", task_uuid
            elif status == 'error' or status == 'failed':
                return None, None, f"task processing failed: {task_data}", task_uuid
            elif status in ['queued', 'processing', 'running', 'created', 'working']:
                # Enhanced progress message with queue info and website promotion
                if progress_callback and task_data and isinstance(task_data, dict):
                    queue_info = task_data.get('queue_info', {})
                    if queue_info and status in ['queued', 'created']:
                        tasks_ahead = queue_info.get('tasks_ahead', 0)
                        current_priority = queue_info.get('current_priority', 0)
                        if tasks_ahead > 0:
                            progress_callback(f"⏳ Queue: {tasks_ahead} tasks ahead | Low priority | Visit website for instant processing → https://omnicreator.net/#generator")
                        else:
                            progress_callback(f"🚀 Processing your image editing request...")
                    elif status == 'processing':
                        progress_callback(f"🎨 AI is processing... Please wait")
                    elif status in ['running', 'working']:
                        progress_callback(f"⚡ Generating... Almost done")
                    else:
                        progress_callback(f"📋 Task status: {status}")
                else:
                    if progress_callback:
                        progress_callback(f"task processing... (status: {status})")
                time.sleep(1)
            else:
                if progress_callback:
                    progress_callback(f"unknown status: {status}")
                time.sleep(1)
        
        return None, None, "task processing timeout", task_uuid
        
    except Exception as e:
        return None, None, f"error occurred during processing: {str(e)}", None


def process_local_image_edit(base_image, layers, prompt, reference_image=None, progress_callback=None, use_example_mask=None):
    """
    处理局部图片编辑的完整流程
    
    Args:
        base_image (PIL.Image): 原始图片
        layers (list): ImageEditor的层数据
        prompt (str): 编辑指令
        progress_callback: 进度回调函数
    """
    try:
        # Generate client IP and timestamp
        client_ip = "127.0.0.1"  # Default IP
        time_id = int(time.time())
        
        if progress_callback:
            progress_callback("creating mask image...")
        
        # Check if we should use example mask (backdoor for example case)
        if use_example_mask:
            # Load local mask file for example
            try:
                from PIL import Image
                import os
                
                # Check if base_image is valid
                if base_image is None:
                    return None, None, "Base image is None, cannot process example mask", None
                
                if os.path.exists(use_example_mask):
                    mask_image = Image.open(use_example_mask)
                    
                    # Ensure mask has same size as base image
                    if hasattr(base_image, 'size') and mask_image.size != base_image.size:
                        mask_image = mask_image.resize(base_image.size)
                    
                    # Ensure mask is in L mode (grayscale)
                    if mask_image.mode != 'L':
                        mask_image = mask_image.convert('L')
                    
                    print(f"🎭 Using example mask from: {use_example_mask}, size: {mask_image.size}")
                else:
                    return None, None, f"Example mask file not found: {use_example_mask}", None
            except Exception as e:
                import traceback
                traceback.print_exc()
                return None, None, f"Failed to load example mask: {str(e)}", None
        else:
            # Normal case: create mask from layers
            mask_image = create_mask_from_layers(base_image, layers)
            
            # 检查mask是否有内容
            mask_array = np.array(mask_image)
            if np.max(mask_array) == 0:
                return None, None, "please draw mask", None
        
        # Print mask statistics
        if not use_example_mask:
            print(f"📝 创建mask图片成功，绘制区域像素数: {np.sum(mask_array > 0)}")
        else:
            mask_array = np.array(mask_image)
            print(f"🎭 Example mask loaded successfully, mask pixels: {np.sum(mask_array > 0)}")
        
        if progress_callback:
            progress_callback("uploading original image...")
        
        # 直接从内存上传原始图片
        uploaded_url = upload_user_img_r2(client_ip, time_id, base_image)
        if not uploaded_url:
            return None, None, "original image upload failed", None
        
        # 从上传 URL 中提取实际的图片 URL
        if "?" in uploaded_url:
            uploaded_url = uploaded_url.split("?")[0]
        
        if progress_callback:
            progress_callback("uploading mask image...")
        
        # 直接从内存上传mask图片
        mask_url = upload_mask_image_r2(client_ip, time_id, mask_image)
        if not mask_url:
            return None, None, "mask image upload failed", None
        
        # 从上传 URL 中提取实际的图片 URL
        if "?" in mask_url:
            mask_url = mask_url.split("?")[0]
        
        reference_url = ""
        if reference_image is not None:
            try:
                if progress_callback:
                    progress_callback("uploading reference image...")

                if hasattr(reference_image, 'save'):
                    reference_pil = reference_image
                else:
                    reference_pil = Image.open(reference_image)

                reference_url = upload_user_img_r2(client_ip, time_id, reference_pil)
                if not reference_url:
                    return None, None, "reference image upload failed", None

                if "?" in reference_url:
                    reference_url = reference_url.split("?")[0]
            except Exception as e:
                return None, None, f"reference image processing failed: {str(e)}", None

        print(f"📤 图片上传成功:")
        print(f"   原始图片: {uploaded_url}")
        print(f"   Mask图片: {mask_url}")
        if reference_url:
            print(f"   参考图片: {reference_url}")
        
        if progress_callback:
            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,
            reference_image_url=reference_url
        )
        if error:
            return None, None, error, None
        
        if progress_callback:
            progress_callback(f"task submitted, ID: {task_id}, processing...")
        
        print(f"🚀 局部编辑任务已提交，任务ID: {task_id}")
        
        # Wait for task completion
        max_attempts = 60  # Wait up to 10 minutes
        task_uuid = None
        for attempt in range(max_attempts):
            status, output_url, task_data = check_task_status(task_id)
            
            # Extract task_uuid from task_data
            if task_data and isinstance(task_data, dict):
                task_uuid = task_data.get('uuid', None)
            
            if status == 'completed':
                if output_url:
                    print(f"✅ 局部编辑任务完成，结果: {output_url}")
                    return uploaded_url, output_url, "local image edit completed", task_uuid
                else:
                    return None, None, "task completed but no result image returned", task_uuid
            elif status == 'error' or status == 'failed':
                return None, None, f"task processing failed: {task_data}", task_uuid
            elif status in ['queued', 'processing', 'running', 'created', 'working']:
                # Enhanced progress message with queue info and website promotion
                if progress_callback and task_data and isinstance(task_data, dict):
                    queue_info = task_data.get('queue_info', {})
                    if queue_info and status in ['queued', 'created']:
                        tasks_ahead = queue_info.get('tasks_ahead', 0)
                        current_priority = queue_info.get('current_priority', 0)
                        if tasks_ahead > 0:
                            progress_callback(f"⏳ Queue: {tasks_ahead} tasks ahead | Low priority | Visit website for instant processing → https://omnicreator.net/#generator")
                        else:
                            progress_callback(f"🚀 Processing your local editing request...")
                    elif status == 'processing':
                        progress_callback(f"🎨 AI is processing... Please wait")
                    elif status in ['running', 'working']:
                        progress_callback(f"⚡ Generating... Almost done")
                    else:
                        progress_callback(f"📋 Task status: {status}")
                else:
                    if progress_callback:
                        progress_callback(f"processing... (status: {status})")
                time.sleep(1)  # Wait 1 second before retry
            else:
                if progress_callback:
                    progress_callback(f"unknown status: {status}")
                time.sleep(1)
        
        return None, None, "task processing timeout", task_uuid
        
    except Exception as e:
        print(f"❌ 局部编辑处理异常: {str(e)}")
        return None, None, f"error occurred during processing: {str(e)}", None


def download_and_check_result_nsfw(image_url, nsfw_detector=None):
    """
    下载结果图片并进行NSFW检测
    
    Args:
        image_url (str): 结果图片URL
        nsfw_detector: NSFW检测器实例
        
    Returns:
        tuple: (is_nsfw, error_message)
    """
    if nsfw_detector is None:
        return False, None
    
    try:
        # 下载图片
        response = requests.get(image_url, timeout=30)
        if response.status_code != 200:
            return False, f"Failed to download result image: HTTP {response.status_code}"
        
        # 将图片数据转换为PIL Image
        image_data = io.BytesIO(response.content)
        result_image = Image.open(image_data)
        
        # 进行NSFW检测
        nsfw_result = nsfw_detector.predict_pil_label_only(result_image)
        
        is_nsfw = nsfw_result.lower() == "nsfw"
        print(f"🔍 结果图片NSFW检测: {'❌❌❌ ' + nsfw_result if is_nsfw else '✅✅✅ ' + nsfw_result}")
        
        return is_nsfw, None
        
    except Exception as e:
        print(f"⚠️ 结果图片NSFW检测失败: {e}")
        return False, f"Failed to check result image: {str(e)}"


if __name__ == "__main__":
    
    pass