Update app.py

app.py

@@ -1,38 +1,14 @@
 import os
-import time
 import numpy as np
 import tempfile
 from scipy.io import wavfile
-from pytube import YouTube
 import gradio as gr
-from moviepy.editor import AudioFileClip
 from inference import EnsembleDemucsMDXMusicSeparationModel, predict_with_model
 import torch
 import librosa
 import librosa.display
 import matplotlib.pyplot as plt
 
-# 下载YouTube视频并转换为WAV格式
-def download_youtube_video_as_wav(youtube_url):
-    output_dir = "downloads"
-    os.makedirs(output_dir, exist_ok=True)
-    output_file = os.path.join(output_dir, "temp.mp4")
-
-    try:
-        yt = YouTube(youtube_url)
-        yt.streams.filter(only_audio=True).first().download(filename=output_file)
-        print("下载成功。")
-    except Exception as e:
-        print(f"下载视频时发生错误: {e}")
-        return None
-
-    # 将mp4音频转换为wav
-    wav_file = os.path.join(output_dir, "mixture.wav")
-    clip = AudioFileClip(output_file)
-    clip.write_audiofile(wav_file)
-
-    return wav_file
-
 # 检查文件是否准备好
 def check_file_readiness(filepath):
     num_same_size_checks = 0
@@ -70,18 +46,8 @@ def generate_spectrograms(audio_files):
     return tuple(output_spectrograms)
 
 # 音乐分离的包装函数
-def separate_music_file_wrapper(input_string, use_cpu, use_single_onnx, large_overlap, small_overlap, chunk_size, use_large_gpu):
-    input_files = []
-    # 验证YouTube URL或目录路径
-    if input_string.startswith("https://www.youtube.com") or input_string.startswith("https://youtu.be"):
-        output_file = download_youtube_video_as_wav(input_string)
-        if output_file is not None:
-            input_files.append(output_file)
-    elif os.path.isdir(input_string):
-        input_directory = input_string
-        input_files = [os.path.join(input_directory, f) for f in os.listdir(input_directory) if f.endswith('.wav')]
-    else:
-        raise ValueError("无效输入！请输入有效的YouTube链接或目录路径。")
+def separate_music_file_wrapper(uploaded_files, use_cpu, use_single_onnx, large_overlap, small_overlap, chunk_size, use_large_gpu):
+    input_files = [uploaded_files.name]  # 获取上传文件的路径
 
     # 验证重叠值
     if not (0 <= large_overlap <= 1) or not (0 <= small_overlap <= 1):
@@ -157,7 +123,7 @@ theme = gr.themes.Base(
 
 with gr.Blocks(theme=theme) as demo:
     gr.Markdown(description)
-    input_string = gr.Text(label="YouTube链接/URL")
+    uploaded_file = gr.File(label="上传音频文件", type="file")
     use_cpu = gr.Checkbox(label="仅使用CPU", value=True)
     use_single_onnx = gr.Checkbox(label="使用单一ONNX模型", value=False)
     large_overlap = gr.Number(label="大块重叠", value=0.6)
@@ -181,7 +147,7 @@ with gr.Blocks(theme=theme) as demo:
     
     process_button.click(
         separate_music_file_wrapper,
-        inputs=[input_string, use_cpu, use_single_onnx, large_overlap, small_overlap, chunk_size, use_large_gpu],
+        inputs=[uploaded_file, use_cpu, use_single_onnx, large_overlap, small_overlap, chunk_size, use_large_gpu],
         outputs=[vocals, instrumental, instrumental2, bass, drums, other, vocals_spectrogram, instrumental_spectrogram, instrumental2_spectrogram, bass_spectrogram, drums_spectrogram, other_spectrogram],
     )