Update app.py

app.py

@@ -330,18 +330,39 @@ def initialize_system():
         os.makedirs(base_dir, exist_ok=True)
         os.makedirs(os.path.join(base_dir, "models"), exist_ok=True)
         
-        # xcodec_mini_infer 다운로드
-        from huggingface_hub import snapshot_download
+        # 작업 디렉토리 변경
+        os.chdir(base_dir)
+        logging.info(f"Working directory changed to: {os.getcwd()}")
         
-        xcodec_path = os.path.join(base_dir, "xcodec_mini_infer")
-        os.makedirs(xcodec_path, exist_ok=True)
+        from huggingface_hub import snapshot_download, hf_hub_download
         
-        # xcodec_mini_infer 모델 다운로드
-        snapshot_download(
-            repo_id="m-a-p/xcodec_mini_infer",
-            local_dir=xcodec_path,
-            resume_download=True
-        )
+        # xcodec_mini_infer 파일들 직접 다운로드
+        xcodec_dir = os.path.join(base_dir, "xcodec_mini_infer")
+        os.makedirs(xcodec_dir, exist_ok=True)
+        
+        # 필수 파일 직접 다운로드
+        required_files = {
+            "config.json": "config.json",
+            "vocal_decoder.pth": "vocal_decoder.pth",
+            "inst_decoder.pth": "inst_decoder.pth"
+        }
+        
+        for file_name in required_files.keys():
+            try:
+                file_path = os.path.join(xcodec_dir, file_name)
+                if not os.path.exists(file_path):
+                    downloaded_path = hf_hub_download(
+                        repo_id="m-a-p/xcodec_mini_infer",
+                        filename=file_name,
+                        local_dir=xcodec_dir,
+                        force_download=True
+                    )
+                    if downloaded_path != file_path:
+                        shutil.copy2(downloaded_path, file_path)
+                    logging.info(f"Downloaded {file_name} to {file_path}")
+            except Exception as e:
+                logging.error(f"Error downloading {file_name}: {e}")
+                raise
         
         # YuE 모델들 다운로드
         models = [
@@ -351,30 +372,38 @@ def initialize_system():
             "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
-            )
+        with ThreadPoolExecutor(max_workers=4) as executor:
+            futures = []
+            
+            # Flash Attention 설치
+            futures.append(executor.submit(install_flash_attn))
+            
+            # 모델 다운로드
+            for model in models:
+                model_name = model.split('/')[-1]
+                model_path = os.path.join(base_dir, "models", model_name)
+                futures.append(executor.submit(
+                    snapshot_download,
+                    repo_id=model,
+                    local_dir=model_path,
+                    force_download=True
+                ))
+            
+            # 모든 작업 완료 대기
+            for future in futures:
+                future.result()
         
-        # 작업 디렉토리 변경
-        os.chdir(base_dir)
-        logging.info(f"Working directory changed to: {os.getcwd()}")
+        # 파일 존재 확인
+        for file_name, _ in required_files.items():
+            file_path = os.path.join(xcodec_dir, file_name)
+            if not os.path.exists(file_path):
+                raise FileNotFoundError(f"Required file still missing after download: {file_path}")
+            else:
+                file_size = os.path.getsize(file_path)
+                logging.info(f"Verified {file_name}: {file_size} bytes")
         
-        # 필요한 파일 존재 확인
-        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")
-        ]
+        logging.info("System initialization completed successfully")
         
-        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"Directory error: {e}")
         raise