Update app.py

app.py

@@ -4,6 +4,20 @@ from typing import Union
 from huggingface_hub import whoami, HfApi
 from fastapi import FastAPI
 from starlette.middleware.sessions import SessionMiddleware
+import sys
+
+# ai-toolkit이 없으면 설치
+if not os.path.exists("ai-toolkit"):
+    subprocess.run("git clone https://github.com/ostris/ai-toolkit.git", shell=True)
+    subprocess.run("cd ai-toolkit && git submodule update --init --recursive", shell=True)
+
+# ai-toolkit 경로 추가
+toolkit_path = os.path.join(os.getcwd(), "ai-toolkit")
+sys.path.append(toolkit_path)
+
+# 필요한 패키지 설치
+subprocess.run("pip install -r ai-toolkit/requirements.txt", shell=True)
+
 
 is_spaces = True if os.environ.get("SPACE_ID") else False
     
@@ -201,6 +215,14 @@ def start_training(
     print("Started training")
     slugged_lora_name = slugify(lora_name)
 
+    try:
+        from toolkit.job import get_job
+    except ImportError:
+        raise gr.Error("Failed to import toolkit. Please check if ai-toolkit is properly installed.")
+        
+    print("Started training")
+    slugged_lora_name = slugify(lora_name)
+
     # Load the default config
     with open("train_lora_flux_24gb.yaml", "r") as f:
         config = yaml.safe_load(f)