import os
import sys

import ffmpeg
import torch
import warnings

#验证SDK token
from modelscope.hub.api import HubApi
api = HubApi()
api.login(os.getenv('token'))

#模型下载
from modelscope import snapshot_download
model_dir = snapshot_download('yuzhouaoyou/uvr')

warnings.filterwarnings("ignore")
os.system('pip install librosa==0.9.1')
os.system('pip install numpy==1.23.5')
#os.system('pip install flask')

sys.path.append(os.getcwd())
from infer.modules.uvr5.mdxnet import MDXNetDereverb
from infer.modules.uvr5.vr import AudioPre

def run_uvr(file):
    out_path = 'opt'
    func = AudioPre(agg=10,model_path=f'{model_dir}/HP5.pth',device='cuda:0' if torch.cuda.is_available() else 'cpu',is_half=False)
    need_reformat = True
    try:
        info = ffmpeg.probe(file, cmd="ffprobe")
        if info["streams"][0]["channels"] == 2 and info["streams"][0]["sample_rate"] == "44100":
            need_reformat = False
    except:
        need_reformat = True

    if need_reformat:
        tmp_path = "%s/%s.reformatted.wav" % (
            os.path.join(os.environ["TEMP"]),
            os.path.basename(file),
        )
        os.system(
            "ffmpeg -i %s -vn -acodec pcm_s16le -ac 2 -ar 44100 %s -y"
            % (file, tmp_path)
        )
        file = tmp_path
    func._path_audio_(file, out_path, out_path, 'mp3')
    return '%s/instrument_%s_10.mp3' % (out_path,os.path.basename(file)), '%s/vocal_%s_10.mp3' % (out_path,os.path.basename(file))
    
import gradio as gr

def f(file):
    i,v = run_uvr(file)
    return [i,v]

demo = gr.Interface(
    fn=f,
    inputs=[gr.File()],
    outputs=[gr.Files()],
    title="UVR 人声分离",
)

demo.launch()