fix: memory leak #11

README.md

@@ -20,6 +20,7 @@ services:
       - models:/root/.cache/huggingface/hub:rw
     environment:
       - DEFAULT_MODEL_NAME
+      - BATCH_SIZE
       - ACCESS_TOKEN
       - DEFAULT_SCORE
       - USE_API_KEYS
@@ -41,6 +42,7 @@ services:
       - models:/root/.cache/huggingface/hub:rw
     environment:
       - DEFAULT_MODEL_NAME
+      - BATCH_SIZE
       - ACCESS_TOKEN
       - DEFAULT_SCORE
       - USE_API_KEYS
@@ -61,6 +63,7 @@ volumes:
 - Create a `.env` file and set the preferred values.
 ```sh
 DEFAULT_MODEL_NAME=Falconsai/nsfw_image_detection
+BATCH_SIZE=5
 DEFAULT_SCORE=0.7
 ACCESS_TOKEN=
 

docker-compose.yml

@@ -7,11 +7,34 @@ services:
       - models:/root/.cache/huggingface/hub:rw
     environment:
       - DEFAULT_MODEL_NAME
+      - BATCH_SIZE
       - ACCESS_TOKEN
       - DEFAULT_SCORE
       - USE_API_KEYS
       - API_KEYS
     restart: unless-stopped
 
+  image_video_classification_cuda:
+    image: ghcr.io/doppeltilde/image_video_classification:latest-cuda
+    ports:
+      - "8000:8000"
+    volumes:
+      - models:/root/.cache/huggingface/hub:rw
+    environment:
+      - DEFAULT_MODEL_NAME
+      - BATCH_SIZE
+      - ACCESS_TOKEN
+      - DEFAULT_SCORE
+      - USE_API_KEYS
+      - API_KEYS
+    restart: unless-stopped
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [ gpu ]
+
 volumes:
   models:

src/routes/api/image_classification.py

@@ -7,6 +7,7 @@
 import filetype
 from src.shared.shared import check_model
 from src.middleware.auth.auth import get_api_key
+import torch
 
 router = APIRouter()
 
@@ -84,6 +85,10 @@ async def image_classification(
         print("File is not a valid image.")
         return {"error": str(e)}
 
+    finally:
+        del classifier
+        torch.cuda.empty_cache()
+
 
 @router.post("/api/multi-image-classification", dependencies=[Depends(get_api_key)])
 async def multi_image_classification(
@@ -139,6 +144,7 @@ async def multi_image_classification(
                         )
                     finally:
                         img.close()
+
             else:
                 img.close()
                 return HTTPException(
@@ -149,4 +155,8 @@ async def multi_image_classification(
             img.close()
             return {"error": str(e)}
 
+        finally:
+            del classifier
+            torch.cuda.empty_cache()
+
     return image_list

src/routes/api/image_query_classification.py

@@ -7,6 +7,7 @@
 from concurrent.futures import ThreadPoolExecutor
 from src.shared.shared import check_model, default_score
 from src.middleware.auth.auth import get_api_key
+import torch
 
 router = APIRouter()
 
@@ -141,10 +142,15 @@ async def image_query_classification(
                         raise HTTPException(
                             status_code=400, detail=f"Error classifying image: {e}"
                         )
+
             except Exception as e:
                 print("File is not a valid image.")
                 return {"error": str(e)}
 
+            finally:
+                del classifier
+                torch.cuda.empty_cache()
+
         return totalResults
 
     except Exception as e:
@@ -216,6 +222,8 @@ async def multi_image_query_classification(
                         )
                     finally:
                         img.close()
+                        del classifier
+                        torch.cuda.empty_cache()
 
                 # Check Static Image
                 else:
@@ -247,11 +255,16 @@ async def multi_image_query_classification(
                         )
                     finally:
                         img.close()
+
             except Exception as e:
                 print("File is not a valid image.")
                 img.close()
                 return {"error": str(e)}
 
+            finally:
+                del classifier
+                torch.cuda.empty_cache()
+
         totalResults.append({index: image_list})
 
     return totalResults

src/routes/api/video_classification.py

@@ -9,6 +9,7 @@
 import filetype
 from src.shared.shared import check_model, default_score
 from src.middleware.auth.auth import get_api_key
+import torch
 
 router = APIRouter()
 
@@ -128,3 +129,5 @@ async def video_classification(
     finally:
         tf.close()
         os.remove(tf.name)
+        del classifier
+        torch.cuda.empty_cache()

src/shared/shared.py

@@ -7,6 +7,7 @@
 
 access_token = os.getenv("ACCESS_TOKEN", None)
 default_model_name = os.getenv("DEFAULT_MODEL_NAME", "Falconsai/nsfw_image_detection")
+default_batch_size = os.getenv("BATCH_SIZE", 5)
 default_score = os.getenv("DEFAULT_SCORE", 0.7)
 device = 0 if torch.cuda.is_available() else -1
 
@@ -25,6 +26,7 @@ def check_model(model_name):
             model=_model_name,
             token=access_token,
             device=device,
+            batch_size=default_batch_size,
         )
 
         return classifier