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]



samples = [
    {
        "image": ["datas/data01/hei_cat01.webp","datas/data01/tom.webp"],
        "prompt": "Let all cats dance together in the park",
        "aspect_ratio": "16:9"
    },
    {
        "image": ["datas/data02/girl.jpg","datas/data02/cloth.jpeg"],
        "prompt": "Let the girl in first image, wear the dress in second image",
        "aspect_ratio": "9:16"
    },
    {
        "image": ["datas/data03/girl.webp","datas/data03/cloth.jpg"],
        "prompt": "Let the girl in first image, wear the bikini in second image, lying on the beach",
        "aspect_ratio": "1:1"
    }
]



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"QwenImageEdit/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_file(self, local_path, cloud_path):
        t1 = time.time()
        head_dict = {
            'jpg': 'image/jpeg',
            'jpeg': 'image/jpeg',
            'png': 'image/png',
            'gif': 'image/gif',
            'bmp': 'image/bmp',
            'webp': 'image/webp',
            'ico': 'image/x-icon'
        }
        ftype = os.path.basename(local_path).split(".")[-1].lower()
        ctype = head_dict.get(ftype, 'application/octet-stream')
        headers = {"Content-Type": ctype}


        cloud_path = f"ImageEdit/Uploads/{str(datetime.date.today())}/{os.path.basename(local_path)}"
        url = self.client.generate_presigned_url(
            "put_object",
            Params={"Bucket": self.R2_BUCKET, "Key": cloud_path, "ContentType": ctype},
            ExpiresIn=604800
        )

        retry_count = 0
        while retry_count < 3:
            try:
                with open(local_path, 'rb') as f:
                    self.session.put(url, data=f.read(), headers=headers, timeout=8)
                break
            except (requests.exceptions.Timeout, requests.exceptions.RequestException):
                retry_count += 1
                if retry_count == 3:
                    raise Exception('Failed to upload file to R2 after 3 retries!')
                continue
        print("upload_file 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="", user_image2_url="", user_image3_url="", width=0, height=0):
    """
    Submit image editing task with improved error handling using API v2
    Supports multi-image editing with user_image2, user_image3, width, height parameters
    """
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {APIKEY}'
    }
    
    data = {
        "user_image": user_image_url,
        "user_image2": user_image2_url,
        "user_image3": user_image3_url,
        "user_mask": mask_image_url,
        "type": task_type,
        "width": width,
        "height": height,
        "text": prompt,
        "user_uuid": APIKEY,
        "priority": 0,
        "secret_key": "219ngu"
    }
    
    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
                    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)
                        print(f"📊 Queue Info - Tasks ahead: {tasks_ahead}, Priority: {current_priority}")
                    
                    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_multi_image_edit(img_inputs, prompt, width=0, height=0, progress_callback=None):
    """
    Complete process for multi-image editing
    
    Args:
        img_inputs: List of image inputs (PIL Image objects or file paths), 2-3 images
        prompt: Editing instructions
        width: Output width (0 for auto)
        height: Output height (0 for auto)
        progress_callback: Progress callback function
    """
    try:
        # Generate client IP and timestamp
        client_ip = "127.0.0.1"  # Default IP
        time_id = int(time.time())
        
        # Validate input images
        if not img_inputs or len(img_inputs) < 2:
            return None, "Please upload at least 2 images", None
        
        if len(img_inputs) > 3:
            return None, "Maximum 3 images allowed", None
        
        # Process input images - supports PIL Image and file path
        uploaded_urls = []
        
        if progress_callback:
            progress_callback("uploading images...")
        
        for i, img_input in enumerate(img_inputs):
            if img_input is None:
                continue
                
            if hasattr(img_input, 'save'):  # PIL Image object
                pil_image = img_input
                print(f"💾 Using PIL Image {i+1} directly from memory")
            else:
                # Load from file path
                pil_image = Image.open(img_input)
                print(f"📁 Loaded image {i+1} from file: {img_input}")
            
            # Upload user image directly from memory
            uploaded_url = upload_user_img_r2(client_ip, time_id + i, pil_image)
            if not uploaded_url:
                return None, f"Image {i+1} upload failed", None
            
            # Extract actual image URL from upload URL
            if "?" in uploaded_url:
                uploaded_url = uploaded_url.split("?")[0]
            
            uploaded_urls.append(uploaded_url)
        
        # Ensure we have the required URLs
        user_image_url = uploaded_urls[0] if len(uploaded_urls) > 0 else ""
        user_image2_url = uploaded_urls[1] if len(uploaded_urls) > 1 else ""
        user_image3_url = uploaded_urls[2] if len(uploaded_urls) > 2 else ""
        
        if progress_callback:
            progress_callback("submitting multi-image edit task...")
        
        # Submit multi-image editing task
        task_id, error = submit_image_edit_task(
            user_image_url, 
            prompt, 
            task_type="80", 
            mask_image_url="",
            user_image2_url=user_image2_url,
            user_image3_url=user_image3_url,
            width=width,
            height=height
        )
        if error:
            return 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 output_url, "multi-image edit completed", task_uuid
                else:
                    return None, "Task completed but no result image returned", task_uuid
            elif status == 'error' or status == 'failed':
                return None, f"task processing failed: {task_data}", task_uuid
            elif status in ['queued', 'processing', 'running', 'created', 'working']:
                # Enhanced progress message with queue info
                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)
                        if tasks_ahead > 0:
                            progress_callback(f"⏳ In queue: {tasks_ahead} tasks ahead. Visit https://omnicreator.net/multi-image-edit#generator for instant processing!")
                        else:
                            progress_callback(f"🚀 Processing your multi-image request...")
                    else:
                        progress_callback(f"multi-image processing... (status: {status})")
                else:
                    if progress_callback:
                        progress_callback(f"multi-image processing... (status: {status})")
                time.sleep(1)
            else:
                if progress_callback:
                    progress_callback(f"unknown status: {status}")
                time.sleep(1)
        
        return None, "task processing timeout", task_uuid
        
    except Exception as e:
        return None, f"error occurred during processing: {str(e)}", None


def process_image_edit(img_input, prompt, progress_callback=None):
    """
    Complete process for single image editing (backward compatibility)
    
    Args:
        img_input: Can be file path (str) or PIL Image object
        prompt: Editing instructions
        progress_callback: Progress callback function
    """
    return process_multi_image_edit([img_input], prompt, 0, 0, progress_callback)


def process_local_image_edit(base_image, layers, prompt, progress_callback=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...")
        
        # 从layers创建mask图片
        mask_image = create_mask_from_layers(base_image, layers)
        
        # 检查mask是否有内容
        mask_array = np.array(mask_image)
        if np.max(mask_array) == 0:
            return None, "please draw mask", None
        
        print(f"📝 创建mask图片成功，绘制区域像素数: {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, "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, "mask image upload failed", None
        
        # 从上传 URL 中提取实际的图片 URL
        if "?" in mask_url:
            mask_url = mask_url.split("?")[0]
        
        print(f"📤 图片上传成功:")
        print(f"   原始图片: {uploaded_url}")
        print(f"   Mask图片: {mask_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)
        if error:
            return 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 output_url, "local image edit completed", task_uuid
                else:
                    return None, "task completed but no result image returned", task_uuid
            elif status == 'error' or status == 'failed':
                return None, f"task processing failed: {task_data}", task_uuid
            elif status in ['queued', 'processing', 'running', 'created', 'working']:
                # Enhanced progress message with queue info
                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)
                        if tasks_ahead > 0:
                            progress_callback(f"⏳ In queue: {tasks_ahead} tasks ahead. Visit https://omnicreator.net/multi-image-edit#generator for instant processing!")
                        else:
                            progress_callback(f"🚀 Processing your local edit request...")
                    else:
                        progress_callback(f"processing... (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, "task processing timeout", task_uuid
        
    except Exception as e:
        print(f"❌ 局部编辑处理异常: {str(e)}")
        return 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