unknown commited on
Commit
346394a
1 Parent(s): 14c3956
Files changed (1) hide show
  1. app.py +5 -5
app.py CHANGED
@@ -64,7 +64,7 @@ class FoleyController:
64
  self.savedir_sample = os.path.join(self.savedir, "sample")
65
  os.makedirs(self.savedir, exist_ok=True)
66
 
67
- self.device = "cuda"
68
 
69
  self.pipeline = None
70
 
@@ -140,10 +140,10 @@ class FoleyController:
140
  seed_textbox,
141
  ):
142
  # move to gpu
143
- self.time_detector.to(self.device)
144
- self.pipeline.to(self.device)
145
- self.vocoder.to(self.device)
146
- self.image_encoder.to(self.device)
147
 
148
  vision_transform_list = [
149
  torchvision.transforms.Resize((128, 128)),
 
64
  self.savedir_sample = os.path.join(self.savedir, "sample")
65
  os.makedirs(self.savedir, exist_ok=True)
66
 
67
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
68
 
69
  self.pipeline = None
70
 
 
140
  seed_textbox,
141
  ):
142
  # move to gpu
143
+ self.time_detector = self.time_detector.to(self.device)
144
+ self.pipeline = self.pipeline.to(self.device)
145
+ self.vocoder = self.vocoder.to(self.device)
146
+ self.image_encoder = self.image_encoder.to(self.device)
147
 
148
  vision_transform_list = [
149
  torchvision.transforms.Resize((128, 128)),