Update app.py

app.py

@@ -8,7 +8,6 @@ if is_spaces:
     import spaces
     from huggingface_hub import whoami
     
-    
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 import sys
 
@@ -31,13 +30,13 @@ from slugify import slugify
 from transformers import AutoProcessor, AutoModelForCausalLM
 
 if not is_spaces:
+    sys.path.insert(0, "ai-toolkit")
     from toolkit.job import get_job
     gr.OAuthProfile = None
     gr.OAuthToken = None
     
 MAX_IMAGES = 150
 
-
 def load_captioning(uploaded_images, concept_sentence):
     gr.Info("Images uploaded!")
     updates = []
@@ -136,8 +135,6 @@ if is_spaces:
     run_captioning = spaces.GPU()(run_captioning)
 
 def start_training(
-    profile: Union[gr.OAuthProfile, None],
-    oauth_token: Union[gr.OAuthToken, None],
     lora_name,
     concept_sentence,
     steps,
@@ -147,6 +144,8 @@ def start_training(
     sample_1,
     sample_2,
     sample_3,
+    profile: Union[gr.OAuthProfile, None],
+    oauth_token: Union[gr.OAuthToken, None],
 ):
     if not lora_name:
         raise gr.Error("You forgot to insert your LoRA name! This name has to be unique.")
@@ -154,7 +153,7 @@ def start_training(
     slugged_lora_name = slugify(lora_name)
 
     # Load the default config
-    with open("train_lora_flux_24gb.yaml" if is_spaces else "config/examples/train_lora_flux_24gb.yaml", "r") as f:
+    with open("train_lora_flux_24gb.yaml" if is_spaces else "ai-toolkit/config/examples/train_lora_flux_24gb.yaml", "r") as f:
         config = yaml.safe_load(f)
 
     # Update the config with user inputs
@@ -347,15 +346,15 @@ with gr.Blocks(theme=theme, css=css) as demo:
         gr.Markdown(
             f"""To use FLUX LoRA Ease locally with this UI, you can clone this repository (yes, HF Spaces are git repos!). You'll need ~23GB of VRAM
         ```bash
-        git clone https://huggingface.co/spaces/flux-train/flux-lora-trainer
-        cd flux-lora-trainer
+        git clone https://huggingface.co/spaces/autotrain-projects/flux-lora-ease
+        cd flux-lora-ease
         ## Optional, start a venv environment ## 
         python3 -m venv venv
         source venv/bin/activate
         # .\venv\Scripts\activate on windows
         # install torch first
         ## End of optional ##
-        pip install requirements_local.txt
+        pip install -r requirements_local.txt
         ```
         
         Then you can install ai-toolkit
@@ -373,7 +372,7 @@ with gr.Blocks(theme=theme, css=css) as demo:
         huggingface-cli login
         ```
         
-        Now you can run FLUX LoRA Ease locally with a UI by doing a simple 
+        Finally, you can run FLUX LoRA Ease locally with a UI by doing a simple 
         ```py
         python app.py
         ```