Update app.py

app.py

@@ -56,7 +56,7 @@ random_seed = 42
 video_length = 201
 W = 1024
 H = W
-device = "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def get_pipe_image_and_video_predictor():
     vae = AutoencoderKLWan.from_pretrained("./model/vae", torch_dtype=torch.float16)
@@ -189,7 +189,8 @@ def inference_and_return_video(dilation_iterations, num_inference_steps, video_s
     images = np.array(images)
     masks = np.array(masks)
     img_tensor, mask_tensor = preprocess_for_removal(images, masks)
-    mask_tensor = mask_tensor[:,:,:,:1]
+    print(mask_tensor.shape)
+    mask_tensor = mask_tensor[:,:,:]
 
     if mask_tensor.shape[1] < mask_tensor.shape[2]:
         height = 480