Update app.py

app.py

@@ -1,12 +1,41 @@
 import os
+import subprocess
+import sys
 
-# Set this environment variable to disable torch.compiler features
+# 필요한 환경 변수 설정
 os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
 os.environ["TRANSFORMERS_COMPILER_DISABLED"] = "1"
 
+# 필요한 모듈 설치 함수
+def install_required_packages():
+    required_packages = [
+        "warmup_scheduler",
+        "cosine_annealing_warmup_restarts"
+    ]
+    
+    for package in required_packages:
+        try:
+            __import__(package)
+            print(f"{package} is already installed")
+        except ImportError:
+            print(f"Installing {package}...")
+            try:
+                subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+                print(f"{package} installed successfully")
+            except subprocess.CalledProcessError:
+                # 일부 패키지는 PyPI에 없을 수 있으므로 GitHub에서 직접 설치
+                if package == "warmup_scheduler":
+                    subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/ildoonet/pytorch-gradual-warmup-lr.git"])
+                    print(f"{package} installed from GitHub successfully")
+                else:
+                    print(f"Failed to install {package}")
+
+# 필요한 모듈 설치
+install_required_packages()
+
+# 그 후 나머지 imports
 import yaml
 import torch
-import sys
 sys.path.append(os.path.abspath('./'))
 from inference.utils import *
 from train import WurstCoreB
@@ -18,7 +47,6 @@ import argparse
 import gradio as gr
 import spaces
 from huggingface_hub import hf_hub_url
-import subprocess
 from huggingface_hub import hf_hub_download
 from transformers import pipeline