cuda

app.py

@@ -154,7 +154,7 @@ class FoleyController:
         frames, duration  = read_frames_with_moviepy(input_video, max_frame_nums=max_frame_nums)
         if duration >= 10:
             duration = 10
-        time_frames = torch.FloatTensor(frames).permute(0, 3, 1, 2).to(self.device)
+        time_frames = torch.FloatTensor(frames).permute(0, 3, 1, 2).to('cuda')
         time_frames = video_transform(time_frames)
         time_frames = {'frames': time_frames.unsqueeze(0).permute(0, 2, 1, 3, 4)}
         preds       = self.time_detector(time_frames)
@@ -166,7 +166,7 @@ class FoleyController:
         # w -> b c h w
         time_condition = torch.FloatTensor(time_condition).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(1, 1, 256, 1)
         
-        images = self.image_processor(images=frames, return_tensors="pt").to(self.device)
+        images = self.image_processor(images=frames, return_tensors="pt").to('cuda')
         image_embeddings = self.image_encoder(**images).image_embeds
         image_embeddings = torch.mean(image_embeddings, dim=0, keepdim=True).unsqueeze(0).unsqueeze(0)
         neg_image_embeddings = torch.zeros_like(image_embeddings)