Update app.py

app.py

@@ -202,64 +202,42 @@ def install_flash_attn():
         return False
 
 
-def check_model_files():
-    base_dir = os.getcwd()
-    xcodec_dir = os.path.join(base_dir, "xcodec_mini_infer")
-    
-    # 필요한 파일 목록
-    required_files = {
-        "config.json": "config.json",
-        "vocal_decoder.pth": "vocal_decoder.pth",
-        "inst_decoder.pth": "inst_decoder.pth"
-    }
-    
-    # 파일 존재 여부 확인
-    missing = False
-    for file_name in required_files.keys():
-        file_path = os.path.join(xcodec_dir, file_name)
-        if not os.path.exists(file_path):
-            missing = True
-            logging.warning(f"Missing file: {file_path}")
-    
-    if missing:
-        logging.info("Downloading missing files...")
-        download_missing_files()
-        
-        # 다운로드 후 파일 재확인
-        for file_name, src_name in required_files.items():
-            src_path = os.path.join(xcodec_dir, src_name)
-            dst_path = os.path.join(xcodec_dir, file_name)
-            
-            if not os.path.exists(dst_path) and os.path.exists(src_path):
-                shutil.copy2(src_path, dst_path)
-            
-            if not os.path.exists(dst_path):
-                raise FileNotFoundError(f"Failed to download or locate required file: {file_name}")
+
 
 def download_missing_files():
     try:
-        from huggingface_hub import snapshot_download
+        from huggingface_hub import hf_hub_download, snapshot_download
         
-        # xcodec_mini_infer 모델 다운로드
-        xcodec_path = snapshot_download(
-            repo_id="m-a-p/xcodec_mini_infer",
-            local_dir="./xcodec_mini_infer",
-            resume_download=True,
-            force_download=True,
-            local_files_only=False
-        )
-        logging.info(f"Downloaded xcodec_mini_infer to: {xcodec_path}")
+        # xcodec_mini_infer 파일들 직접 다운로드
+        files_to_download = {
+            "config.json": "config.json",
+            "vocal_decoder.pth": "vocal_decoder.pth",
+            "inst_decoder.pth": "inst_decoder.pth"
+        }
         
-        # 파일 존재 확인 및 복사
         xcodec_dir = "./xcodec_mini_infer"
-        for root, _, files in os.walk(xcodec_dir):
-            for file in files:
-                if file in ["config.json", "vocal_decoder.pth", "inst_decoder.pth"]:
-                    src_path = os.path.join(root, file)
-                    dst_path = os.path.join(xcodec_dir, file)
-                    if src_path != dst_path:
-                        shutil.copy2(src_path, dst_path)
-                        logging.info(f"Copied {file} to {dst_path}")
+        os.makedirs(xcodec_dir, exist_ok=True)
+        
+        for file_name in files_to_download.keys():
+            try:
+                downloaded_path = hf_hub_download(
+                    repo_id="m-a-p/xcodec_mini_infer",
+                    filename=file_name,
+                    local_dir=xcodec_dir,
+                    force_download=True,
+                    local_files_only=False
+                )
+                logging.info(f"Downloaded {file_name} 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}")
+                
+            except Exception as e:
+                logging.error(f"Error downloading {file_name}: {e}")
+                raise
         
         # YuE 모델들 다운로드
         models = [
@@ -281,10 +259,54 @@ def download_missing_files():
             logging.info(f"Downloaded {model_name} to: {model_path}")
             
         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
 
+def check_model_files():
+    base_dir = os.getcwd()
+    xcodec_dir = os.path.join(base_dir, "xcodec_mini_infer")
+    
+    # 필요한 파일 목록
+    required_files = {
+        "config.json": "config.json",
+        "vocal_decoder.pth": "vocal_decoder.pth",
+        "inst_decoder.pth": "inst_decoder.pth"
+    }
+    
+    # 파일 존재 여부 확인
+    missing = False
+    for file_name in required_files.keys():
+        file_path = os.path.join(xcodec_dir, file_name)
+        if not os.path.exists(file_path):
+            missing = True
+            logging.warning(f"Missing file: {file_path}")
+    
+    if missing:
+        logging.info("Downloading missing files...")
+        download_missing_files()
+        
+        # 다운로드 후 파일 재확인
+        for file_name in required_files.keys():
+            file_path = os.path.join(xcodec_dir, file_name)
+            if not os.path.exists(file_path):
+                raise FileNotFoundError(f"Failed to download or locate required file: {file_name}")
+            else:
+                file_size = os.path.getsize(file_path)
+                if file_size == 0:
+                    raise FileNotFoundError(f"Downloaded file is empty: {file_name}")
+                logging.info(f"Verified {file_name}: {file_size} bytes")
+
 def initialize_system():
     optimize_gpu_settings()