Update app.py

app.py

@@ -59,7 +59,7 @@ H = W
 #device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def get_pipe_image_and_video_predictor():
-    device="cpu"
+    device="cuda"
     vae = AutoencoderKLWan.from_pretrained("./model/vae", torch_dtype=torch.float16)
     transformer = Transformer3DModel.from_pretrained("./model/transformer", torch_dtype=torch.float16)
     scheduler = UniPCMultistepScheduler.from_pretrained("./model/scheduler")
@@ -195,6 +195,8 @@ def inference_and_return_video(dilation_iterations, num_inference_steps, video_s
     print(mask_tensor.shape)
     mask_tensor = mask_tensor[:,:,:]
 
+    pipe.to("cuda")
+
     if mask_tensor.shape[1] < mask_tensor.shape[2]:
         height = 480
         width = 832