Update app.py

app.py

@@ -20,6 +20,70 @@ logging.basicConfig(
     ]
 )
 
+# 가사 분석 함수
+def analyze_lyrics(lyrics):
+    # 줄 단위로 분리
+    lines = [line.strip() for line in lyrics.split('\n') if line.strip()]
+    
+    # 섹션 카운트
+    sections = {
+        'verse': 0,
+        'chorus': 0,
+        'bridge': 0,
+        'total_lines': len(lines)
+    }
+    
+    current_section = None
+    section_lines = {
+        'verse': 0,
+        'chorus': 0,
+        'bridge': 0
+    }
+    
+    for line in lines:
+        lower_line = line.lower()
+        if '[verse]' in lower_line:
+            current_section = 'verse'
+            sections['verse'] += 1
+        elif '[chorus]' in lower_line:
+            current_section = 'chorus'
+            sections['chorus'] += 1
+        elif '[bridge]' in lower_line:
+            current_section = 'bridge'
+            sections['bridge'] += 1
+        elif current_section and line.strip():
+            section_lines[current_section] += 1
+            
+    # 총 섹션 수 계산
+    total_sections = sections['verse'] + sections['chorus'] + sections['bridge']
+    
+    return sections, total_sections, len(lines), section_lines
+
+def calculate_generation_params(lyrics):
+    sections, total_sections, total_lines, section_lines = analyze_lyrics(lyrics)
+    
+    # 기본 토큰 수 계산
+    base_tokens_per_line = 200
+    verse_tokens = section_lines['verse'] * base_tokens_per_line
+    chorus_tokens = section_lines['chorus'] * (base_tokens_per_line * 1.5)  # 코러스는 50% 더 많은 토큰
+    bridge_tokens = section_lines['bridge'] * base_tokens_per_line
+    
+    # 총 토큰 수 계산
+    total_tokens = int(verse_tokens + chorus_tokens + bridge_tokens)
+    
+    # 섹션 기반 세그먼트 수 계산
+    num_segments = max(2, min(4, total_sections))
+    
+    # 토큰 수 제한
+    max_tokens = min(32000, max(3000, total_tokens))
+    
+    return {
+        'max_tokens': max_tokens,
+        'num_segments': num_segments,
+        'sections': sections,
+        'section_lines': section_lines
+    }
+
 # 언어 감지 및 모델 선택 함수
 def detect_and_select_model(text):
     if re.search(r'[\u3131-\u318E\uAC00-\uD7A3]', text):  # 한글
@@ -33,26 +97,33 @@ def detect_and_select_model(text):
 
 def optimize_model_selection(lyrics, genre):
     model_path = detect_and_select_model(lyrics)
+    params = calculate_generation_params(lyrics)
     
     model_config = {
         "m-a-p/YuE-s1-7B-anneal-en-cot": {
-            "max_tokens": 24000,
+            "max_tokens": params['max_tokens'],
             "temperature": 0.8,
-            "batch_size": 8
+            "batch_size": 8,
+            "num_segments": params['num_segments'],
+            "chorus_strength": 1.2 if params['sections']['chorus'] > 0 else 1.0
         },
         "m-a-p/YuE-s1-7B-anneal-jp-kr-cot": {
-            "max_tokens": 24000,
+            "max_tokens": params['max_tokens'],
             "temperature": 0.7,
-            "batch_size": 8
+            "batch_size": 8,
+            "num_segments": params['num_segments'],
+            "chorus_strength": 1.2 if params['sections']['chorus'] > 0 else 1.0
         },
         "m-a-p/YuE-s1-7B-anneal-zh-cot": {
-            "max_tokens": 24000,
+            "max_tokens": params['max_tokens'],
             "temperature": 0.7,
-            "batch_size": 8
+            "batch_size": 8,
+            "num_segments": params['num_segments'],
+            "chorus_strength": 1.2 if params['sections']['chorus'] > 0 else 1.0
         }
     }
     
-    return model_path, model_config[model_path]
+    return model_path, model_config[model_path], params
 
 # GPU 설정 최적화
 def optimize_gpu_settings():
@@ -138,11 +209,22 @@ def get_last_mp3_file(output_dir):
     mp3_files_with_path.sort(key=os.path.getmtime, reverse=True)
     return mp3_files_with_path[0]
 
+
+
+
+
 def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
     try:
         # 모델 선택 및 설정
-        model_path, config = optimize_model_selection(lyrics_txt_content, genre_txt_content)
+        model_path, config, params = optimize_model_selection(lyrics_txt_content, genre_txt_content)
         logging.info(f"Selected model: {model_path}")
+        logging.info(f"Lyrics analysis: {params}")
+        
+        # 실제 사용할 파라미터
+        actual_num_segments = config['num_segments']
+        actual_max_tokens = config['max_tokens']
+        
+        logging.info(f"Using segments: {actual_num_segments}, tokens: {actual_max_tokens}")
         
         # 임시 파일 생성
         genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_")
@@ -159,15 +241,16 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
             "--stage2_model", "m-a-p/YuE-s2-1B-general",
             "--genre_txt", genre_txt_path,
             "--lyrics_txt", lyrics_txt_path,
-            "--run_n_segments", str(num_segments),
+            "--run_n_segments", str(actual_num_segments),
             "--stage2_batch_size", str(config['batch_size']),
             "--output_dir", output_dir,
             "--cuda_idx", "0",
-            "--max_new_tokens", str(config['max_tokens']),
+            "--max_new_tokens", str(actual_max_tokens),
             "--temperature", str(config['temperature']),
             "--disable_offload_model",
             "--use_flash_attention_2",
-            "--bf16"
+            "--bf16",
+            "--chorus_strength", str(config['chorus_strength'])
         ]
 
         # CUDA 환경 변수 설정
@@ -228,31 +311,31 @@ with gr.Blocks() as demo:
                 )
                 lyrics_txt = gr.Textbox(
                     label="Lyrics (Supports English, Korean, Japanese, Chinese)",
-                    placeholder="Enter song lyrics...",
+                    placeholder="Enter song lyrics with [verse], [chorus], [bridge] tags...",
                     lines=10
                 )
                 
             with gr.Column():
                 num_segments = gr.Number(
-                    label="Number of Song Segments",
+                    label="Number of Song Segments (Auto-adjusted based on lyrics)",
                     value=2,
                     minimum=1,
                     maximum=4,
                     step=1,
-                    interactive=True
+                    interactive=False
                 )
                 max_new_tokens = gr.Slider(
-                    label="Max New Tokens",
+                    label="Max New Tokens (Auto-adjusted based on lyrics)",
                     minimum=500,
                     maximum=32000,
                     step=500,
                     value=4000,
-                    interactive=True
+                    interactive=False
                 )
                 submit_btn = gr.Button("Generate Music", variant="primary")
                 music_out = gr.Audio(label="Generated Audio")
 
-        # 다국어 예제 추가
+        # 다국어 예제
         gr.Examples(
             examples=[
                 # 영어 예제
@@ -264,6 +347,18 @@ Whispers of the night wind echo through the hall
 Lost within the silence, I hear your gentle voice
 Guiding me back homeward, making my heart rejoice
 
+[chorus]
+Don't let this moment fade, hold me close tonight
+With you here beside me, everything's alright
+Can't imagine life alone, don't want to let you go
+Stay with me forever, let our love just flow
+
+[verse]
+Morning light is breaking, through the window pane
+Memories of yesterday, like soft summer rain
+In your arms I'm finding, all I'm dreaming of
+Every day beside you, fills my heart with love
+
 [chorus]
 Don't let this moment fade, hold me close tonight
 With you here beside me, everything's alright
@@ -285,38 +380,18 @@ Stay with me forever, let our love just flow
 두려움은 없어 너와 함께라면
 영원히 계속될 우리의 노래
 이 순간을 기억해 forever
-                    """
-                ],
-                # 일본어 예제
-                [
-                    "J-pop melodic soft piano emotional",
-                    """[verse]
-春の風に乗って
-思い出が流れる
-あの日の約束を
-今でも覚えてる
 
-[chorus]
-君と見た空は
-今も変わらないよ
-どこまでも続く
-この道の先で
-                    """
-                ],
-                # 중국어 예제
-                [
-                    "Chinese pop traditional fusion modern",
-                    """[verse]
-晨光照亮天际
-新的一天开始
-追逐着梦想前进
-不停歇的脚步
+[verse]
+새로운 내일을 향해 나아가
+우리만의 길을 만들어가
+믿음으로 가득한 우리의 맘
+절대 멈추지 않아 계속해서
 
 [chorus]
-让希望照亮前方
-让勇气伴随身旁
-这一路有你相伴
-永远不会孤单
+달려가자 더 높이 더 멀리
+두려움은 없어 너와 함께라면
+영원히 계속될 우리의 노래
+이 순간을 기억해 forever
                     """
                 ]
             ],
@@ -341,4 +416,4 @@ demo.queue(concurrency_count=2).launch(
     enable_queue=True,
     show_api=True,
     show_error=True
-)
+)