Update app.py

app.py

@@ -1,5 +1,5 @@
 from fastapi import FastAPI, Request
-import requests, os, random
+import requests, os, random, httpx, asyncio
 import psycopg2
 import logging
 from fastapi import FastAPI, HTTPException
@@ -73,6 +73,18 @@ def insert_batch(prompt: str, width: int, height: int, model: str, urls: list[st
         if conn is not None:
             conn.close()
 
+async def upload_to_imagebb(image_url):
+        upload_url = "https://api.imgbb.com/1/upload"
+        params = {
+            "key": IMGBB_API_KEY,
+            "image": image_url,
+        }
+        async with httpx.AsyncClient() as client:
+            response = await client.post(upload_url, params=params)
+            if response.is_success:
+                return response.json()["data"]["url"]
+            return None
+
 @app.post("/generate-image")
 async def generate_image(request: ImageRequest):
     try:
@@ -90,9 +102,25 @@ async def generate_image(request: ImageRequest):
         images = await runware.imageInference(requestImage=request_image)
         image_urls = [image.imageURL for image in images]
         
-        batch_id = insert_batch(request.prompt, request.width, request.height, request.model, image_urls)
+        # Upload images to ImgBB concurrently
+        upload_tasks = [upload_to_imagebb(url) for url in image_urls]
+        imgbb_urls = await asyncio.gather(*upload_tasks)
         
-        response = {"batch": {"id": batch_id, "prompt": request.prompt, "width": request.width, "height": request.height, "model": request.model, "images": [{"url": url} for url in image_urls]}}
+        # Filter out any failed uploads (None values)
+        imgbb_urls = [url for url in imgbb_urls if url is not None]
+
+        batch_id = insert_batch(request.prompt, request.width, request.height, request.model, imgbb_urls)
+        
+        response = {
+            "batch": {
+                "id": batch_id,
+                "prompt": request.prompt,
+                "width": request.width,
+                "height": request.height,
+                "model": request.model,
+                "images": [{"url": url} for url in imgbb_urls]
+            }
+        }
         return response
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed to generate image: {str(e)}")