upload edited images to imgbb

flux_kontext_lib/image_generator.py

@@ -1,6 +1,7 @@
 import requests
 import base64
 import json
+import os
 from io import BytesIO
 from typing import Optional, Dict
 
@@ -12,6 +13,41 @@ except ImportError:
     print("Pillow library not found. Please install it using: pip install Pillow")
     exit()
 
+# --- Configuration ---
+IMGBB_API_KEY = os.getenv("IMGBB_API_KEY")
+
+def upload_to_imgbb(image_path: str, file_name: str) -> Optional[str]:
+    """
+    Uploads the image at image_path to ImgBB.
+    Returns the public URL or None on failure.
+    """
+    if not IMGBB_API_KEY:
+        print("Warning: IMGBB_API_KEY not set, skipping upload")
+        return None
+
+    print(f"Uploading {file_name} to ImgBB...")
+    try:
+        with open(image_path, 'rb') as f:
+            files = {"image": (file_name, f.read())}
+        resp = requests.post(
+            "https://api.imgbb.com/1/upload",
+            params={"key": IMGBB_API_KEY},
+            files=files,
+            timeout=20
+        )
+        resp.raise_for_status()
+        data = resp.json().get("data", {})
+        url = data.get("url")
+        if url:
+            print(f"Successfully uploaded to ImgBB: {url}")
+            return url
+        else:
+            print(f"Error: ImgBB API response missing 'url'. Response: {resp.json()}")
+            return None
+    except requests.exceptions.RequestException as e:
+        print(f"Error: ImgBB upload failed: {e}")
+        return None
+
 def _download_image_from_url(image_url: str, save_path: str) -> bool:
     """
     Downloads an image from a URL and saves it to a local path.
@@ -25,7 +61,7 @@ def _download_image_from_url(image_url: str, save_path: str) -> bool:
     """
     print(f"Downloading generated image from: {image_url}")
     try:
-        image_response = requests.get(image_url, stream=True)
+        image_response = requests.get(image_url, stream=True, timeout=30)
         # Check if the download request was successful.
         if image_response.status_code == 200:
             content_type = image_response.headers.get('Content-Type', '')
@@ -51,7 +87,8 @@ def generate_image(
     download_path: Optional[str] = None
 ) -> Optional[Dict]:
     """
-    Sends a request to the image generation API and optionally downloads the result.
+    Sends a request to the image generation API, optionally downloads the result,
+    and uploads it to ImgBB.
 
     Args:
         prompt_text: The instructional text for image modification.
@@ -59,7 +96,8 @@ def generate_image(
         download_path: If provided, the path to save the generated image.
 
     Returns:
-        A dictionary of the JSON response from the API, or None on error.
+        A dictionary of the JSON response from the API, including the ImgBB URL,
+        or None on error.
     """
     url = "https://kontext-chat.replicate.dev/generate-image"
 
@@ -101,19 +139,28 @@ def generate_image(
     }
 
     try:
-        response = requests.post(url, headers=headers, data=json.dumps(payload))
+        response = requests.post(url, headers=headers, data=json.dumps(payload), timeout=45)
         response.raise_for_status()
         api_response_data = response.json()
 
-        # --- Optional Download Logic ---
+        imgbb_url = None  # Initialize imgbb_url
+
+        # --- Download & Upload Logic ---
         if download_path and isinstance(api_response_data, dict):
             image_url = api_response_data.get("imageUrl")
             if image_url:
-                _download_image_from_url(image_url, download_path)
+                if _download_image_from_url(image_url, download_path):
+                    # If download is successful, upload the saved file to ImgBB
+                    file_name = os.path.basename(download_path)
+                    imgbb_url = upload_to_imgbb(download_path, file_name)
             else:
-                print("Warning: 'imageUrl' not found in response, could not download image.")
+                print("Warning: 'imageUrl' not found in API response, cannot download or upload image.")
+        
+        # Add the ImgBB URL to the response dictionary if it was generated
+        if isinstance(api_response_data, dict):
+            api_response_data['imgbb_url'] = imgbb_url
 
         return api_response_data
     except requests.exceptions.RequestException as e:
         print(f"API request failed: {e}")
-        return None
+        return None