import gradio as gr
import torchaudio
import soundfile as sf
from demucs.pretrained import get_model
from demucs.apply import apply_model
import numpy as np
import subprocess

def process(voice_file_path, song_file_path):
    # Загрузка файла песни
    mixture, sr = torchaudio.load(song_file_path)
    num_channels = mixture.size(0)
    
    # Проверка частоты дискретизации
    if sr != 44100:
        raise ValueError("Частота дискретизации должна быть 44100 Гц")
    
    # Загрузка модели Demucs
    model = get_model('htdemucs')
    
    # Применение модели для разделения
    tracks = apply_model(model, mixture, device='cpu')
    
    # Извлечение вокала и музыкального трека
    vocal_track = tracks[2].numpy()  # вокал обычно третий по порядку в htdemucs
    music_track = (tracks[0] + tracks[1] + tracks[3]).numpy()  # drums + bass + other
    
    # Сохранение музыкального трека
    if num_channels == 1:
        music_track = music_track[0]  # Преобразование в моно
    sf.write('music_track.wav', music_track.T, sr)
    
    # Подготовка вокала для SEED-VC: преобразование в моно, если стерео
    if vocal_track.shape[0] == 2:
        vocal_mono = np.mean(vocal_track, axis=0)
        sf.write('vocal_track_mono.wav', vocal_mono, sr)
        source_path = 'vocal_track_mono.wav'
    else:
        vocal_mono = vocal_track[0]  # Уже моно
        sf.write('vocal_track_mono.wav', vocal_mono, sr)
        source_path = 'vocal_track_mono.wav'
    
    # Запуск SEED-VC для конверсии голоса
    subprocess.run([
        'python', 'seed-vc/inference.py', 
        '--source', source_path, 
        '--target', voice_file_path, 
        '--output', 'converted_vocal', 
        '--checkpoint', 'seed-uvit-whisper-base'
    ])
    
    # Загрузка преобразованного вокала (моно)
    converted_vocal, sr_vocal = sf.read('converted_vocal/output.wav')
    
    # Проверка соответствия частот дискретизации
    if sr_vocal != sr:
        raise ValueError("Частоты дискретизации не совпадают")
    
    # Если музыкальный трек стерео, преобразовать вокал в стерео
    if num_channels == 2:
        converted_vocal_stereo = np.array([converted_vocal, converted_vocal])
    else:
        converted_vocal_stereo = converted_vocal
    
    # Загрузка музыкального трека
    music_track_loaded, _ = sf.read('music_track.wav')
    
    # Объединение вокала и музыки
    final_song = converted_vocal_stereo.T + music_track_loaded
    
    # Сохранение финальной песни
    sf.write('final_song.wav', final_song, sr)
    
    return 'final_song.wav'

# Создание интерфейса Gradio
with gr.Blocks() as demo:
    voice_input = gr.File(label="Загрузите запись вашего голоса")
    song_input = gr.File(label="Загрузите песню для модификации")
    output = gr.File(label="Финальная песня")
    process_button = gr.Button("Обработать")
    process_button.click(fn=process, inputs=[voice_input, song_input], outputs=output)

if __name__ == "__main__":
    demo.launch()