Update app.py

app.py

@@ -3,12 +3,12 @@ import uuid
 import os
 from huggingface_hub import snapshot_download
 from flux_train import build_job
-
 import sys
-sys.path.append("/app/ai-toolkit")  # Tell Python to look here
 
-from toolkit.job import run_job
+# Add ai-toolkit to sys.path for toolkit imports
+sys.path.append("/app/ai-toolkit")
 
+from toolkit.job import run_job
 
 app = FastAPI()
 
@@ -20,10 +20,11 @@ HF_TOKEN = os.environ.get("HF_TOKEN", "")
 
 status = {"running": False, "last_job": None, "error": None}
 
-def run_lora_training():
+def run_lora_training(push_to_hub: bool = False):
     try:
         status.update({"running": True, "error": None})
         local_dir = f"/tmp/{LORA_NAME}-{uuid.uuid4()}"
+        
         snapshot_download(
             repo_id=REPO_ID,
             repo_type="dataset",
@@ -31,9 +32,11 @@ def run_lora_training():
             local_dir=local_dir,
             local_dir_use_symlinks=False
         )
+        
         training_path = os.path.join(local_dir, FOLDER_IN_REPO)
-        job = build_job(CONCEPT_SENTENCE, training_path, LORA_NAME)
+        job = build_job(CONCEPT_SENTENCE, training_path, LORA_NAME, push_to_hub=push_to_hub)
         run_job(job)
+        
         status.update({"running": False, "last_job": job})
     except Exception as e:
         status.update({"running": False, "error": str(e)})
@@ -46,9 +49,14 @@ def root():
 def get_status():
     return status
 
+from pydantic import BaseModel
+
+class TrainRequest(BaseModel):
+    push_to_hub: bool = False
+
 @app.post("/train")
-def start_training(background_tasks: BackgroundTasks):
+def start_training(background_tasks: BackgroundTasks, request: TrainRequest):
     if status["running"]:
         return {"message": "A training job is already running."}
-    background_tasks.add_task(run_lora_training)
-    return {"message": "Training started in background."}
+    background_tasks.add_task(run_lora_training, push_to_hub=request.push_to_hub)
+    return {"message": "Training started in background.", "push_to_hub": request.push_to_hub}