Update app.py

app.py

@@ -202,32 +202,75 @@ def install_flash_attn():
 def initialize_system():
     optimize_gpu_settings()
     
-    with ThreadPoolExecutor(max_workers=4) as executor:
-        futures = []
+    try:
+        # 기본 디렉토리 구조 생성
+        base_dir = os.path.abspath("./inference")
+        os.makedirs(base_dir, exist_ok=True)
+        os.makedirs(os.path.join(base_dir, "models"), exist_ok=True)
         
-        futures.append(executor.submit(install_flash_attn))
+        # 작업 디렉토리 변경
+        os.chdir(base_dir)
+        logging.info(f"Working directory changed to: {os.getcwd()}")
         
-        from huggingface_hub import snapshot_download
+        from huggingface_hub import snapshot_download, hf_hub_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,
+        # xcodec_mini_infer 다운로드
+        xcodec_dir = "xcodec_mini_infer"
+        os.makedirs(xcodec_dir, exist_ok=True)
+        
+        # 필수 파일 다운로드
+        snapshot_download(
             repo_id="m-a-p/xcodec_mini_infer",
-            local_dir="./inference/xcodec_mini_infer",
-            resume_download=True
-        ))
+            local_dir=xcodec_dir,
+            force_download=True
+        )
         
-        for future in futures:
-            future.result()
-
-    try:
-        os.chdir("./inference")
-        logging.info(f"Working directory changed to: {os.getcwd()}")
-    except FileNotFoundError as e:
-        logging.error(f"Directory error: {e}")
+        # infer.py 파일 다운로드
+        infer_script = hf_hub_download(
+            repo_id="m-a-p/xcodec_mini_infer",
+            filename="infer.py",
+            force_download=True
+        )
+        
+        # infer.py 파일을 현재 디렉토리로 복사
+        shutil.copy2(infer_script, "infer.py")
+        
+        # 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("models", model_name)
+            snapshot_download(
+                repo_id=model,
+                local_dir=model_path,
+                force_download=True
+            )
+        
+        # 필요한 파일들 존재 확인
+        required_files = [
+            "infer.py",
+            os.path.join(xcodec_dir, "config.json"),
+            os.path.join(xcodec_dir, "vocal_decoder.pth"),
+            os.path.join(xcodec_dir, "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}")
+            else:
+                file_size = os.path.getsize(file_path)
+                logging.info(f"Verified {file_path}: {file_size} bytes")
+        
+        logging.info("System initialization completed successfully")
+        
+    except Exception as e:
+        logging.error(f"Initialization error: {e}")
         raise
 
 @lru_cache(maxsize=100)
@@ -305,21 +348,22 @@ 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",
+            "python", "./infer.py",  # 상대 경로로 수정
             "--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", output_dir,
+            "--output_dir", "./output",
             "--cuda_idx", "0",
             "--max_new_tokens", str(actual_max_tokens),
-            "--disable_offload_model"  # GPU 메모리 최적화를 위해 추가
+            "--disable_offload_model"
         ]
 
+
+
         env = os.environ.copy()
         if torch.cuda.is_available():
             env.update({
@@ -496,7 +540,7 @@ def main():
             with gr.Row(elem_id="header"):
                 gr.Markdown(
                     """
-                    # 🎵 AI Song Creator
+                    # 🎵 AI Song Creator 'Open SUNO'
                     ### Transform Your Lyrics into Complete Songs with Music
                     Create professional songs from your lyrics in multiple languages
                     """