added image ai-moderation

ai-moderation/app/ml/model_manager.py

@@ -3,6 +3,8 @@ from transformers import AutoModelForSequenceClassification
 
 from app.config.settings import settings
 from app.config.logging import logger
+from transformers import AutoImageProcessor
+from transformers import AutoModelForImageClassification
 
 import torch
 
@@ -13,6 +15,8 @@ class ModelManager:
     self.device = None
     self.text_model = None
     self.text_tokenizer = None
+    self.image_model = None
+    self.image_processor = None
 
   def load_device(self):
     if settings.DEVICE == "auto":
@@ -51,6 +55,35 @@ class ModelManager:
     print(f"Using device: {self.device}")
     self.load_tokenizer()
     self.load_model()
+    self.load_image_model()
+
+  def load_image_model(self):
+
+    print("Loading image model...")
+
+
+    self.image_processor = (
+      AutoImageProcessor.from_pretrained(
+          settings.IMAGE_MODEL,
+          cache_dir=settings.MODEL_CACHE_DIR
+      )
+    )
+
+
+    self.image_model = (
+      AutoModelForImageClassification.from_pretrained(
+          settings.IMAGE_MODEL,
+          cache_dir=settings.MODEL_CACHE_DIR
+      )
+    )
+
+
+    self.image_model.to(self.device)
+
+    self.image_model.eval()
+
+
+    print("Image model loaded.")
 
 
 model_manager = ModelManager()