add back CUDA

app.py

@@ -14,15 +14,15 @@ example_indexes = {
     "Uber 10k 2021": "./uber_index",
 }
 
-# device = "cpu"
-# if torch.cuda.is_available():
-#     device = "cuda"
-# elif torch.backends.mps.is_available():
-#     device = "mps"
+device = "cpu"
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
 
 image_embed_model = HuggingFaceEmbedding(
     model_name="llamaindex/vdr-2b-multi-v1",
-    device="cpu",
+    device=device,
     trust_remote_code=True,
     token=os.getenv("HUGGINGFACE_TOKEN"),
     model_kwargs={"torch_dtype": torch.float16},
@@ -31,10 +31,10 @@ image_embed_model = HuggingFaceEmbedding(
 
 text_embed_model = HuggingFaceEmbedding(
     model_name="BAAI/bge-small-en",
-    device="cpu",
+    device=device,
     trust_remote_code=True,
     token=os.getenv("HUGGINGFACE_TOKEN"),
-    embed_batch_size=2,
+    embed_batch_size=1,
 )
 
 def load_index(index_path: str) -> MultiModalVectorStoreIndex: