Update app.py

app.py

@@ -325,34 +325,58 @@ def initialize_system():
     optimize_gpu_settings()
     
     try:
-        # 현재 작업 디렉토리 저장
-        original_dir = os.getcwd()
-        
         # 기본 디렉토리 구조 생성
-        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.makedirs("./inference/xcodec_mini_infer/checkpoints", exist_ok=True)
+        base_dir = os.path.abspath("./inference")
+        os.makedirs(base_dir, exist_ok=True)
+        os.makedirs(os.path.join(base_dir, "models"), exist_ok=True)
         
-        # inference 디렉토리로 이동
-        os.chdir("./inference")
-        logging.info(f"Working directory changed to: {os.getcwd()}")
+        # xcodec_mini_infer 다운로드
+        from huggingface_hub import snapshot_download
         
-        # 환경 변수 설정
-        cache_dir = os.path.abspath("./models/cache")
-        os.environ["HF_HOME"] = cache_dir
-        os.environ["TRANSFORMERS_CACHE"] = cache_dir
-        os.environ["HF_HUB_CACHE"] = cache_dir
+        xcodec_path = os.path.join(base_dir, "xcodec_mini_infer")
+        os.makedirs(xcodec_path, exist_ok=True)
         
-        # 모델 파일 체크 및 다운로드
-        check_model_files()
+        # xcodec_mini_infer 모델 다운로드
+        snapshot_download(
+            repo_id="m-a-p/xcodec_mini_infer",
+            local_dir=xcodec_path,
+            resume_download=True
+        )
         
-        logging.info("System initialization completed successfully")
+        # 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:
+            model_name = model.split('/')[-1]
+            model_path = os.path.join(base_dir, "models", model_name)
+            snapshot_download(
+                repo_id=model,
+                local_dir=model_path,
+                resume_download=True
+            )
+        
+        # 작업 디렉토리 변경
+        os.chdir(base_dir)
+        logging.info(f"Working directory changed to: {os.getcwd()}")
+        
+        # 필요한 파일 존재 확인
+        required_files = [
+            os.path.join("xcodec_mini_infer", "config.json"),
+            os.path.join("xcodec_mini_infer", "vocal_decoder.pth"),
+            os.path.join("xcodec_mini_infer", "inst_decoder.pth")
+        ]
+        
+        for file_path in required_files:
+            if not os.path.exists(file_path):
+                raise FileNotFoundError(f"Required file not found: {file_path}")
+            
     except Exception as e:
-        logging.error(f"Initialization error: {e}")
-        os.chdir(original_dir)
+        logging.error(f"Directory error: {e}")
         raise
 
 @lru_cache(maxsize=100)
@@ -435,24 +459,19 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
         os.makedirs(output_dir, exist_ok=True)
         empty_output_folder(output_dir)
 
-
         command = [
             "python", "infer.py",
-            "--stage1_model", f"./models/{model_path.split('/')[-1]}",
-            "--stage2_model", "./models/YuE-s2-1B-general",
-            "--genre_txt", os.path.abspath(genre_txt_path),
-            "--lyrics_txt", os.path.abspath(lyrics_txt_path),
+            "--stage1_model", model_path,  # 원래 모델 경로 사용
+            "--stage2_model", "m-a-p/YuE-s2-1B-general",
+            "--genre_txt", genre_txt_path,
+            "--lyrics_txt", lyrics_txt_path,
             "--run_n_segments", str(actual_num_segments),
             "--stage2_batch_size", "16",
-            "--output_dir", os.path.abspath(output_dir),
+            "--output_dir", output_dir,
             "--cuda_idx", "0",
             "--max_new_tokens", str(actual_max_tokens),
-            "--disable_offload_model",
-            "--basic_model_config", os.path.abspath("./xcodec_mini_infer/config.json"),
-            "--vocal_decoder_path", os.path.abspath("./xcodec_mini_infer/vocal_decoder.pth"),
-            "--inst_decoder_path", os.path.abspath("./xcodec_mini_infer/inst_decoder.pth")
+            "--disable_offload_model"
         ]
-        
 
 
         env = os.environ.copy()