Make device dependable of the machine capacity

app.py

@@ -3,6 +3,7 @@ import random
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer, util
 import logging
+import torch
 from PIL import Image
 # Create a custom logger
 logger = logging.getLogger(__name__)
@@ -22,9 +23,10 @@ c_handler.setFormatter(c_format)
 logger.addHandler(c_handler)
 
 class SearchEngine:
-    def __init__(self):
+    def __init__(self, device="cpu"):
+        self.device = device if torch.cuda.is_available() else "cpu"
         self.model = SentenceTransformer('clip-ViT-B-32')
-        self.embedding_dataset = load_dataset("JLD/unsplash25k-image-embeddings", trust_remote_code=True, split="train").with_format("torch", device="cuda:0")
+        self.embedding_dataset = load_dataset("JLD/unsplash25k-image-embeddings", trust_remote_code=True, split="train").with_format("torch", device=self.device)
         image_dataset = load_dataset("jamescalam/unsplash-25k-photos", trust_remote_code=True, revision="refs/pr/3")
         self.image_dataset = {image["photo_id"]: image["photo_image_url"] for image in image_dataset["train"]}
 
@@ -35,12 +37,12 @@ class SearchEngine:
 
     def search_images_from_text(self, text):
         logger.info("Searching images from text")
-        emb = self.model.encode(text, convert_to_tensor=True, device="cuda:0")
+        emb = self.model.encode(text, convert_to_tensor=True, device=self.device)
         return self.get_candidates(query_embedding=emb)
 
     def search_images_from_image(self, image):
         logger.info("Searching images from image")
-        emb = self.model.encode(Image.fromarray(image), convert_to_tensor=True, device="cuda:0")
+        emb = self.model.encode(Image.fromarray(image), convert_to_tensor=True, device=self.device)
         return self.get_candidates(query_embedding=emb)
 
 def main():