Spaces:
Runtime error
Runtime error
Fix for cpu
Browse files
app.py
CHANGED
@@ -18,7 +18,7 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
18 |
yolov8 = YOLO(hf_hub_download('stamps-labs/yolov8-finetuned', filename='best.torchscript'), task='detect')
|
19 |
|
20 |
yolo_stamp = YOLOStamp()
|
21 |
-
yolo_stamp.load_state_dict(torch.load(hf_hub_download('stamps-labs/yolo-stamp', filename='state_dict.pth')))
|
22 |
yolo_stamp = yolo_stamp.to(device)
|
23 |
yolo_stamp.eval()
|
24 |
transform = A.Compose([
|
@@ -31,7 +31,7 @@ vits8 = vits8.to(device)
|
|
31 |
vits8.eval()
|
32 |
|
33 |
encoder = Encoder()
|
34 |
-
encoder.load_state_dict(torch.load(hf_hub_download('stamps-labs/vae-encoder', filename='encoder.pth')))
|
35 |
encoder = encoder.to(device)
|
36 |
encoder.eval()
|
37 |
|
|
|
18 |
yolov8 = YOLO(hf_hub_download('stamps-labs/yolov8-finetuned', filename='best.torchscript'), task='detect')
|
19 |
|
20 |
yolo_stamp = YOLOStamp()
|
21 |
+
yolo_stamp.load_state_dict(torch.load(hf_hub_download('stamps-labs/yolo-stamp', filename='state_dict.pth'), map_location='cpu'))
|
22 |
yolo_stamp = yolo_stamp.to(device)
|
23 |
yolo_stamp.eval()
|
24 |
transform = A.Compose([
|
|
|
31 |
vits8.eval()
|
32 |
|
33 |
encoder = Encoder()
|
34 |
+
encoder.load_state_dict(torch.load(hf_hub_download('stamps-labs/vae-encoder', filename='encoder.pth'), map_location='cpu'))
|
35 |
encoder = encoder.to(device)
|
36 |
encoder.eval()
|
37 |
|