Update app.py

app.py

@@ -211,16 +211,22 @@ def initialize_system():
         
         from huggingface_hub import snapshot_download
         
-        folder_path = './inference/xcodec_mini_infer'
-        os.makedirs(folder_path, exist_ok=True)
-        logging.info(f"Created folder at: {folder_path}")
-
-        futures.append(executor.submit(
-            snapshot_download,
-            repo_id="m-a-p/xcodec_mini_infer",
-            local_dir="./inference/xcodec_mini_infer",
-            resume_download=True
-        ))
+        # 필요한 모든 모델 다운로드
+        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:
+            futures.append(executor.submit(
+                snapshot_download,
+                repo_id=model,
+                local_dir=f"./inference/models/{model.split('/')[-1]}",
+                resume_download=True
+            ))
         
         for future in futures:
             future.result()
@@ -307,11 +313,10 @@ 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 - 지원되지 않는 인수 제거
         command = [
             "python", "infer.py",
-            "--stage1_model", model_path,
-            "--stage2_model", "m-a-p/YuE-s2-1B-general",
+            "--stage1_model", f"./models/{model_path.split('/')[-1]}",
+            "--stage2_model", "./models/YuE-s2-1B-general",
             "--genre_txt", genre_txt_path,
             "--lyrics_txt", lyrics_txt_path,
             "--run_n_segments", str(actual_num_segments),
@@ -319,7 +324,10 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
             "--output_dir", output_dir,
             "--cuda_idx", "0",
             "--max_new_tokens", str(actual_max_tokens),
-            "--disable_offload_model"  # GPU 메모리 최적화를 위해 추가
+            "--disable_offload_model",
+            "--basic_model_config", "./xcodec_mini_infer/config.json",
+            "--vocal_decoder_path", "./xcodec_mini_infer/vocal_decoder.pth",
+            "--inst_decoder_path", "./xcodec_mini_infer/inst_decoder.pth"
         ]
 
         env = os.environ.copy()
@@ -330,9 +338,14 @@ def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
                 "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",
-                "CUDA_LAUNCH_BLOCKING": "0"
+                "CUDA_LAUNCH_BLOCKING": "0",
+                "TRANSFORMERS_CACHE": "./models/cache",
+                "HF_HOME": "./models/cache"
             })
 
+            
+
+
         # transformers 캐시 마이그레이션 처리
         try:
             from transformers.utils import move_cache