Spaces:

Zitang
/

Self-attention-based-V1MT-motion-model

Running on Zero

sunana commited on Nov 9, 2023

Commit

b61d593

1 Parent(s): 7430431

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -7,7 +7,7 @@ import sys
 import FFV1MT_MS
 import flow_tools
 import transformers
 def process_images(videos, x, y):
     # read video file
     cap = cv2.VideoCapture(videos)
@@ -87,7 +87,7 @@ if __name__ =='__main__':
     # # Tesla T4
     model = FFV1MT_MS.FFV1DNN()
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print('Number fo parameters: {}'.format(model.num_parameters()))
     model.to(device)

 import FFV1MT_MS
 import flow_tools
 import transformers
+@spaces.GPU
 def process_images(videos, x, y):
     # read video file
     cap = cv2.VideoCapture(videos)
     # # Tesla T4
     model = FFV1MT_MS.FFV1DNN()
+    device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
     print('Number fo parameters: {}'.format(model.num_parameters()))
     model.to(device)