app.py

import os  # 파일 및 디렉토리 작업을 위한 모듈
import torch  # 딥러닝 프레임워크 PyTorch
import librosa  # 오디오 처리를 위한 모듈
import binascii  # 이진 데이터를 다루는 모듈
import warnings  # 경고 메시지를 출력하는 모듈
import midi2audio  # MIDI 파일을 WAV 파일로 변환하는 모듈
import numpy as np  # 다차원 배열을 다루는 모듈
import pytube as pt  # YouTube 동영상을 다운로드하는 모듈
import gradio as gr  # 인터랙티브한 UI를 만들기 위한 모듈
import soundfile as sf  # 사운드 파일을 다루는 모듈
from transformers import Pop2PianoForConditionalGeneration, Pop2PianoProcessor  # Pop2Piano 모델과 전처리기

yt_video_dir = "./yt_dir"  # YouTube 동영상 다운로드 디렉토리 경로
outputs_dir = "./midi_wav_outputs"  # MIDI 및 WAV 파일 출력 디렉토리 경로
os.makedirs(outputs_dir, exist_ok=True)  # 출력 디렉토리 생성 (이미 존재하는 경우 무시)
os.makedirs(yt_video_dir, exist_ok=True)  # YouTube 동영상 다운로드 디렉토리 생성 (이미 존재하는 경우 무시)

device = "cuda" if torch.cuda.is_available() else "cpu"  # CUDA가 사용 가능한 경우 GPU를 사용하고, 그렇지 않은 경우 CPU를 사용
model = Pop2PianoForConditionalGeneration.from_pretrained("sweetcocoa/pop2piano").to(device)  # 사전 학습된 Pop2Piano 모델 로드
processor = Pop2PianoProcessor.from_pretrained("sweetcocoa/pop2piano")  # 사전 학습된 Pop2Piano 전처리기 로드
composers = model.generation_config.composer_to_feature_token.keys()  # 작곡가 목록 가져오기

def get_audio_from_yt_video(yt_link):
        try:
                yt = pt.YouTube(yt_link)  # YouTube 동영상 객체 생성
                t = yt.streams.filter(only_audio=True)  # 오디오 스트림 필터링
                filename = os.path.join(yt_video_dir, binascii.hexlify(os.urandom(8)).decode() + ".mp4")  # 랜덤 파일 이름 생성
                t[0].download(filename=filename)  # 동영상 다운로드
        except:
                warnings.warn(f"Video Not Found at {yt_link}")  # 경고 메시지 출력
                filename = None
        
        return filename, filename
        
def inference(file_uploaded, composer):
        waveform, sr = librosa.load(file_uploaded, sr=None)  # 파일에서 오디오 데이터와 샘플링 주파수 로드
        
        inputs = processor(audio=waveform, sampling_rate=sr, return_tensors="pt").to(device)  # 입력 데이터 전처리
        model_output = model.generate(input_features=inputs["input_features"], composer=composer)  # 모델에 입력하여 출력 생성
        tokenizer_output = processor.batch_decode(token_ids=model_output.to("cpu"), feature_extractor_output=inputs.to("cpu"))["pretty_midi_objects"]  # 토큰 디코딩
        
        return prepare_output_file(tokenizer_output, sr)  # 출력 파일 준비 함수 호출

def prepare_output_file(tokenizer_output, sr):
        output_file_name = "output_" + binascii.hexlify(os.urandom(8)).decode()  # 랜덤 출력 파일 이름 생성
        midi_output = os.path.join(outputs_dir, output_file_name + ".mid")  # MIDI 출력 파일 경로
        
        tokenizer_output[0].write(midi_output)  # MIDI 파일 작성
        
        wav_output = midi_output.replace(".mid", ".wav")  # WAV 출력 파일 경로
        midi2audio.FluidSynth().midi_to_audio(midi_output, wav_output)  # MIDI를 WAV로 변환
        
        return wav_output, wav_output, midi_output  # WAV 및 MIDI 파일 경로 반환

def get_stereo(pop_path, midi, pop_scale=0.5):
        pop_y, sr = librosa.load(pop_path, sr=None)  # 팝 음악 파일 로드
        midi_y, _ = librosa.load(midi.name, sr=None)  # MIDI 파일 로드

        if len(pop_y) > len(midi_y):
                midi_y = np.pad(midi_y, (0, len(pop_y) - len(midi_y)))  # MIDI 길이를 팝 음악 길이에 맞춤
        elif len(pop_y) < len(midi_y):
                pop_y = np.pad(pop_y, (0, -len(pop_y) + len(midi_y)))  # 팝 음악 길이를 MIDI 길이에 맞춤
        stereo = np.stack((midi_y, pop_y * pop_scale))  # 스테레오 믹스 생성
        
        stereo_mix_path = pop_path.replace("output", "output_stereo_mix")  # 스테레오 믹스 파일 경로
        sf.write(file=stereo_mix_path, data=stereo.T, samplerate=sr, format="wav")  # 스테레오 믹스 파일 작성
        
        return stereo_mix_path, stereo_mix_path  # 스테레오 믹스 파일 경로 반환

block = gr.Blocks(theme="Taithrah/Deep")  # Gradio 블록 생성

with block:
    gr.HTML(
        """
        <div style="text-align: center; max-width: 800px; margin: 0 auto;">
            <div
                style="
                    display: inline-flex;
                    align-items: center;
                    gap: 0.8rem;
                    font-size: 1.75rem;
                "
            >
                <h1 style="font-weight: 900; margin-bottom: 12px;">
                    🎹 Pop2Piano : 피아노 커버곡 생성기 🎹
                </h1>
            </div>
            <p style="margin-bottom: 10px; font-size: 90%">
                A demo for Pop2Piano: Pop Audio-based Piano Cover Generation. <br>
                Please select the composer (Arranger) and upload the pop audio or enter the YouTube link and then click Generate.
            </p>
        </div>
        """
    )
    with gr.Group():
        with gr.Row(equal_height=True):
            with gr.Column():
                file_uploaded = gr.Audio(label="오디오 업로드", type="filepath")
            with gr.Column():
                with gr.Row():
                    yt_link = gr.Textbox(label="유튜브 링크를 입력하세요.", autofocus=True, lines=3)
                    yt_btn = gr.Button("유튜브 링크에서 오디오를 다운 받습니다.", size="lg")

        yt_audio_path = gr.Audio(label="유튜브 동영상에서 추출한 오디오", interactive=False)
    yt_btn.click(get_audio_from_yt_video, inputs=[yt_link], outputs=[yt_audio_path, file_uploaded])

    with gr.Group():
        with gr.Column():
            composer = gr.Dropdown(label="편곡자 선택", choices=composers, value="composer1")
            generate_btn = gr.Button("나만의 피아노 커버곡 만들기 🎹  🎵")

        with gr.Row().style(mobile_collapse=False, equal_height=True):
            wav_output2 = gr.File(label="나만의 피아노 커버곡을 다운로드 (.wav)")
            wav_output1 = gr.Audio(label="나만의 피아노 커버곡 듣기")
            midi_output = gr.File(label="생성한 midi 파일 다운로드 (.mid)")
            generate_btn.click(
                inference, inputs=[file_uploaded, composer], outputs=[wav_output1, wav_output2, midi_output]
            )

    with gr.Group():
        gr.HTML(
            """
            <div> <h3> <center> 원본 오디오와 MIDI를 혼합하여 스테레오 믹스 하기. </h3> </div>
            """
        )
        pop_scale = gr.Slider(
            0, 1, value=0.5, label="원곡과 MIDI 사이에서 비율을 조정해 보세요.", info="1.0 = 원곡, 0.0=.mid", interactive=True
        ),
        stereo_btn = gr.Button("Get Stereo Mix")
        with gr.Row():
            stereo_mix1 = gr.Audio(label="스테레오 믹스 듣기")
            stereo_mix2 = gr.File(label="스테레오 믹스 다운로드")

        stereo_btn.click(get_stereo, inputs=[file_uploaded, wav_output2, pop_scale[0]], outputs=[stereo_mix1, stereo_mix2])

    gr.HTML(
        """
        <div class="footer">
            <center>The design for this Space is taken from DataBassist </a>
        </div>
        """
    )

    gr.HTML(
        """
        <div class="footer">
            <center><p><a href="http://sweetcocoa.github.io/pop2piano_samples" style="text-decoration: underline;" target="_blank">Project Page</a> 
            <center><a href="https://huggingface.co/docs/transformers/main/model_doc/pop2piano" style="text-decoration: underline;" target="_blank">HuggingFace Model Docs</a>
            <center><a href="https://github.com/sweetcocoa/pop2piano" style="text-decoration: underline;" target="_blank">Github</a>
            </p>
        </div>
        """
    )

block.launch(debug=False)