Update main.py

main.py

@@ -4,6 +4,7 @@ import base64
 import io
 import os
 import logging
+import gc  # Import garbage collector
 from PIL import Image, UnidentifiedImageError
 import torch
 import asyncio
@@ -35,7 +36,7 @@ try:
         "weights/icon_caption_florence",
         torch_dtype=torch.float16,
         trust_remote_code=True,
-    ).to("cuda" if torch.cuda.is_available() else "cpu")
+    ).to(device)
 except Exception as e:
     logger.warning(f"Failed to load caption model on GPU: {e}. Falling back to CPU.")
     model = AutoModelForCausalLM.from_pretrained(
@@ -50,7 +51,7 @@ logger.info("Finished loading models!")
 # Initialize FastAPI app
 app = FastAPI()
 
-MAX_QUEUE_SIZE = 20  # Set a reasonable limit based on your system capacity
+MAX_QUEUE_SIZE = 10  # Set a reasonable limit based on your system capacity
 request_queue = asyncio.Queue(maxsize=MAX_QUEUE_SIZE)
 
 # Define response model
@@ -79,7 +80,7 @@ async def startup_event():
     asyncio.create_task(worker())
 
 
-# Image processing function
+# Image processing function with memory cleanup
 async def process(image_input: Image.Image, box_threshold: float, iou_threshold: float) -> ProcessResponse:
     try:
         # Define save path
@@ -132,11 +133,18 @@ async def process(image_input: Image.Image, box_threshold: float, iou_threshold:
         # Join parsed content list
         parsed_content_list_str = "\n".join([str(item) for item in parsed_content_list])
 
-        return ProcessResponse(
+        response = ProcessResponse(
             image=img_str,
             parsed_content_list=parsed_content_list_str,
             label_coordinates=str(label_coordinates),
         )
+
+        # **Memory Cleanup**
+        del image_input, text, ocr_bbox, dino_labled_img, label_coordinates, parsed_content_list
+        torch.cuda.empty_cache()  # Free GPU memory
+        gc.collect()  # Free CPU memory
+
+        return response
     except Exception as e:
         logger.error(f"Error in process function: {e}")
         raise HTTPException(status_code=500, detail=f"Failed to process the image: {e}")