Update app.py

app.py

@@ -202,42 +202,55 @@ def install_flash_attn():
         return False
 
 
-
-
 def download_missing_files():
     try:
         from huggingface_hub import hf_hub_download, snapshot_download
         
         # xcodec_mini_infer 파일들 직접 다운로드
+        repo_id = "hf-internal-testing/xcodec_mini_infer"  # 저장소 경로 수정
         files_to_download = {
-            "config.json": "config.json",
-            "vocal_decoder.pth": "vocal_decoder.pth",
-            "inst_decoder.pth": "inst_decoder.pth"
+            "config.json": "config/config.json",
+            "vocal_decoder.pth": "checkpoints/vocal_decoder.pth",
+            "inst_decoder.pth": "checkpoints/inst_decoder.pth"
         }
         
         xcodec_dir = "./xcodec_mini_infer"
         os.makedirs(xcodec_dir, exist_ok=True)
+        os.makedirs(os.path.join(xcodec_dir, "checkpoints"), exist_ok=True)
         
-        for file_name in files_to_download.keys():
+        for target_name, source_path in files_to_download.items():
             try:
                 downloaded_path = hf_hub_download(
-                    repo_id="m-a-p/xcodec_mini_infer",
-                    filename=file_name,
-                    local_dir=xcodec_dir,
+                    repo_id=repo_id,
+                    filename=source_path,
+                    cache_dir="./models/cache",
                     force_download=True,
                     local_files_only=False
                 )
-                logging.info(f"Downloaded {file_name} to: {downloaded_path}")
+                logging.info(f"Downloaded {source_path} to: {downloaded_path}")
                 
-                # 파일 위치 확인 및 필요한 경우 복사
-                target_path = os.path.join(xcodec_dir, file_name)
-                if downloaded_path != target_path:
-                    shutil.copy2(downloaded_path, target_path)
-                    logging.info(f"Copied {file_name} to correct location: {target_path}")
+                # 파일을 올바른 위치로 복사
+                target_path = os.path.join(xcodec_dir, target_name)
+                shutil.copy2(downloaded_path, target_path)
+                logging.info(f"Copied to: {target_path}")
                 
             except Exception as e:
-                logging.error(f"Error downloading {file_name}: {e}")
-                raise
+                logging.error(f"Error downloading {source_path}: {e}")
+                # 대체 저장소 시도
+                try:
+                    alt_repo_id = "facebook/musicgen-small"
+                    downloaded_path = hf_hub_download(
+                        repo_id=alt_repo_id,
+                        filename=source_path,
+                        cache_dir="./models/cache",
+                        force_download=True
+                    )
+                    target_path = os.path.join(xcodec_dir, target_name)
+                    shutil.copy2(downloaded_path, target_path)
+                    logging.info(f"Downloaded from alternate source to: {target_path}")
+                except Exception as alt_e:
+                    logging.error(f"Error with alternate download: {alt_e}")
+                    raise
         
         # YuE 모델들 다운로드
         models = [
@@ -252,23 +265,24 @@ def download_missing_files():
             model_path = snapshot_download(
                 repo_id=model,
                 local_dir=f"./models/{model_name}",
+                cache_dir="./models/cache",
                 resume_download=True,
-                force_download=True,
-                local_files_only=False
+                force_download=True
             )
             logging.info(f"Downloaded {model_name} to: {model_path}")
+        
+        # 파일 존재 및 크기 확인
+        for target_name in files_to_download.keys():
+            file_path = os.path.join(xcodec_dir, target_name)
+            if not os.path.exists(file_path):
+                raise FileNotFoundError(f"Failed to download {target_name}")
+            file_size = os.path.getsize(file_path)
+            if file_size == 0:
+                raise FileNotFoundError(f"Downloaded file is empty: {target_name}")
+            logging.info(f"Verified {target_name}: {file_size} bytes")
             
         logging.info("All required models downloaded successfully")
         
-        # 파일 존재 확인
-        for file_name in files_to_download.keys():
-            file_path = os.path.join(xcodec_dir, file_name)
-            if not os.path.exists(file_path):
-                raise FileNotFoundError(f"Failed to download {file_name}")
-            else:
-                file_size = os.path.getsize(file_path)
-                logging.info(f"Verified {file_name}: {file_size} bytes")
-                
     except Exception as e:
         logging.error(f"Error downloading models: {e}")
         raise
@@ -319,25 +333,21 @@ def initialize_system():
         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)
         
         # inference 디렉토리로 이동
         os.chdir("./inference")
         logging.info(f"Working directory changed to: {os.getcwd()}")
         
         # 환경 변수 설정
-        os.environ["HF_HOME"] = os.path.abspath("./models/cache")
-        os.environ["TRANSFORMERS_CACHE"] = os.path.abspath("./models/cache")
+        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
         
         # 모델 파일 체크 및 다운로드
         check_model_files()
         
-        # 파일 권한 확인
-        xcodec_dir = "./xcodec_mini_infer"
-        for file_name in ["config.json", "vocal_decoder.pth", "inst_decoder.pth"]:
-            file_path = os.path.join(xcodec_dir, file_name)
-            if os.path.exists(file_path):
-                os.chmod(file_path, 0o644)  # 읽기 권한 확인
-        
         logging.info("System initialization completed successfully")
         
     except Exception as e: