Updated so it works on CPU

app.py

@@ -12,7 +12,7 @@ from braceexpand import braceexpand
 
 # Load model
 checkpoint_path = "ViT-B/16"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"
 model, preprocess = clip.load(checkpoint_path, device=device, jit=False)
 
 
@@ -78,18 +78,18 @@ def estimate_price_and_usage(img):
         bias=False
     )
     # Load weights for the linear layer as a tensor
-    linear_data = torch.load("files/reuse_linear.pth")
+    linear_data = torch.load("files/reuse_linear.pth", map_location="cpu")
     probe.weight.data = linear_data["weight"]
 
     # Do inference
-    probe.eval()
-    probe = probe.to(device)
-    output = probe(query_features)
-    print(output)
-    output = torch.softmax(output, dim=-1)
-    output = output.cpu().detach().numpy().astype("float32")
-    reuse = output.argmax(axis=-1)[0]
-    reuse_classes = ["Reuse", "Export"]
+    with torch.autocast("cpu"):
+        probe.eval()
+        probe = probe.to(device)
+        output = probe(query_features)
+        output = torch.softmax(output, dim=-1)
+        #output = output.cpu().detach().numpy().astype("float32")
+        reuse = output.argmax(axis=-1)[0]
+        reuse_classes = ["Reuse", "Export"]
 
     # Estimate price
     num_classes = 4
@@ -101,17 +101,18 @@ def estimate_price_and_usage(img):
     )
     # Print output shape for the linear layer
     # Load weights for the linear layer as a tensor
-    linear_data = torch.load("files/price_linear.pth")
+    linear_data = torch.load("files/price_linear.pth", map_location="cpu")
     probe.weight.data = linear_data["weight"]
 
     # Do inference
-    probe.eval()
-    probe = probe.to(device)
-    output = probe(query_features)
-    output = torch.softmax(output, dim=-1)
-    output = output.cpu().detach().numpy().astype("float32")
-    price = output.argmax(axis=-1)[0]
-    price_classes = ["<50", "50-100", "100-150", ">150"]
+    with torch.autocast("cpu"):
+        probe.eval()
+        probe = probe.to(device)
+        output = probe(query_features)
+        output = torch.softmax(output, dim=-1)
+        #output = output.cpu().detach().numpy().astype("float32")
+        price = output.argmax(axis=-1)[0]
+        price_classes = ["<50", "50-100", "100-150", ">150"]
 
     return f"Estimated price: {price_classes[price]} SEK - Usage: {reuse_classes[reuse]}"