Update app.py

app.py

@@ -56,7 +56,7 @@ random_seed = 42
 video_length = 201
 W = 1024
 H = W
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"
 
 def get_pipe_image_and_video_predictor():
     vae = AutoencoderKLWan.from_pretrained("./model/vae", torch_dtype=torch.float16)
@@ -177,7 +177,7 @@ def preprocess_for_removal(images, masks):
         out_masks.append(msk_resized)
     arr_images = np.stack(out_images)
     arr_masks = np.stack(out_masks)
-    return torch.from_numpy(arr_images).half(), torch.from_numpy(arr_masks).half()
+    return torch.from_numpy(arr_images).half().to(device), torch.from_numpy(arr_masks).half().to(device)
 
 @spaces.GPU(duration=300)
 def inference_and_return_video(dilation_iterations, num_inference_steps, video_state=None):
@@ -189,8 +189,7 @@ 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)
-    img_tensor=img_tensor.to(device)
-    mask_tensor = mask_tensor[:,:,:,:1].to(device)
+    mask_tensor = mask_tensor[:,:,:,:1]
 
     if mask_tensor.shape[1] < mask_tensor.shape[2]:
         height = 480
@@ -207,7 +206,7 @@ def inference_and_return_video(dilation_iterations, num_inference_steps, video_s
                 height=height,
                 width=width,
                 num_inference_steps=int(num_inference_steps),
-                generator=torch.Generator(device=device).manual_seed(random_seed),
+                generator=torch.Generator(device="cuda").manual_seed(random_seed),
                 iterations=int(dilation_iterations)
         ).frames[0]
 
@@ -403,4 +402,4 @@ with gr.Blocks() as demo:
         clear_btn.click(clear_clicks, inputs=video_state, outputs=image_output)
         track_btn.click(track_video, inputs=[n_frames_slider, video_state], outputs=video_output)
 
-demo.launch()
+demo.launch()