init

util.py

@@ -12,6 +12,7 @@ import numpy as np
 import gradio as gr
 import boto3
 import tempfile
+import io
 from botocore.client import Config
 from PIL import Image
 
@@ -31,43 +32,45 @@ R2_SECRET_KEY = OneKey[5]
 R2_ENDPOINT = OneKey[6]
 
 
-tmpFolder = "tmp"
-os.makedirs(tmpFolder, exist_ok=True)
-
-
-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
+# 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:
@@ -90,6 +93,41 @@ class R2Api:
 
         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:
+                self.session.put(url, data=image_data, 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_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 = {
@@ -127,16 +165,29 @@ class R2Api:
         print("upload_file time is ====>", time.time() - t1)
         return f"{self.domain}{cloud_path}"
     
-def upload_user_img_r2(clientIp, timeId, img):
+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
+    """
     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)
-
-    res = R2Api().upload_file(local_path, fileName)
-
-    if os.path.exists(local_path):
-        os.remove(local_path)
+    
+    # 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
 
 
@@ -203,7 +254,7 @@ def create_mask_from_layers(base_image, layers):
 
 def upload_mask_image_r2(client_ip, time_id, mask_image):
     """
-    Upload mask image to R2
+    Upload mask image to R2 directly from memory
     
     Args:
         client_ip (str): Client IP
@@ -214,23 +265,20 @@ def upload_mask_image_r2(client_ip, time_id, mask_image):
         str: Uploaded URL
     """
     file_name = f"{client_ip.replace('.', '')}{time_id}_mask.png"
-    local_path = os.path.join(tmpFolder, file_name)
     
     try:
-        # Save mask image as PNG format (supports transparency)
-        mask_image.save(local_path, 'PNG')
+        # Convert mask image to bytes
+        img_buffer = io.BytesIO()
+        mask_image.save(img_buffer, format='PNG')
+        img_data = img_buffer.getvalue()
         
-        # Upload to R2
-        res = R2Api().upload_file(local_path, file_name)
+        # 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
-    finally:
-        # Clean up local files
-        if os.path.exists(local_path):
-            os.remove(local_path)
 
 
 
@@ -313,7 +361,6 @@ def process_image_edit(img_input, prompt, progress_callback=None):
         prompt: Editing instructions
         progress_callback: Progress callback function
     """
-    temp_img_path = None
     try:
         # Generate client IP and timestamp
         client_ip = "127.0.0.1"  # Default IP
@@ -321,26 +368,18 @@ def process_image_edit(img_input, prompt, progress_callback=None):
         
         # Process input image - supports PIL Image and file path
         if hasattr(img_input, 'save'):  # PIL Image object
-            # Create temporary file
-            temp_dir = tempfile.mkdtemp()
-            temp_img_path = os.path.join(temp_dir, f"temp_img_{time_id}.jpg")
-            
-            # Save PIL Image as temporary file
-            if img_input.mode != 'RGB':
-                img_input = img_input.convert('RGB')
-            img_input.save(temp_img_path, 'JPEG', quality=95)
-            
-            img_path = temp_img_path
-            print(f"💾 PIL Image saved as temporary file: {temp_img_path}")
+            pil_image = img_input
+            print(f"💾 Using PIL Image directly from memory")
         else:
-            # Assume it's a file path
-            img_path = img_input
+            # 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
-        uploaded_url = upload_user_img_r2(client_ip, time_id, img_path)
+        # Upload user image directly from memory
+        uploaded_url = upload_user_img_r2(client_ip, time_id, pil_image)
         if not uploaded_url:
             return None, "image upload failed"
         
@@ -374,7 +413,7 @@ def process_image_edit(img_input, prompt, progress_callback=None):
             elif status in ['queued', 'processing', 'running', 'created', 'working']:
                 if progress_callback:
                     progress_callback(f"task processing... (status: {status})")
-                time.sleep(1)  # 等待10秒后重试
+                time.sleep(1)
             else:
                 if progress_callback:
                     progress_callback(f"unknown status: {status}")
@@ -384,19 +423,6 @@ def process_image_edit(img_input, prompt, progress_callback=None):
         
     except Exception as e:
         return None, f"error occurred during processing: {str(e)}"
-    
-    finally:
-        # 清理临时文件
-        if temp_img_path and os.path.exists(temp_img_path):
-            try:
-                os.remove(temp_img_path)
-                # 尝试删除临时目录（如果为空）
-                temp_dir = os.path.dirname(temp_img_path)
-                if os.path.exists(temp_dir):
-                    os.rmdir(temp_dir)
-                print(f"🗑️ Cleaned up temporary file: {temp_img_path}")
-            except Exception as cleanup_error:
-                print(f"⚠️ Failed to clean up temporary file: {cleanup_error}")
 
 
 def process_local_image_edit(base_image, layers, prompt, progress_callback=None):
@@ -409,9 +435,6 @@ def process_local_image_edit(base_image, layers, prompt, progress_callback=None)
         prompt (str): 编辑指令
         progress_callback: 进度回调函数
     """
-    temp_img_path = None
-    temp_mask_path = None
-    
     try:
         # Generate client IP and timestamp
         client_ip = "127.0.0.1"  # Default IP
@@ -433,17 +456,8 @@ def process_local_image_edit(base_image, layers, prompt, progress_callback=None)
         if progress_callback:
             progress_callback("uploading original image...")
         
-        # 处理并上传原始图片
-        temp_dir = tempfile.mkdtemp()
-        temp_img_path = os.path.join(temp_dir, f"temp_img_{time_id}.jpg")
-        
-        # 保存原始图片
-        if base_image.mode != 'RGB':
-            base_image = base_image.convert('RGB')
-        base_image.save(temp_img_path, 'JPEG', quality=95)
-        
-        # 上传原始图片
-        uploaded_url = upload_user_img_r2(client_ip, time_id, temp_img_path)
+        # 直接从内存上传原始图片
+        uploaded_url = upload_user_img_r2(client_ip, time_id, base_image)
         if not uploaded_url:
             return None, "original image upload failed"
         
@@ -454,7 +468,7 @@ def process_local_image_edit(base_image, layers, prompt, progress_callback=None)
         if progress_callback:
             progress_callback("uploading mask image...")
         
-        # 上传mask图片
+        # 直接从内存上传mask图片
         mask_url = upload_mask_image_r2(client_ip, time_id, mask_image)
         if not mask_url:
             return None, "mask image upload failed"
@@ -507,18 +521,6 @@ def process_local_image_edit(base_image, layers, prompt, progress_callback=None)
     except Exception as e:
         print(f"❌ 局部编辑处理异常: {str(e)}")
         return None, f"error occurred during processing: {str(e)}"
-    
-    finally:
-        # 清理临时文件
-        if temp_img_path and os.path.exists(temp_img_path):
-            try:
-                os.remove(temp_img_path)
-                temp_dir = os.path.dirname(temp_img_path)
-                if os.path.exists(temp_dir):
-                    os.rmdir(temp_dir)
-                print(f"🗑️ Cleaned up temporary file: {temp_img_path}")
-            except Exception as cleanup_error:
-                print(f"⚠️ Failed to clean up temporary file: {cleanup_error}")
 
 
 if __name__ == "__main__":