Update sd.py

sd.py

@@ -5,8 +5,10 @@ from datetime import datetime
 import io
 import yaml
 
+# Configuration
 IMAGE_GENERATION_URL = os.getenv("SD_API_URL")
 CACHE_DIR = '/tmp/cache'  # Use /tmp directory which is writable
+IMGBB_API_KEY = os.getenv("IMGBB_API_KEY")  # Replace with your ImgBB API key
 
 # Ensure the cache directory exists
 os.makedirs(CACHE_DIR, exist_ok=True)
@@ -15,6 +17,41 @@ def load_config():
     with open('config.yaml', 'r') as f:
         return yaml.safe_load(f)
 
+def upload_to_imgbb(image_data):
+    """
+    Uploads image data to ImgBB and returns the URL.
+
+    Args:
+        image_data (bytes): Raw image data as bytes.
+
+    Returns:
+        str: URL of the uploaded image on ImgBB.
+    """
+    try:
+        response = requests.post(
+            "https://api.imgbb.com/1/upload",
+            params={
+                "key": IMGBB_API_KEY,
+                "expiration": 600  # Optional: Set expiration time in seconds
+            },
+            files={
+                "image": ("image.jpg", image_data)
+            }
+        )
+        response.raise_for_status()
+        result = response.json()
+        if result["data"] and "url" in result["data"]:
+            return result["data"]["url"]
+        else:
+            print("Failed to upload image to ImgBB.")
+            return None
+    except requests.RequestException as e:
+        print(f"Error uploading image to ImgBB: {e}")
+        return None
+    except Exception as e:
+        print(f"Unexpected error uploading image to ImgBB: {e}")
+        return None
+
 def generate_sd(prompt):
     """
     Generate an image using the specified prompt and save it to the cache directory.
@@ -22,9 +59,11 @@ def generate_sd(prompt):
     Returns:
         image_data (bytes): Raw image data as bytes.
         image_path (str): Absolute path to the saved image file using forward slashes.
+        image_url (str): URL of the uploaded image on ImgBB.
     """
     config = load_config()
     config = config['config']
+    
     # Define payload for the POST request
     payload = {
         "prompt": prompt,
@@ -58,19 +97,21 @@ def generate_sd(prompt):
 
             # Save the image as a JPEG file
             with open(image_path, 'wb') as f:
-                image = io.BytesIO(image_data)
-                f.write(image.getbuffer())
+                f.write(image_data)
+
+            # Upload image to ImgBB
+            image_url = upload_to_imgbb(image_data)
 
             # Convert the path to use forward slashes
             image_path = image_path.replace('\\', '/')
 
-            return image_data, image_path
+            return image_data, image_path, image_url
         else:
             print("No images found in the response.")
-            return None, None
+            return None, None, None
     except requests.RequestException as e:
         print(f"HTTP request error: {e}")
-        return None, None
+        return None, None, None
     except Exception as e:
         print(f"Error generating image: {e}")
-        return None, None
+        return None, None, None