Update app.py

app.py

@@ -31,6 +31,10 @@ if True:
     file_path = snapshot_download(repo_id="lshzhm/Video-to-Audio-and-Piano", local_dir=model_path)
 
     print(f"Model saved at: {file_path}")
+    
+    device = "cpu"
+else:
+    device = "cuda"
 
 log = logging.getLogger()
 
@@ -94,7 +98,7 @@ def read_audio_from_video(video_path):
     return waveform
 
 
-def load():
+def load(device):
     #duration_predictor = DurationPredictor(
     #    transformer = dict(
     #        dim = 512,
@@ -140,7 +144,7 @@ def load():
         num_channels = 128,
         sampling_rate = 24000,
     )
-    e2tts = e2tts.to("cuda")
+    e2tts = e2tts.to(device)
 
     #checkpoint = torch.load("/ckptstorage/zhanghaomin/e2/e2_tts_experiment_v2a_encodec/3000.pt", map_location="cpu")
     #checkpoint = torch.load("/ckptstorage/zhanghaomin/e2/e2_tts_experiment_v2a_encodec_more/500.pt", map_location="cpu")
@@ -159,7 +163,7 @@ def load():
     for param in e2tts.vocos.parameters():
         param.requires_grad = False
     e2tts.vocos.eval()
-    e2tts.vocos.to("cuda")
+    e2tts.vocos.to(device)
     
     #dataset = HFDataset(load_dataset("parquet", data_files={"test": "/ckptstorage/zhanghaomin/tts/GLOBE/data/test-*.parquet"})["test"])
     #sample = dataset[1]
@@ -190,7 +194,7 @@ def load():
     return e2tts, stft
 
 
-e2tts, stft = load()
+e2tts, stft = load(device)
 
 
 def run(e2tts, stft, arg1, arg2, arg3, arg4):