Update app.py

app.py

@@ -2,15 +2,14 @@ import torch
 import spaces
 import gradio as gr
 import os
+from pyannote.audio import Pipeline
 from pydub import AudioSegment
-from pyannote.audio.pipelines import SpeakerDiarization
-
 
 # 初始化 pyannote/speaker-diarization 模型
 HF_TOKEN = os.environ.get("HUGGINGFACE_READ_TOKEN")
 pipeline = None
 try:
-    pipeline = SpeakerDiarization.from_pretrained(
+    pipeline = Pipeline.from_pretrained(
         "pyannote/speaker-diarization-3.1", use_auth_token=HF_TOKEN
     )
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -19,16 +18,16 @@ except Exception as e:
     print(f"Error initializing pipeline: {e}")
     pipeline = None
 
-# 音频处理函数：拼接目标音频和混合音频
+# 音频拼接函数：拼接目标音频和混合音频
 def combine_audio_with_time(target_audio, mixed_audio):
     if pipeline is None:
         return "错误: 模型未初始化"
 
     # 加载目标说话人的样本音频
-    target_audio_segment = AudioSegment.from_wav(target_audio.name)
+    target_audio_segment = AudioSegment.from_wav(target_audio)
 
     # 加载混合音频
-    mixed_audio_segment = AudioSegment.from_wav(mixed_audio.name)
+    mixed_audio_segment = AudioSegment.from_wav(mixed_audio)
 
     # 记录目标说话人音频的时间点（精确到0.01秒）
     target_start_time = len(mixed_audio_segment) / 1000  # 秒为单位，精确到 0.01 秒
@@ -42,7 +41,7 @@ def combine_audio_with_time(target_audio, mixed_audio):
     return "final_output.wav", target_start_time
 
 # 使用 pyannote/speaker-diarization 对拼接后的音频进行说话人分离
-@spaces.GPU(duration=60 * 2)
+@spaces.GPU(duration=60 * 2)  # 使用 GPU 加速，限制执行时间为 120 秒
 def diarize_audio(temp_file):
     if pipeline is None:
         return "错误: 模型未初始化"
@@ -55,7 +54,7 @@ def diarize_audio(temp_file):
     # 返回 diarization 输出
     return str(diarization)
 
-# 处理并生成标签文件
+# 生成标签文件的函数
 def generate_labels_from_diarization(diarization_output):
     labels_path = 'labels.txt'
     successful_lines = 0
@@ -89,47 +88,46 @@ def timestamp_to_seconds(timestamp):
     except ValueError as e:
         print(f"转换时间戳时出错: '{timestamp}'. 错误: {e}")
         return None
-        
-@spaces.GPU(duration=60 * 2)
-# 处理音频文件
-def process_audio(audio):
-    diarization_result = diarize_audio(save_audio(audio))
+
+# 处理音频文件并返回输出
+def process_audio(target_audio, mixed_audio):
+    # 进行音频拼接
+    final_audio_path, target_start_time = combine_audio_with_time(target_audio, mixed_audio)
+    
+    # 执行说话人分离
+    diarization_result = diarize_audio(final_audio_path)
+    
     if diarization_result.startswith("错误"):
-        return diarization_result, None  # 如果出错，返回错误信息和空的标签文件
+        return diarization_result, None, None  # 出错时返回错误信息
     else:
+        # 生成标签文件
         label_file = generate_labels_from_diarization(diarization_result)
-        return diarization_result, label_file
-
-# 保存上传的音频
-def save_audio(audio):
-    with open(audio.name, "rb") as f:
-        audio_data = f.read()
-
-    # 保存上传的音频文件到临时位置
-    with open("temp.wav", "wb") as f:
-        f.write(audio_data)
-
-    return "temp.wav"
+        return diarization_result, label_file, final_audio_path  # 返回说话人分离结果、标签文件和剪辑后的音频路径
 
 # Gradio 接口
 with gr.Blocks() as demo:
     gr.Markdown("""
     # 🗣️ 音频拼接与说话人分类 🗣️
-    上传目标说话人音频和混合音频，拼接并进行说话人分类。
+    上传目标说话人音频和混合音频，拼接并进行说话人分类。结果包括说话人分离输出、标签文件和剪辑后的音频文件。
     """)
 
-    audio_input = gr.Audio(type="filepath", label="上传目标说话人音频")
+    target_audio_input = gr.Audio(type="filepath", label="上传目标说话人音频")
     mixed_audio_input = gr.Audio(type="filepath", label="上传混合音频")
-
+    
     process_button = gr.Button("处理音频")
+    
+    # 输出结果
     diarization_output = gr.Textbox(label="说话人分离结果")
     label_file_link = gr.File(label="下载标签文件")
+    
+    # 修改为 gr.Audio 组件来返回音频
+    final_audio_link = gr.Audio(label="下载剪辑后的音频", type="file")
 
-    # 处理音频
+    # 点击按钮时触发处理音频
     process_button.click(
         fn=process_audio,
-        inputs=[audio_input],
-        outputs=[diarization_output, label_file_link]
+        inputs=[target_audio_input, mixed_audio_input],
+        outputs=[diarization_output, label_file_link, final_audio_link]
     )
 
 demo.launch(share=False)