Update app.py

app.py

@@ -62,7 +62,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):
+def main(image_path, b_lora_trained_folder, instance_prompt, training_type):
 
     if is_shared_ui:
         raise gr.Error("This Space only works in duplicated instances")
@@ -86,9 +86,13 @@ def main(image_path, b_lora_trained_folder, instance_prompt):
 
     shutil.copy(image_path, local_dir)
     print(f"source image has been copied in {local_dir} directory")
-    
-    max_train_steps = 2000
-    checkpoint_steps = 1000
+
+    if training_type == "style":
+        max_train_steps = 1000
+        checkpoint_steps = 500
+    elif training_type == "concept" :
+        max_train_steps = 2000
+        checkpoint_steps = 1000
     
     train_dreambooth_blora_sdxl(local_dir, b_lora_trained_folder, instance_prompt, max_train_steps, checkpoint_steps)
     
@@ -196,6 +200,7 @@ with gr.Blocks(css=css) as demo:
             image = gr.Image(label="Image Reference", sources=["upload"], type="filepath")
 
             with gr.Column():
+                training_type = gr.Radio(label="Training type", choices=["style", "concept"], value="style")
                 b_lora_name = gr.Textbox(label="Name your B-LoRa model", placeholder="b_lora_trained_folder")
                 instance_prompt = gr.Textbox(label="Create instance prompt", placeholder="[v42]")
         train_btn = gr.Button("Train B-LoRa")