Update app.py

app.py

@@ -1,42 +1,40 @@
-import torch
-
-import spaces
-import os
 import gradio as gr
-from huggingface_hub import HfApi
-from gradio_client.exceptions import AuthenticationError
-
+import os
+import torch
 from model import Wav2Vec2BERT_Llama  # 自定义模型模块
 import dataset  # 自定义数据集模块
+from huggingface_hub import hf_hub_download
 
-@spaces.GPU
-def dummy(): # just a dummy
-    pass
 
 # 初始化设备
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # 初始化模型
-# def load_model():
-#     model = Wav2Vec2BERT_Llama().to(device)
-#     checkpoint_path = "ckpt/model_checkpoint.pth"
-#     if os.path.exists(checkpoint_path):
-#         checkpoint = torch.load(checkpoint_path)
-#         model_state_dict = checkpoint['model_state_dict']
-
-#         # 处理模型状态字典的 key
-#         if hasattr(model, 'module') and not any(key.startswith('module.') for key in model_state_dict.keys()):
-#             model_state_dict = {'module.' + key: value for key, value in model_state_dict.items()}
-#         elif not hasattr(model, 'module') and any(key.startswith('module.') for key in model_state_dict.keys()):
-#             model_state_dict = {key.replace('module.', ''): value for key, value in model_state_dict.items()}
-
-#         model.load_state_dict(model_state_dict)
-#         model.eval()
-#     else:
-#         raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
-#     return model
-
-# model = load_model()
+def load_model():
+    model = Wav2Vec2BERT_Llama().to(device)
+    checkpoint_path = hf_hub_download(
+        repo_id="amphion/deepfake_detection", 
+        filename="checkpoints_wav2vec2bert_ft_llama_labels_ASVspoof2019_RandomPrompts_6/model_checkpoint.pth"
+    )
+    # checkpoint_path = "ckpt/model_checkpoint.pth"
+    if os.path.exists(checkpoint_path):
+        checkpoint = torch.load(checkpoint_path)
+        model_state_dict = checkpoint['model_state_dict']
+        threshold = 0.9996
+
+        # 处理模型状态字典的 key
+        if hasattr(model, 'module') and not any(key.startswith('module.') for key in model_state_dict.keys()):
+            model_state_dict = {'module.' + key: value for key, value in model_state_dict.items()}
+        elif not hasattr(model, 'module') and any(key.startswith('module.') for key in model_state_dict.keys()):
+            model_state_dict = {key.replace('module.', ''): value for key, value in model_state_dict.items()}
+
+        model.load_state_dict(model_state_dict)
+        model.eval()
+    else:
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+    return model, threshold
+
+model, threshold = load_model()
 
 # 检测函数
 def detect(dataset, model):
@@ -96,13 +94,13 @@ def gradio_ui():
     interface = gr.Interface(
         fn=detection_wrapper,  # 主函数
         inputs=[
-            gr.Audio(type="filepath", label="Demonstration Audio 1"),
+            gr.Audio(source="upload", type="filepath", label="Demonstration Audio 1"),
             gr.Dropdown(choices=["bonafide", "spoof"], value="bonafide", label="Label 1"),
-            gr.Audio(type="filepath", label="Demonstration Audio 2"),
+            gr.Audio(source="upload", type="filepath", label="Demonstration Audio 2"),
             gr.Dropdown(choices=["bonafide", "spoof"], value="bonafide", label="Label 2"),
-            gr.Audio(type="filepath", label="Demonstration Audio 3"),
+            gr.Audio(source="upload", type="filepath", label="Demonstration Audio 3"),
             gr.Dropdown(choices=["bonafide", "spoof"], value="bonafide", label="Label 3"),
-            gr.Audio(type="filepath", label="Query Audio (Audio for Detection)")
+            gr.Audio(source="upload", type="filepath", label="Query Audio (Audio for Detection)")
         ],
         outputs=gr.JSON(label="Detection Results"),
         title="Audio Deepfake Detection System",
@@ -110,7 +108,6 @@ def gradio_ui():
     )
     return interface
 
-
 if __name__ == "__main__":
     demo = gradio_ui()
-    demo.launch(share=False)
+    demo.launch()