Update app.py

app.py

@@ -131,46 +131,7 @@ def detect_and_select_model(text):
     else:  # 영어/기타
         return "m-a-p/YuE-s1-7B-anneal-en-cot"
 
-def optimize_model_selection(lyrics, genre):
-    model_path = detect_and_select_model(lyrics)
-    params = calculate_generation_params(lyrics)
-    
-    # 코러스 존재 여부에 따른 설정 조정
-    has_chorus = params['sections']['chorus'] > 0
-    
-    model_config = {
-        "m-a-p/YuE-s1-7B-anneal-en-cot": {
-            "max_tokens": params['max_tokens'],
-            "temperature": 0.8,
-            "batch_size": 8,
-            "num_segments": params['num_segments'],
-            "tokens_per_segment": params['tokens_per_segment'],
-            "estimated_duration": params['estimated_duration']
-        },
-        "m-a-p/YuE-s1-7B-anneal-jp-kr-cot": {
-            "max_tokens": params['max_tokens'],
-            "temperature": 0.7,
-            "batch_size": 8,
-            "num_segments": params['num_segments'],
-            "tokens_per_segment": params['tokens_per_segment'],
-            "estimated_duration": params['estimated_duration']
-        },
-        "m-a-p/YuE-s1-7B-anneal-zh-cot": {
-            "max_tokens": params['max_tokens'],
-            "temperature": 0.7,
-            "batch_size": 8,
-            "num_segments": params['num_segments'],
-            "tokens_per_segment": params['tokens_per_segment'],
-            "estimated_duration": params['estimated_duration']
-        }
-    }
-    
-    # 코러스가 있는 경우 토큰 수 증가
-    if has_chorus:
-        for config in model_config.values():
-            config['max_tokens'] = int(config['max_tokens'] * 1.5)  # 50% 더 많은 토큰 할당
-    
-    return model_path, model_config[model_path], params
+
 
 # GPU 설정 최적화
 def optimize_gpu_settings():
@@ -279,7 +240,51 @@ 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 optimize_model_selection(lyrics, genre):
+    model_path = detect_and_select_model(lyrics)
+    params = calculate_generation_params(lyrics)
+    
+    # 코러스 존재 여부에 따른 설정 조정
+    has_chorus = params['sections']['chorus'] > 0
+    
+    # 토큰 수 계산
+    tokens_per_segment = params['max_tokens'] // params['num_segments']
+    
+    model_config = {
+        "m-a-p/YuE-s1-7B-anneal-en-cot": {
+            "max_tokens": params['max_tokens'],
+            "temperature": 0.8,
+            "batch_size": 8,
+            "num_segments": params['num_segments'],
+            "estimated_duration": params['estimated_duration']
+        },
+        "m-a-p/YuE-s1-7B-anneal-jp-kr-cot": {
+            "max_tokens": params['max_tokens'],
+            "temperature": 0.7,
+            "batch_size": 8,
+            "num_segments": params['num_segments'],
+            "estimated_duration": params['estimated_duration']
+        },
+        "m-a-p/YuE-s1-7B-anneal-zh-cot": {
+            "max_tokens": params['max_tokens'],
+            "temperature": 0.7,
+            "batch_size": 8,
+            "num_segments": params['num_segments'],
+            "estimated_duration": params['estimated_duration']
+        }
+    }
+    
+    # 코러스가 있는 경우 토큰 수 증가
+    if has_chorus:
+        for config in model_config.values():
+            config['max_tokens'] = int(config['max_tokens'] * 1.5)  # 50% 더 많은 토큰 할당
+    
+    return model_path, model_config[model_path], params
+
 def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
+    genre_txt_path = None
+    lyrics_txt_path = None
+    
     try:
         # 모델 선택 및 설정
         model_path, config, params = optimize_model_selection(lyrics_txt_content, genre_txt_content)
@@ -287,23 +292,20 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         logging.info(f"Lyrics analysis: {params}")
         
         # 코러스 섹션 확인 및 로깅
-        has_chorus = params['has_chorus']
+        has_chorus = params['sections']['chorus'] > 0
         estimated_duration = params.get('estimated_duration', 90)
         
         # 토큰 수와 세그먼트 수 조정
         if has_chorus:
-            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
+            actual_max_tokens = int(config['max_tokens'] * 1.5)  # 50% 더 많은 토큰
+            actual_num_segments = max(3, config['num_segments'])  # 최소 3개 세그먼트
         else:
-            actual_max_tokens = params['max_tokens']
-            actual_num_segments = params['num_segments']
-            tokens_per_segment = actual_max_tokens // actual_num_segments
+            actual_max_tokens = config['max_tokens']
+            actual_num_segments = config['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_")
@@ -330,8 +332,7 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         # GPU 설정
         if torch.cuda.is_available():
             command.extend([
-                "--disable_offload_model",
-                "--use_bf16"  # 더 빠른 처리를 위한 BF16 사용
+                "--disable_offload_model"
             ])
 
         # CUDA 환경 변수 설정
@@ -396,12 +397,19 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         raise
     finally:
         # 임시 파일 정리
-        for file in [genre_txt_path, lyrics_txt_path]:
+        if genre_txt_path and os.path.exists(genre_txt_path):
+            try:
+                os.remove(genre_txt_path)
+                logging.debug(f"Removed temporary file: {genre_txt_path}")
+            except Exception as e:
+                logging.warning(f"Failed to remove temporary file {genre_txt_path}: {e}")
+        
+        if lyrics_txt_path and os.path.exists(lyrics_txt_path):
             try:
-                os.remove(file)
-                logging.debug(f"Removed temporary file: {file}")
+                os.remove(lyrics_txt_path)
+                logging.debug(f"Removed temporary file: {lyrics_txt_path}")
             except Exception as e:
-                logging.warning(f"Failed to remove temporary file {file}: {e}")
+                logging.warning(f"Failed to remove temporary file {lyrics_txt_path}: {e}")
 
 def main():
     # Gradio 인터페이스