add: async image upload, nto optimize

app.py

@@ -22,4 +22,9 @@ app.add_middleware(
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
-)
+)
+
+if __name__ == '__main__':
+    import uvicorn
+
+    uvicorn.run(app)

src/api/nto_api.py

@@ -26,6 +26,9 @@ import replicate
 import requests
 from src.utils.logger import logger
 import secrets
+import aiohttp
+import asyncio
+import gc
 
 pipeline = Pipeline()
 
@@ -360,31 +363,64 @@ async def parse_necklace_try_on_id(necklaceImageId: str = Form(...),
     )
 
 
-def supabase_upload_and_return_url(prefix: str, image: Image.Image,quality: int = 85) :
-
+async def supabase_upload_and_return_url(prefix: str, image: Image.Image, quality: int = 85):
     try:
-        filename = f"{prefix}_{secrets.token_hex(16)}.webp"
-
-
+        filename = f"{prefix}_{secrets.token_hex(8)}.webp"
+
+        loop = asyncio.get_event_loop()
+        image_bytes = await loop.run_in_executor(
+            None,
+            process_image,
+            image,
+            quality
+        )
 
+        async with aiohttp.ClientSession() as session:
+            headers = {
+                "Authorization": f"Bearer {key}",
+                "Content-Type": "image/webp"
+            }
 
-        buffer = BytesIO()
-        image.save(buffer, format='WEBP', quality=quality, optimize=True)
-        image_bytes = buffer.getvalue()
+            upload_url = f"{url}/storage/v1/object/JewelMirrorOutputs/{filename}"
 
-        bucket.upload(
-            path=filename,
-            file=image_bytes,
-        )
+            async with session.post(
+                    upload_url,
+                    data=image_bytes,
+                    headers=headers
+            ) as response:
+                if response.status != 200:
+                    raise Exception(f"Upload failed with status {response.status}")
 
         return bucket.get_public_url(filename)
 
     except Exception as e:
-        print(f"Failed to upload image: {str(e)}")
+        logger.error(f"Failed to upload image: {str(e)}")
         return None
-    finally:
-        if 'buffer' in locals():
-            buffer.close()
+
+
+def process_image(image: Image.Image, quality: int) -> bytes:
+    try:
+        if image.mode in ['RGBA', 'P']:
+            image = image.convert('RGB')
+
+        max_size = 3000
+        if image.width > max_size or image.height > max_size:
+            ratio = min(max_size / image.width, max_size / image.height)
+            new_size = (int(image.width * ratio), int(image.height * ratio))
+            image = image.resize(new_size, Image.Resampling.LANCZOS)
+
+        with BytesIO() as buffer:
+            image.save(
+                buffer,
+                format='WEBP',
+                quality=quality,
+                optimize=True,
+                method=6
+            )
+            return buffer.getvalue()
+    except Exception as e:
+        logger.error(f"Image processing failed: {str(e)}")
+        raise
 
 
 @nto_cto_router.post("/necklaceTryOnID")
@@ -394,18 +430,6 @@ async def necklace_try_on_id(necklace_try_on_id: NecklaceTryOnIDEntity = Depends
     logger.info(f">>> NECKLACE TRY ON ID STARTED :: {necklace_try_on_id.storename} <<<")
     start_time = time.time()
 
-    try:
-        data, _ = supabase.table("APIKeyList").select("*").filter("API_KEY", "eq",
-                                                                  necklace_try_on_id.api_token).execute()
-        api_key_actual = data[1][0]['API_KEY']
-        if api_key_actual != necklace_try_on_id.api_token:
-            logger.error(">>> API KEY VALIDATION FAILED <<<")
-            return JSONResponse(content={"error": "Invalid API Key"}, status_code=401)
-        logger.info(">>> API KEY VALIDATION SUCCESSFUL <<<")
-    except Exception as e:
-        logger.error(f">>> API KEY VALIDATION ERROR: {str(e)} <<<")
-        return JSONResponse(content={"error": f"Error validating API key", "code": 500}, status_code=500)
-
     try:
         imageBytes = await image.read()
         jewellery_url = f"https://lvuhhlrkcuexzqtsbqyu.supabase.co/storage/v1/object/public/Stores/{necklace_try_on_id.storename}/{necklace_try_on_id.necklaceCategory}/image/{necklace_try_on_id.necklaceImageId}.png"
@@ -436,39 +460,47 @@ async def necklace_try_on_id(necklace_try_on_id: NecklaceTryOnIDEntity = Depends
         logger.info(">>> SAVING RESULT IMAGES <<<")
         start_time_saving = time.time()
 
-        result_url = supabase_upload_and_return_url(prefix="necklace_try_on", image=result)
-        mask_url = supabase_upload_and_return_url(prefix="necklace_try_on_mask", image=mask)
+        # Upload both images concurrently
+        upload_tasks = [
+            supabase_upload_and_return_url(prefix="necklace_try_on", image=result),
+            supabase_upload_and_return_url(prefix="necklace_try_on_mask", image=mask)
+        ]
+        result_url, mask_url = await asyncio.gather(*upload_tasks)
 
-        logger.info(f">>> RESULT IMAGES SAVED IN {round((time.time() - start_time_saving), 2)}s <<<")
+        if not result_url or not mask_url:
+            raise Exception("Failed to upload one or both images")
 
+        logger.info(f">>> RESULT IMAGES SAVED IN {round((time.time() - start_time_saving), 2)}s <<<")
         logger.info(">>> RESULT IMAGES SAVED <<<")
     except Exception as e:
         logger.error(f">>> RESULT SAVING ERROR: {str(e)} <<<")
         return JSONResponse(content={"error": f"Error saving result images", "code": 500}, status_code=500)
-
     try:
-        creditResponse = deductAndTrackCredit(storename=necklace_try_on_id.storename, endpoint="/necklaceTryOnID")
-        total_backend_time = round((time.time() - start_time), 2)
-        response = {
-            "code": 200,
-            "output": f"{result_url}",
-            "mask": f"{mask_url}",
-            "inference_time": total_backend_time
-        }
-        if creditResponse == "No Credits Available":
-            logger.error(">>> NO CREDITS REMAINING <<<")
-            response = {"error": "No Credits Remaining"}
-            return JSONResponse(content=response, status_code=402)
-        logger.info(">>> CREDITS DEDUCTED SUCCESSFULLY <<<")
-    except Exception as e:
-        logger.error(f">>> CREDIT DEDUCTION ERROR: {str(e)} <<<")
-        return JSONResponse(content={"error": f"Error deducting credits", "code": 500}, status_code=500)
-
-    logger.info(f">>> TOTAL INFERENCE TIME: {total_backend_time}s <<<")
-    logger.info(f">>> NECKLACE TRY ON COMPLETED :: {necklace_try_on_id.storename} <<<")
-    logger.info("-" * 50)
-
-    return JSONResponse(content=response, status_code=200)
+        try:
+            total_backend_time = round((time.time() - start_time), 2)
+            response = {
+                "code": 200,
+                "output": f"{result_url}",
+                "mask": f"{mask_url}",
+                "inference_time": total_backend_time
+            }
+
+    
+
+        except Exception as e:
+            logger.error(f">>> RESPONSE GENERATION ERROR: {str(e)} <<<")
+            return JSONResponse(content={"error": f"Error generating response", "code": 500}, status_code=500)
+
+        logger.info(f">>> TOTAL INFERENCE TIME: {total_backend_time}s <<<")
+        logger.info(f">>> NECKLACE TRY ON COMPLETED :: {necklace_try_on_id.storename} <<<")
+        logger.info("-" * 50)
+
+        return JSONResponse(content=response, status_code=200)
+    
+    finally:
+        if 'result' in locals(): del result
+        if 'mask' in locals(): del mask
+        gc.collect()
 
 
 @nto_cto_router.post("/canvasPoints")