Update app.py

app.py

@@ -117,8 +117,8 @@ def real_fake_check(list_dir, path, model):
     prob = {}
     for i in list_dir:      # real / fake 선택
         input_data = extract_mfcc_path(os.path.join(path, i))
-        #input_data = torch.tensor(input_data).unsqueeze(0).to('cuda')  # 배치 차원을 추가하여 (1, input_dim, sequence_length)로 맞춤
-        input_data = torch.tensor(input_data).unsqueeze(0).to('cpu') 
+        input_data = torch.tensor(input_data).unsqueeze(0).to('cuda')  # 배치 차원을 추가하여 (1, input_dim, sequence_length)로 맞춤
+        #input_data = torch.tensor(input_data).unsqueeze(0).to('cpu') 
         result = model(input_data.float())
         probabilities = F.softmax(result, dim=1)
         prob[i]='%.2f'%probabilities[0][1].item()
@@ -140,8 +140,8 @@ def main(file_name):
         
     hf_token = os.getenv("HUGGINGFACE_TOKEN")
     pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1",use_auth_token=hf_token)
-    #device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
-    device = torch.device('cpu')
+    device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+    #device = torch.device('cpu')
 
     video_file = file_name #deepfake #meganfox.mp4'
     #current_path = os.getcwd()
@@ -161,7 +161,7 @@ def main(file_name):
     # 모델
     model_name = hf_hub_download(repo_id="sssssungk/deepfake_voice", filename="deepvoice_model_girl.pth")
     model = DeepVoiceModel(input_dim, hidden_dim, num_classes, dropout_rate, l2_reg).to(device)
-    model.load_state_dict(torch.load(model_name, map_location=torch.device('cpu')))
+    model.load_state_dict(torch.load(model_name))#, map_location=torch.device('cpu')))
     model.eval()  # 평가 모드로 설정