Update app.py

app.py

@@ -1,137 +1,98 @@
-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 모델과 전처리기
+import os
+import torch
+import librosa
+import binascii
+import warnings
+import midi2audio  # MIDI 파일을 WAV 파일로 변환
+import numpy as np
+import pytube as pt  # YouTube 비디오를 오디오로 다운로드
+import gradio as gr
+import soundfile as sf
+from transformers import Pop2PianoForConditionalGeneration, Pop2PianoProcessor
 
-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 동영상 다운로드 디렉토리 생성 (이미 존재하는 경우 무시)
+# 디렉토리 생성
+yt_video_dir = "./yt_dir"  # 유튜브 비디오 다운로드 경로
+outputs_dir = "./midi_wav_outputs"  # 출력 파일 경로
+os.makedirs(outputs_dir, exist_ok=True)
+os.makedirs(yt_video_dir, exist_ok=True)
 
-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()  # 작곡가 목록 가져오기
+# 모델 설정
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = Pop2PianoForConditionalGeneration.from_pretrained("sweetcocoa/pop2piano").to(device)
+processor = Pop2PianoProcessor.from_pretrained("sweetcocoa/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)  # 동영상 다운로드
+        yt = pt.YouTube(yt_link)
+        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}")  # 경고 메시지 출력
+        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 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"]
 
-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 파일 로드
+    return prepare_output_file(tokenizer_output, sr)    
 
-    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  # 스테레오 믹스 파일 경로 반환
+# 출력 파일 준비 함수
+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")
+    tokenizer_output[0].write(midi_output)
+    wav_output = midi_output.replace(".mid", ".wav")
+    midi2audio.FluidSynth().midi_to_audio(midi_output, wav_output)
+    
+    return wav_output, wav_output, midi_output
 
-block = gr.Blocks("Taithrah/Minimal")  # Gradio 블록 생성
+# Gradio UI 설정
+block = gr.Blocks(theme="Taithrah/Minimal")
 
 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>
+            <h1 style="font-weight: 900; margin-bottom: 12px;">
+                🎹 Pop2Piano : 피아노 커버곡 생성기 🎹
+            </h1>
             <p style="margin-bottom: 12px; 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.
+                Pop2Piano 데모: 팝 오디오 기반 피아노 커버곡 생성. <br>
+                작곡가(편곡자)를 선택하고 팝 오디오를 업로드하거나 유튜브 링크를 입력한 후 생성 버튼을 클릭하세요.
             </p>
         </div>
         """
     )
     with gr.Group():
-        with gr.Row(equal_height=True):
+        with gr.Row():
             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])
+                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="Arranger", choices=composers, value="composer1")
+            composer = gr.Dropdown(label="편곡자", choices=composers, value="composer1")
             generate_btn = gr.Button("나만의 피아노 커버곡 만들기🎹🎵")
-
-
-        
-        with gr.Row().style(mobile_collapse=False, equal_height=True):
+        with gr.Row():
             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])
-        
-            
-
-    
-
-        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>
-        """
-        )
+            midi_output = gr.File(label="생성한 midi 파일 다운로드 (.mid)")
+            generate_btn.click(
+                inference, 
+                inputs=[file_uploaded, composer], 
+                outputs=[wav_output1, wav_output2, midi_output])
 
-block.launch(debug=False)
+block.launch(debug=False)