Update app.py

app.py

@@ -201,41 +201,104 @@ def install_flash_attn():
         logging.warning(f"Failed to install flash-attn: {e}")
         return False
 
-def initialize_system():
-    optimize_gpu_settings()
+
+def check_model_files():
+    required_files = [
+        "./xcodec_mini_infer/config.json",
+        "./xcodec_mini_infer/vocal_decoder.pth",
+        "./xcodec_mini_infer/inst_decoder.pth"
+    ]
     
-    with ThreadPoolExecutor(max_workers=4) as executor:
-        futures = []
-        
-        futures.append(executor.submit(install_flash_attn))
-        
+    missing_files = [f for f in required_files if not os.path.exists(f)]
+    if missing_files:
+        logging.warning(f"Missing required files: {missing_files}")
+        download_missing_files()
+
+def download_missing_files():
+    try:
         from huggingface_hub import snapshot_download
         
-        # 필요한 모든 모델 다운로드
-        models_to_download = [
-            "m-a-p/xcodec_mini_infer",
+        # xcodec_mini_infer 모델 다운로드
+        snapshot_download(
+            repo_id="m-a-p/xcodec_mini_infer",
+            local_dir="./xcodec_mini_infer",
+            resume_download=True
+        )
+        
+        # YuE 모델들 다운로드
+        models = [
             "m-a-p/YuE-s1-7B-anneal-jp-kr-cot",
             "m-a-p/YuE-s1-7B-anneal-en-cot",
             "m-a-p/YuE-s1-7B-anneal-zh-cot",
             "m-a-p/YuE-s2-1B-general"
         ]
         
-        for model in models_to_download:
-            futures.append(executor.submit(
-                snapshot_download,
+        for model in models:
+            model_name = model.split('/')[-1]
+            snapshot_download(
                 repo_id=model,
-                local_dir=f"./inference/models/{model.split('/')[-1]}",
+                local_dir=f"./models/{model_name}",
                 resume_download=True
-            ))
-        
-        for future in futures:
-            future.result()
-
+            )
+            
+        logging.info("All required models downloaded successfully")
+    except Exception as e:
+        logging.error(f"Error downloading models: {e}")
+        raise
+def initialize_system():
+    optimize_gpu_settings()
+    
     try:
+        # 기본 디렉토리 구조 생성
+        os.makedirs("./inference", exist_ok=True)
+        os.makedirs("./inference/models", exist_ok=True)
+        os.makedirs("./inference/models/cache", exist_ok=True)
+        os.makedirs("./inference/xcodec_mini_infer", exist_ok=True)
+        
+        # 작업 디렉토리 변경
         os.chdir("./inference")
         logging.info(f"Working directory changed to: {os.getcwd()}")
-    except FileNotFoundError as e:
-        logging.error(f"Directory error: {e}")
+        
+        # 모델 파일 체크 및 다운로드
+        check_model_files()
+        
+        with ThreadPoolExecutor(max_workers=4) as executor:
+            futures = []
+            
+            # Flash Attention 설치
+            futures.append(executor.submit(install_flash_attn))
+            
+            from huggingface_hub import snapshot_download
+            
+            # 필요한 모든 모델 다운로드
+            models_to_download = [
+                "m-a-p/xcodec_mini_infer",
+                "m-a-p/YuE-s1-7B-anneal-jp-kr-cot",
+                "m-a-p/YuE-s1-7B-anneal-en-cot",
+                "m-a-p/YuE-s1-7B-anneal-zh-cot",
+                "m-a-p/YuE-s2-1B-general"
+            ]
+            
+            for model in models_to_download:
+                model_name = model.split('/')[-1]
+                model_dir = f"./models/{model_name}"
+                
+                if not os.path.exists(model_dir):
+                    futures.append(executor.submit(
+                        snapshot_download,
+                        repo_id=model,
+                        local_dir=model_dir,
+                        resume_download=True
+                    ))
+            
+            # 모든 다운로드 완료 대기
+            for future in futures:
+                future.result()
+        
+        logging.info("System initialization completed successfully")
+        
+    except Exception as e:
+        logging.error(f"Initialization error: {e}")
         raise
 
 @lru_cache(maxsize=100)
@@ -279,11 +342,16 @@ def get_audio_duration(file_path):
         logging.error(f"Failed to get audio duration: {e}")
         return None
 
+
+
 def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
+    check_model_files()  # 필요한 파일 체크 및 다운로드
+    
     genre_txt_path = None
     lyrics_txt_path = None
     
     try:
+
         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}")