Update app.py

app.py

@@ -6,8 +6,8 @@ from datasets import Dataset
 from torch.nn import CosineSimilarity
 
 image_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
-image_encoder = ViTModel.from_pretrained("model/image_encoder/epoch_29").eval().to("cuda")
-scribble_encoder = ViTModel.from_pretrained("model/scibble_encoder/epoch_29").eval().to("cuda")
+image_encoder = ViTModel.from_pretrained("model/image_encoder/epoch_29").eval()
+scribble_encoder = ViTModel.from_pretrained("model/scibble_encoder/epoch_29").eval()
 
 candidates: Dataset = None
 
@@ -18,7 +18,7 @@ cosinesimilarity = CosineSimilarity()
 def load_candidates(candidate_dir):
     def preprocess(examples):
         images = [image.convert("RGB") for image in examples["image"]]
-        examples["image_embedding"] = image_encoder(image_processor(images, return_tensors="pt")["pixel_values"].to("cuda"))["pooler_output"]
+        examples["image_embedding"] = image_encoder(image_processor(images, return_tensors="pt")["pixel_values"])["pooler_output"]
         return examples
 
     dataset = [dict(image=Image.open(tempfile.name).convert("RGB").resize((224, 224))) for tempfile in candidate_dir]
@@ -38,7 +38,7 @@ def scribble_matching(input_img: Image):
     input_img = ImageOps.invert(input_img)
 
     scribble = input_img
-    scribble_embedding = scribble_encoder(image_processor(scribble, return_tensors="pt")["pixel_values"].to("cuda"))["pooler_output"].to("cpu")
+    scribble_embedding = scribble_encoder(image_processor(scribble, return_tensors="pt")["pixel_values"])["pooler_output"].to("cpu")
     image_embeddings = torch.tensor(candidates["image_embedding"], dtype=torch.float32)