Update app.py

app.py

@@ -200,37 +200,89 @@ def install_flash_attn():
         return False
 
 
+# 전역 변수로 경로 설정
+APP_DIR = os.path.abspath(os.path.dirname(__file__))
+INFERENCE_DIR = os.path.join(APP_DIR, "inference")
+INFER_SCRIPT = os.path.join(INFERENCE_DIR, "infer.py")
+
 def initialize_system():
     optimize_gpu_settings()
     
-    with ThreadPoolExecutor(max_workers=4) as executor:
-        futures = []
+    try:
+        # 디렉토리 구조 생성
+        os.makedirs(INFERENCE_DIR, exist_ok=True)
+        os.makedirs(os.path.join(INFERENCE_DIR, "models"), exist_ok=True)
+        os.makedirs(os.path.join(INFERENCE_DIR, "xcodec_mini_infer"), exist_ok=True)
         
-        futures.append(executor.submit(install_flash_attn))
+        from huggingface_hub import snapshot_download, hf_hub_download
         
-        from huggingface_hub import snapshot_download
+        # infer.py 파일 다운로드
+        try:
+            infer_script_download = hf_hub_download(
+                repo_id="m-a-p/xcodec_mini_infer",
+                filename="infer.py",
+                local_dir=INFERENCE_DIR,
+                force_download=True
+            )
+            if not os.path.exists(INFER_SCRIPT):
+                shutil.copy2(infer_script_download, INFER_SCRIPT)
+            logging.info(f"infer.py available at: {INFER_SCRIPT}")
+        except Exception as e:
+            logging.error(f"Failed to download infer.py: {e}")
+            raise
         
-        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_path = os.path.join(INFERENCE_DIR, "xcodec_mini_infer")
+        snapshot_download(
             repo_id="m-a-p/xcodec_mini_infer",
-            local_dir="./inference/xcodec_mini_infer",
-            resume_download=True
-        ))
+            local_dir=xcodec_path,
+            force_download=True
+        )
         
-        for future in futures:
-            future.result()
-
-    try:
-        os.chdir("./inference")
+        # 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(INFERENCE_DIR, "models", model_name)
+            snapshot_download(
+                repo_id=model,
+                local_dir=model_path,
+                force_download=True
+            )
+        
+        # 작업 디렉토리 변경
+        os.chdir(INFERENCE_DIR)
         logging.info(f"Working directory changed to: {os.getcwd()}")
-    except FileNotFoundError as e:
-        logging.error(f"Directory error: {e}")
+        
+        # 파일 존재 확인
+        required_files = [
+            INFER_SCRIPT,
+            os.path.join(xcodec_path, "config.json"),
+            os.path.join(xcodec_path, "vocal_decoder.pth"),
+            os.path.join(xcodec_path, "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 {os.path.basename(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)
 def get_cached_file_path(content_hash, prefix):
     return create_temp_file(content_hash, prefix)
@@ -272,15 +324,21 @@ def get_audio_duration(file_path):
         logging.error(f"Failed to get audio duration: {e}")
         return None
 
+
 def infer(genre_txt_content, lyrics_txt_content, num_segments, max_new_tokens):
     genre_txt_path = None
     lyrics_txt_path = None
     
     try:
+        if not os.path.exists(INFER_SCRIPT):
+            raise FileNotFoundError(f"infer.py not found at: {INFER_SCRIPT}")
+        
         model_path, config, params = optimize_model_selection(lyrics_txt_content, genre_txt_content)
         logging.info(f"Selected model: {model_path}")
         logging.info(f"Lyrics analysis: {params}")
         
+
+        
         has_chorus = params['sections']['chorus'] > 0
         estimated_duration = params.get('estimated_duration', 90)
 
@@ -306,20 +364,24 @@ 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",
+            sys.executable,
+            INFER_SCRIPT,
             "--stage1_model", model_path,
             "--stage2_model", "m-a-p/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"),
             "--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():