device handling

app.py

@@ -5,11 +5,19 @@ from PIL import Image, ImageFilter
 import uuid
 from scipy.interpolate import interp1d, PchipInterpolator
 import torchvision
+import torch
 from utils import *
 
 output_dir = "outputs"
 ensure_dirname(output_dir)
 
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cpu"
+
 def interpolate_trajectory(points, n_points):
     x = [point[0] for point in points]
     y = [point[1] for point in points]
@@ -219,7 +227,7 @@ with gr.Blocks() as demo:
                     2.4. Click "Delete last step" to delete the lastest clicked control point.<br>
                 3. Animate the image according the path with a click on "Run" button. <br>""")
     
-    DragNUWA_net = Drag("cuda:0", 'drag_nuwa.pth', 'DragNUWA_net.py', 320, 576, 14)
+    DragNUWA_net = Drag(device, 'drag_nuwa.pth', 'DragNUWA_net.py', 320, 576, 14)
     first_frame_path = gr.State()
     tracking_points = gr.State([])
 

requirements.txt

@@ -24,11 +24,11 @@ streamlit>=0.73.1
 tensorboardx==2.6
 timm>=0.9.2
 tokenizers==0.12.1
-torch>=2.0.1
-torchaudio>=2.0.2
+#torch>=2.0.1
+#torchaudio>=2.0.2
 torchdata==0.6.1
 torchmetrics>=1.0.1
-torchvision>=0.15.2
+#torchvision>=0.15.2
 tqdm>=4.65.0
 transformers==4.19.1
 triton==2.0.0
@@ -36,7 +36,7 @@ urllib3<1.27,>=1.25.4
 wandb>=0.15.6
 webdataset>=0.2.33
 wheel>=0.41.0
-xformers>=0.0.20
+#xformers>=0.0.20
 colorlog 
 deepdish  
 json_lines 
@@ -51,4 +51,4 @@ pyyaml
 pandas 
 einops 
 deepspeed
-gradio==3.50.2
+gradio==3.50.2