Update app.py

app.py

@@ -65,40 +65,50 @@ def calculate_generation_params(lyrics):
     # 기본 시간 계산 (초 단위)
     time_per_line = {
         'verse': 4,    # verse는 한 줄당 4초
-        'chorus': 6,   # chorus는 한 줄당 6초 (더 긴 시간 할당)
+        'chorus': 6,   # chorus는 한 줄당 6초
         'bridge': 5    # bridge는 한 줄당 5초
     }
     
     # 각 섹션별 예상 시간 계산
-    total_duration = 0
-    for section_type, lines in section_lines.items():
-        total_duration += lines * time_per_line[section_type]
+    section_durations = {
+        'verse': section_lines['verse'] * time_per_line['verse'],
+        'chorus': section_lines['chorus'] * time_per_line['chorus'],
+        'bridge': section_lines['bridge'] * time_per_line['bridge']
+    }
+    
+    total_duration = sum(section_durations.values())
     
-    # 최소 지속 시간 보장 (60초)
-    total_duration = max(60, total_duration)
+    # 최소 지속 시간 보장 (90초)
+    total_duration = max(90, total_duration)
     
-    # 토큰 계산 (1초당 약 50토큰으로 계산)
-    tokens_per_second = 50
-    total_tokens = int(total_duration * tokens_per_second)
+    # 토큰 계산 (1초당 약 100토큰으로 증가)
+    tokens_per_second = 100
+    base_tokens = int(total_duration * tokens_per_second)
+    
+    # 코러스가 있는 경우 추가 토큰 할당
+    if sections['chorus'] > 0:
+        chorus_tokens = int(section_durations['chorus'] * tokens_per_second * 1.5)
+        total_tokens = base_tokens + chorus_tokens
+    else:
+        total_tokens = base_tokens
     
     # 섹션 기반 세그먼트 수 계산
-    if total_duration > 180:    # 3분 이상
-        num_segments = 4
-    elif total_duration > 120:  # 2분 이상
-        num_segments = 3
-    else:                      # 2분 미만
-        num_segments = 2
+    if sections['chorus'] > 0:
+        num_segments = max(3, sections['verse'] + sections['chorus'])
+    else:
+        num_segments = max(2, total_sections)
     
-    # 토큰 수 제한 (최소 6000토큰 보장)
-    max_tokens = min(32000, max(6000, total_tokens))
+    # 토큰 수 제한 (최소 8000토큰 보장)
+    max_tokens = min(32000, max(8000, total_tokens))
     
     return {
         'max_tokens': max_tokens,
-        'num_segments': num_segments,
+        'num_segments': min(4, num_segments),  # 최대 4개 세그먼트로 제한
         'sections': sections,
         'section_lines': section_lines,
         'estimated_duration': total_duration,
-        'tokens_per_segment': max_tokens // num_segments
+        'section_durations': section_durations,
+        'has_chorus': sections['chorus'] > 0
     }
 
 def get_audio_duration(file_path):
@@ -277,20 +287,23 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         logging.info(f"Lyrics analysis: {params}")
         
         # 코러스 섹션 확인 및 로깅
-        has_chorus = params['sections']['chorus'] > 0
-        estimated_duration = params.get('estimated_duration', 60)
+        has_chorus = params['has_chorus']
+        estimated_duration = params.get('estimated_duration', 90)
         
-        # 토큰 수 조정 (코러스가 있는 경우 더 많은 토큰 할당)
+        # 토큰 수와 세그먼트 수 조정
         if has_chorus:
-            actual_max_tokens = int(config['max_tokens'] * 1.5)  # 50% 더 많은 토큰
-            actual_num_segments = max(3, config['num_segments'])  # 최소 3개 세그먼트 보장
+            actual_max_tokens = int(params['max_tokens'] * 1.5)  # 50% 더 많은 토큰
+            actual_num_segments = max(3, params['num_segments'])  # 최소 3개 세그먼트
+            tokens_per_segment = actual_max_tokens // actual_num_segments
         else:
-            actual_max_tokens = config['max_tokens']
-            actual_num_segments = config['num_segments']
+            actual_max_tokens = params['max_tokens']
+            actual_num_segments = params['num_segments']
+            tokens_per_segment = actual_max_tokens // actual_num_segments
         
         logging.info(f"Estimated duration: {estimated_duration} seconds")
         logging.info(f"Has chorus sections: {has_chorus}")
         logging.info(f"Using segments: {actual_num_segments}, tokens: {actual_max_tokens}")
+        logging.info(f"Tokens per segment: {tokens_per_segment}")
         
         # 임시 파일 생성
         genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_")
@@ -314,9 +327,12 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
             "--max_new_tokens", str(actual_max_tokens)
         ]
 
-        # GPU가 있는 경우에만 추가 옵션 적용
+        # GPU 설정
         if torch.cuda.is_available():
-            command.append("--disable_offload_model")
+            command.extend([
+                "--disable_offload_model",
+                "--use_bf16"  # 더 빠른 처리를 위한 BF16 사용
+            ])
 
         # CUDA 환경 변수 설정
         env = os.environ.copy()
@@ -326,7 +342,7 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
                 "CUDA_HOME": "/usr/local/cuda",
                 "PATH": f"/usr/local/cuda/bin:{env.get('PATH', '')}",
                 "LD_LIBRARY_PATH": f"/usr/local/cuda/lib64:{env.get('LD_LIBRARY_PATH', '')}",
-                "PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:512"
+                "PYTORCH_CUDA_ALLOC_CONF": f"max_split_size_mb:512"
             })
 
         # transformers 캐시 마이그레이션 처리
@@ -366,7 +382,7 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
                     logging.info(f"Expected duration: {estimated_duration} seconds")
                     
                     # 생성된 음악이 너무 짧은 경우 경고
-                    if duration < estimated_duration * 0.8:  # 예상 길이의 80% 미만인 경우
+                    if duration < estimated_duration * 0.8:
                         logging.warning(f"Generated audio is shorter than expected: {duration:.2f}s < {estimated_duration:.2f}s")
             except Exception as e:
                 logging.warning(f"Failed to get audio duration: {e}")