Update app.py

app.py

@@ -68,7 +68,7 @@ def train_dreambooth_blora_sdxl(instance_data_dir, b_lora_trained_folder, instan
     except subprocess.CalledProcessError as e:
         print(f"An error occurred: {e}")
 
-def main(image_path, b_lora_trained_folder, instance_prompt, training_type):
+def main(image_path, b_lora_trained_folder, instance_prompt, training_type, training_steps):
 
     if is_shared_ui:
         raise gr.Error("This Space only works in duplicated instances")
@@ -94,11 +94,11 @@ def main(image_path, b_lora_trained_folder, instance_prompt, training_type):
     print(f"source image has been copied in {local_dir} directory")
 
     if training_type == "style":
-        max_train_steps = 1000
         checkpoint_steps = 500
     elif training_type == "concept" :
-        max_train_steps = 2000
         checkpoint_steps = 1000
+
+    max_train_steps = training_steps
     
     train_dreambooth_blora_sdxl(local_dir, b_lora_trained_folder, instance_prompt, max_train_steps, checkpoint_steps)
     
@@ -222,7 +222,7 @@ with gr.Blocks(css=css) as demo:
     
     train_btn.click(
         fn = main,
-        inputs = [image, b_lora_name, instance_prompt, training_type],
+        inputs = [image, b_lora_name, instance_prompt, training_type, training_steps],
         outputs = [status]
     )