Update app.py

app.py

@@ -77,33 +77,26 @@ def calculate_generation_params(lyrics):
     }
     
     total_duration = sum(section_durations.values())
+    total_duration = max(60, total_duration)  # 최소 60초
     
-    # 최소 지속 시간 보장 (90초)
-    total_duration = max(90, total_duration)
+    # 토큰 계산 (더 보수적인 값 사용)
+    base_tokens = 3000  # 기본 토큰 수
+    tokens_per_line = 200  # 줄당 토큰 수
     
-    # 토큰 계산 (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
+    total_tokens = base_tokens + (total_lines * tokens_per_line)
     
     # 섹션 기반 세그먼트 수 계산
     if sections['chorus'] > 0:
-        num_segments = max(3, sections['verse'] + sections['chorus'])
+        num_segments = 3  # 코러스가 있는 경우 3개 세그먼트
     else:
-        num_segments = max(2, total_sections)
+        num_segments = 2  # 코러스가 없는 경우 2개 세그먼트
     
-    # 토큰 수 제한 (최소 8000토큰 보장)
-    max_tokens = min(32000, max(8000, total_tokens))
+    # 토큰 수 제한
+    max_tokens = min(8000, total_tokens)  # 최대 8000 토큰으로 제한
     
     return {
         'max_tokens': max_tokens,
-        'num_segments': min(4, num_segments),  # 최대 4개 세그먼트로 제한
+        'num_segments': num_segments,
         'sections': sections,
         'section_lines': section_lines,
         'estimated_duration': total_duration,
@@ -294,14 +287,17 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         # 코러스 섹션 확인 및 로깅
         has_chorus = params['sections']['chorus'] > 0
         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 = min(8000, int(config['max_tokens'] * 1.2))  # 20% 증가, 최대 8000
+            actual_num_segments = 3
         else:
             actual_max_tokens = config['max_tokens']
-            actual_num_segments = config['num_segments']
+            actual_num_segments = 2
+        
+
         
         logging.info(f"Estimated duration: {estimated_duration} seconds")
         logging.info(f"Has chorus sections: {has_chorus}")
@@ -314,7 +310,6 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         output_dir = "./output"
         os.makedirs(output_dir, exist_ok=True)
         empty_output_folder(output_dir)
-
         # 기본 명령어 구성
         command = [
             "python", "infer.py",
@@ -323,7 +318,7 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
             "--genre_txt", genre_txt_path,
             "--lyrics_txt", lyrics_txt_path,
             "--run_n_segments", str(actual_num_segments),
-            "--stage2_batch_size", str(config['batch_size']),
+            "--stage2_batch_size", "4",  # 배치 사이즈 감소
             "--output_dir", output_dir,
             "--cuda_idx", "0",
             "--max_new_tokens", str(actual_max_tokens)
@@ -331,9 +326,9 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
 
         # GPU 설정
         if torch.cuda.is_available():
-            command.extend([
-                "--disable_offload_model"
-            ])
+            command.append("--disable_offload_model")
+        # GPU 설정
+
 
         # CUDA 환경 변수 설정
         env = os.environ.copy()