feat: implement Direct-to-R2 upload on worker to bypass Base64 bottleneck

Author: effame

requirements.txt

@@ -2,3 +2,4 @@ numpy<2
 opencv-python-headless==4.10.0.84
 Pillow==10.4.0
 tqdm==4.64.0
+boto3==1.34.106

rp_handler.py

@@ -2,12 +2,56 @@
 import sys
 import base64
 import urllib.request
+import boto3
+from botocore.config import Config
 
 import cv2
 import numpy as np
 import torch
 import runpod
 
+# Initialize S3 client for R2 (Lazy)
+s3_client = None
+
+def get_s3_client():
+    global s3_client
+    if s3_client is None:
+        access_key = os.environ.get("R2_ACCESS_KEY_ID")
+        secret_key = os.environ.get("R2_SECRET_ACCESS_KEY")
+        endpoint = os.environ.get("R2_ENDPOINT")
+        if all([access_key, secret_key, endpoint]):
+            s3_client = boto3.client(
+                "s3",
+                endpoint_url=endpoint,
+                aws_access_key_id=access_key,
+                aws_secret_access_key=secret_key,
+                config=Config(signature_version="s3v4"),
+                region_name="auto",
+            )
+    return s3_client
+
+def upload_to_r2(buffer, key, content_type="image/jpeg"):
+    client = get_s3_client()
+    if client is None:
+        return None
+    
+    bucket = os.environ.get("R2_BUCKET_NAME", "stockgen-ai")
+    try:
+        client.put_object(
+            Bucket=bucket,
+            Key=key,
+            Body=buffer,
+            ContentType=content_type,
+            CacheControl="public, max-age=31536000"
+        )
+        custom_domain = os.environ.get("R2_CUSTOM_DOMAIN")
+        if custom_domain:
+            return f"{custom_domain.rstrip('/')}/{key}"
+        return f"{os.environ.get('R2_ENDPOINT').rstrip('/')}/{bucket}/{key}"
+    except Exception as e:
+        print(f"❌ R2 Upload Failed: {str(e)}")
+        return None
+
 # Patch basicsr compatibility with newer PyTorch/torchvision
 import torchvision.transforms.functional as F_tv
 import types
@@ -240,13 +284,24 @@ def handler(job):
     if not success:
         return {"error": "Failed to encode output image"}
 
+    # Direct-to-R2 Upload Logic
+    r2_url = None
+    r2_key = job_input.get("r2_key") # Target path like "users/uid/batches/bid/jid.jpg"
+    
+    if r2_key:
+        print(f"☁️ [Direct-to-R2] Uploading result to: {r2_key}")
+        r2_url = upload_to_r2(encoded_img.tobytes(), r2_key)
+        if r2_url:
+            print(f"✅ [Direct-to-R2] Success: {r2_url}")
+
     b64 = base64.b64encode(encoded_img).decode("utf-8")
 
     h, w = img.shape[:2]
     oh, ow = output.shape[:2]
 
     return {
-        "image": b64,
+        "image": b64 if not r2_url else None, # Skip Base64 if R2 success to save bandwidth
+        "r2_url": r2_url,
         "image_format": "jpg",
         "model": model_name,
         "face_enhance_applied": face_enhance,