Add support for multiple GPU optimization,

Add support for enabling and disabling directories

Author: merfanian

backend/core/embedders.py

@@ -1,8 +1,8 @@
 import torch
 from timm import create_model, data
-
 from core.singleton import Singleton
 from settings import settings
+import torch.nn as nn
 
 
 class ImageEmbedder:
@@ -12,39 +12,45 @@ def __init__(self, name, model_name, weight, device=torch.device("cpu")):
         self._device = device
         self._weight = weight
 
-        self.model = create_model(model_name, pretrained=True, num_classes=0).to(device)
+        # Create and move the model to the device.
+        model = create_model(model_name, pretrained=True, num_classes=0).to(device)
+
+        # Wrap the model with DataParallel if more than one GPU is available.
+        if torch.cuda.is_available() and torch.cuda.device_count() > 1 and settings.service.use_cuda:
+            self.model = nn.DataParallel(model)
+        else:
+            self.model = model
+
         self.model.eval()
 
+        # Use the unwrapped model for configuration
         self.preprocess = self.get_preprocess()
-        self._weight = 1.0
         self._embedding_dim = self._determine_embedding_dim()
 
     def get_preprocess(self):
-        data_config = data.resolve_model_data_config(self.model)
+        # Unwrap the model if wrapped in DataParallel
+        model_for_config = self.model.module if hasattr(self.model, 'module') else self.model
+        data_config = data.resolve_model_data_config(model_for_config)
         return data.create_transform(**data_config, is_training=False)
 
-    @property
-    def name(self) -> str:
-        return self._name
-
-    @property
-    def model_name(self) -> str:
-        return self._model_name
-
-    @property
-    def device(self) -> torch.device:
-        return self._device
-
     def embed(self, img_binary):
-        img_binary = self.preprocess(img_binary)
-        img_binary = img_binary.unsqueeze(0).to(self.device)
+        # Preprocess the image and add batch dimension.
+        img_tensor = self.preprocess(img_binary)
+        img_tensor = img_tensor.unsqueeze(0).to(self.device)
         with torch.no_grad():
-            embedding = self.model(img_binary).squeeze(0).cpu().numpy()
+            # DataParallel will split the batch across GPUs.
+            embedding = self.model(img_tensor).squeeze(0).cpu().numpy()
         return embedding
 
     def _determine_embedding_dim(self):
-        # Generate a dummy image tensor to determine the output dimension of the embedder
-        dummy_input = torch.zeros((3, 224, 224)).to(self.device)  # Assuming the input size is 3x224x224
+        # Unwrap the model to get the proper configuration.
+        model_for_config = self.model.module if hasattr(self.model, 'module') else self.model
+        data_config = data.resolve_model_data_config(model_for_config)
+        # Get the expected input size from the configuration; defaults to (3,224,224)
+        input_size = data_config.get("input_size", (3, 224, 224))
+
+        # Create a dummy input tensor with the correct size.
+        dummy_input = torch.zeros(input_size).to(self.device)
         dummy_input = self.preprocess(dummy_input).unsqueeze(0).to(self.device)
         with torch.no_grad():
             embedding = self.model(dummy_input).squeeze(0).cpu().numpy()
@@ -58,6 +64,10 @@ def embedding_dim(self):
     def weight(self):
         return self._weight
 
+    @property
+    def device(self):
+        return self._device
+
 
 @Singleton
 class EmbedderManager:

backend/core/singleton.py

@@ -1,34 +1,12 @@
 class Singleton:
-    """
-    A non-thread-safe helper class to ease implementing singletons.
-    This should be used as a decorator -- not a metaclass -- to the
-    class that should be a singleton.
-
-    The decorated class can define one `__init__` function that
-    takes only the `self` argument. Also, the decorated class cannot be
-    inherited from. Other than that, there are no restrictions that apply
-    to the decorated class.
-
-    To get the singleton instance, use the `instance` method. Trying
-    to use `__call__` will result in a `TypeError` being raised.
-
-    """
-
     def __init__(self, decorated):
         self._decorated = decorated
 
     def instance(self):
-        """
-        Returns the singleton instance. Upon its first call, it creates a
-        new instance of the decorated class and calls its `__init__` method.
-        On all subsequent calls, the already created instance is returned.
-
-        """
-        try:
-            return self._instance
-        except AttributeError:
-            self._instance = self._decorated()
-            return self._instance
+        # Store the instance on the decorated class itself.
+        if not hasattr(self._decorated, '_instance'):
+            self._decorated._instance = self._decorated()
+        return self._decorated._instance
 
     def __call__(self):
         raise TypeError('Singletons must be accessed through `instance()`.')

backend/database/__init__.py

@@ -0,0 +1,5 @@
+from .services.image_indexing_service import ImageIndexingService
+
+image_indexing_service = ImageIndexingService.instance()
+
+__all__ = ["image_indexing_service"]

backend/indexing/__init__.py

@@ -0,0 +1,82 @@
+import os
+import threading
+import time
+
+from sqlalchemy.orm import Session
+
+from core import embedder_manager
+from models.models import SessionLocal, Directory, Image
+from indexing.repositories.repositories import DirectoryRepository, ImageRepository, MilvusRepository
+from monitoring import logger
+from settings import settings
+
+
+class ConsistencyChecker:
+    def __init__(self, interval: int = 3600):
+        self.interval = interval
+        self.thread = threading.Thread(target=self.run, daemon=True)
+
+    def start(self):
+        self.thread.start()
+
+    def run(self):
+        while True:
+            time.sleep(self.interval)
+            self.check_consistency()
+
+    def check_consistency(self):
+        logger.info("Running system-wide consistency check")
+        session = SessionLocal()
+        try:
+            directory_repo = DirectoryRepository(session)
+            directories = directory_repo.get_all()
+            for directory in directories:
+                self.check_directory(session, directory)
+            logger.info("Consistency check completed")
+        except Exception as e:
+            logger.error(f"Consistency check error: {e}", exc_info=True)
+            session.rollback()
+        finally:
+            session.close()
+
+    def check_directory(self, session: Session, directory: Directory):
+        logger.info(f"Checking consistency for directory {directory.path} (ID: {directory.id})")
+        if not os.path.exists(directory.path):
+            logger.warning(f"Directory missing: {directory.path}. Removing from system.")
+            DirectoryRepository(session).delete(directory)
+            return
+
+        # Gather filesystem image paths
+        fs_paths = set()
+        for entry in os.scandir(directory.path):
+            if entry.is_file() and entry.name.lower().endswith(('.png', '.jpg', '.jpeg')):
+                fs_paths.add(entry.path)
+            elif entry.is_dir() and settings.directory.recursive_indexing:
+                for root, _, files in os.walk(entry.path):
+                    for file in files:
+                        if file.lower().endswith(('.png', '.jpg', '.jpeg')):
+                            fs_paths.add(os.path.join(root, file))
+
+        # Get database image paths
+        image_repo = ImageRepository(session)
+        db_images = session.query(Image).filter(Image.directory_id == directory.id).all()
+        db_paths = {img.path for img in db_images}
+
+        new_paths = fs_paths - db_paths
+        deleted_paths = db_paths - fs_paths
+        logger.info(f"Directory {directory.path}: {len(new_paths)} new images, {len(deleted_paths)} missing images")
+
+        # Add new images to DB
+        for path in new_paths:
+            if not image_repo.get_by_path(path):
+                session.add(Image(path=path, directory_id=directory.id, is_indexed=False))
+        session.commit()
+
+        # Remove deleted images from DB and Milvus
+        for path in deleted_paths:
+            image = image_repo.get_by_path(path)
+            if image:
+                for embedder_name in embedder_manager.get_image_embedders().keys():
+                    MilvusRepository().delete_entries(embedder_name, f"image_path == '{path}'")
+                image_repo.delete(image)
+        session.commit()

backend/indexing/consistency/__init__.py

@@ -0,0 +1,44 @@
+import queue
+import threading
+from monitoring import logger
+
+from concurrent.futures import ThreadPoolExecutor
+
+from core.singleton import Singleton
+from models.models import SessionLocal
+from indexing.repositories.repositories import MilvusRepository
+from indexing.services.directory_indexer import DirectoryIndexer
+from indexing.services.embedder_service import EmbedderService
+from settings import settings
+
+
+@Singleton
+class IndexQueueManager:
+    def __init__(self):
+        self.index_queue = queue.PriorityQueue()
+        self.processing_paths = set()
+        self.queue_lock = threading.Lock()
+        self.index_workers = ThreadPoolExecutor(max_workers=settings.directory.num_watcher_workers)
+        self.embedder_service = EmbedderService()
+        self.milvus_repo = MilvusRepository()
+        self.directory_indexer = DirectoryIndexer(self.embedder_service, self.milvus_repo)
+
+    def add_to_queue(self, directory_id: int, path: str, priority: int = 0):
+        with self.queue_lock:
+            if (directory_id, path) not in self.processing_paths:
+                self.index_queue.put((priority, (directory_id, path)))
+                self.processing_paths.add((directory_id, path))
+                logger.debug(f"Queued directory {path} (ID: {directory_id}) with priority {priority}")
+                self.index_workers.submit(self._process_queue)
+
+    def _process_queue(self):
+        while not self.index_queue.empty():
+            priority, (directory_id, path) = self.index_queue.get()
+            session = SessionLocal()
+            try:
+                self.directory_indexer.index_directory(directory_id, path, session)
+            finally:
+                session.close()
+                with self.queue_lock:
+                    self.processing_paths.discard((directory_id, path))
+                logger.debug(f"Finished processing directory {path} (ID: {directory_id})")

backend/indexing/consistency/consistency_checker.py

@@ -0,0 +1,79 @@
+from monitoring import logger
+
+from typing import List, Dict
+
+from pymilvus import Collection
+from sqlalchemy.orm import Session
+
+from models.models import Directory, Image
+
+
+class DirectoryRepository:
+    def __init__(self, session: Session):
+        self.session = session
+
+    def get_by_path(self, path: str) -> Directory:
+        return self.session.query(Directory).filter(Directory.path == path).first()
+
+    def create(self, path: str) -> Directory:
+        directory = Directory(path=path, is_indexed=False)
+        self.session.add(directory)
+        self.session.commit()
+        self.session.refresh(directory)
+        logger.debug(f"Created directory entry with ID {directory.id} for path {path}")
+        return directory
+
+    def get_all(self) -> List[Directory]:
+        return self.session.query(Directory).all()
+
+    def delete(self, directory: Directory):
+        self.session.delete(directory)
+        self.session.commit()
+
+
+class ImageRepository:
+    def __init__(self, session: Session):
+        self.session = session
+
+    def get_by_path(self, path: str) -> Image:
+        return self.session.query(Image).filter(Image.path == path).first()
+
+    def add_new_images(self, directory_id: int, image_paths: List[str]) -> List[Image]:
+        new_images = []
+        for path in image_paths:
+            if not self.get_by_path(path):
+                image = Image(path=path, directory_id=directory_id, is_indexed=False)
+                self.session.add(image)
+                new_images.append(image)
+        self.session.commit()
+        logger.info(f"Added {len(new_images)} new images to database for directory {directory_id}")
+        return new_images
+
+    def get_unindexed_images(self, directory_id: int) -> List[Image]:
+        return self.session.query(Image).filter(
+            Image.directory_id == directory_id,
+            Image.is_indexed == False
+        ).all()
+
+    def delete(self, image: Image):
+        self.session.delete(image)
+        self.session.commit()
+
+
+class MilvusRepository:
+    def delete_entries(self, embedder_name: str, expr: str):
+        collection = Collection(embedder_name)
+        result = collection.delete(expr)
+        collection.flush()
+        logger.info(f"Deleted {result.delete_count} entries in Milvus collection '{embedder_name}' using expr {expr}")
+        return result
+
+    def insert_entries(self, embedder_name: str, entries: List[Dict]):
+        collection = Collection(embedder_name)
+        collection.insert(entries)
+        collection.flush()
+        logger.debug(f"Inserted {len(entries)} entries into Milvus collection '{embedder_name}'")
+
+    def query_entries(self, embedder_name: str, expr: str, output_fields: List[str], batch_size: int = 1000):
+        collection = Collection(embedder_name)
+        return collection.query_iterator(expr=expr, output_fields=output_fields, batch_size=batch_size)