Update app.py

app.py

@@ -203,27 +203,35 @@ def install_flash_attn():
 
 
 def check_model_files():
+    base_dir = os.getcwd()
     required_files = [
-        "./xcodec_mini_infer/config.json",
-        "./xcodec_mini_infer/vocal_decoder.pth",
-        "./xcodec_mini_infer/inst_decoder.pth"
+        os.path.join(base_dir, "xcodec_mini_infer", "config.json"),
+        os.path.join(base_dir, "xcodec_mini_infer", "vocal_decoder.pth"),
+        os.path.join(base_dir, "xcodec_mini_infer", "inst_decoder.pth")
     ]
     
     missing_files = [f for f in required_files if not os.path.exists(f)]
     if missing_files:
         logging.warning(f"Missing required files: {missing_files}")
         download_missing_files()
+        
+    # 파일 존재 확인
+    for file_path in required_files:
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"Required file still missing after download: {file_path}")
 
 def download_missing_files():
     try:
         from huggingface_hub import snapshot_download
         
         # xcodec_mini_infer 모델 다운로드
-        snapshot_download(
+        xcodec_path = snapshot_download(
             repo_id="m-a-p/xcodec_mini_infer",
             local_dir="./xcodec_mini_infer",
-            resume_download=True
+            resume_download=True,
+            force_download=True
         )
+        logging.info(f"Downloaded xcodec_mini_infer to: {xcodec_path}")
         
         # YuE 모델들 다운로드
         models = [
@@ -235,70 +243,57 @@ def download_missing_files():
         
         for model in models:
             model_name = model.split('/')[-1]
-            snapshot_download(
+            model_path = snapshot_download(
                 repo_id=model,
                 local_dir=f"./models/{model_name}",
-                resume_download=True
+                resume_download=True,
+                force_download=True
             )
+            logging.info(f"Downloaded {model_name} to: {model_path}")
             
         logging.info("All required models downloaded successfully")
     except Exception as e:
         logging.error(f"Error downloading models: {e}")
         raise
+
 def initialize_system():
     optimize_gpu_settings()
     
     try:
+        # 현재 작업 디렉토리 저장
+        original_dir = os.getcwd()
+        
         # 기본 디렉토리 구조 생성
         os.makedirs("./inference", exist_ok=True)
         os.makedirs("./inference/models", exist_ok=True)
         os.makedirs("./inference/models/cache", exist_ok=True)
         os.makedirs("./inference/xcodec_mini_infer", 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")
+        
         # 모델 파일 체크 및 다운로드
         check_model_files()
         
-        with ThreadPoolExecutor(max_workers=4) as executor:
-            futures = []
-            
-            # Flash Attention 설치
-            futures.append(executor.submit(install_flash_attn))
-            
-            from huggingface_hub import snapshot_download
-            
-            # 필요한 모든 모델 다운로드
-            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:
-                model_name = model.split('/')[-1]
-                model_dir = f"./models/{model_name}"
-                
-                if not os.path.exists(model_dir):
-                    futures.append(executor.submit(
-                        snapshot_download,
-                        repo_id=model,
-                        local_dir=model_dir,
-                        resume_download=True
-                    ))
-            
-            # 모든 다운로드 완료 대기
-            for future in futures:
-                future.result()
+        # GPU 설정 확인
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            device = torch.device("cuda")
+            logging.info(f"Using GPU device: {torch.cuda.get_device_name(0)}")
+        else:
+            device = torch.device("cpu")
+            logging.warning("GPU not available, using CPU")
         
         logging.info("System initialization completed successfully")
         
     except Exception as e:
         logging.error(f"Initialization error: {e}")
+        os.chdir(original_dir)  # 오류 발생 시 원래 디렉토리로 복귀
         raise
 
 @lru_cache(maxsize=100)
@@ -381,22 +376,25 @@ 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 = [
             "python", "infer.py",
             "--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,
+            "--genre_txt", os.path.abspath(genre_txt_path),
+            "--lyrics_txt", os.path.abspath(lyrics_txt_path),
             "--run_n_segments", str(actual_num_segments),
             "--stage2_batch_size", "16",
-            "--output_dir", output_dir,
+            "--output_dir", os.path.abspath(output_dir),
             "--cuda_idx", "0",
             "--max_new_tokens", str(actual_max_tokens),
             "--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"
+            "--basic_model_config", os.path.abspath("./xcodec_mini_infer/config.json"),
+            "--vocal_decoder_path", os.path.abspath("./xcodec_mini_infer/vocal_decoder.pth"),
+            "--inst_decoder_path", os.path.abspath("./xcodec_mini_infer/inst_decoder.pth")
         ]
+        
+
 
         env = os.environ.copy()
         if torch.cuda.is_available():