Update app.py

app.py

@@ -30,7 +30,7 @@ def change_training_setup(training_type):
     elif training_type == "concept" :
         return 2000, 1000
 
-def train_dreambooth_blora_sdxl(instance_data_dir, b_lora_trained_folder, instance_prompt, max_train_steps, checkpoint_steps):
+def train_dreambooth_blora_sdxl(instance_data_dir, b_lora_trained_folder, instance_prompt, class_prompt, max_train_steps, checkpoint_steps):
     
     script_filename = "train_dreambooth_b-lora_sdxl.py"  # Assuming it's in the same folder
 
@@ -42,6 +42,7 @@ def train_dreambooth_blora_sdxl(instance_data_dir, b_lora_trained_folder, instan
         f"--instance_data_dir={instance_data_dir}",
         f"--output_dir={b_lora_trained_folder}",
         f"--instance_prompt='{instance_prompt}'",
+        f"--class_prompt={class_prompt}",
         #f"--validation_prompt=a teddy bear in {instance_prompt} style",
         "--num_validation_images=1",
         "--validation_epochs=500",
@@ -68,7 +69,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, training_steps):
+def main(image_path, b_lora_trained_folder, instance_prompt, class_prompt, training_type, training_steps):
 
     if is_shared_ui:
         raise gr.Error("This Space only works in duplicated instances")
@@ -100,7 +101,7 @@ def main(image_path, b_lora_trained_folder, instance_prompt, training_type, trai
 
     max_train_steps = training_steps
     
-    train_dreambooth_blora_sdxl(local_dir, b_lora_trained_folder, instance_prompt, max_train_steps, checkpoint_steps)
+    train_dreambooth_blora_sdxl(local_dir, b_lora_trained_folder, instance_prompt, class_prompt, max_train_steps, checkpoint_steps)
     
     your_username = api.whoami(token=hf_token)["name"]
     
@@ -208,7 +209,9 @@ with gr.Blocks(css=css) as demo:
             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="A [v42]")
+                with gr.Row():
+                    instance_prompt = gr.Textbox(label="Create instance prompt", placeholder="A [v42] <class_prompt>")
+                    class_prompt = gr.Textbox(label="Specify class prompt", placeholder="style | person | dog ")
                 training_steps = gr.Number(label="Training steps", value=1000, interactive=False)
                 checkpoint_step = gr.Number(label="checkpoint step", visible=False, value=500)
         train_btn = gr.Button("Train B-LoRa")
@@ -222,7 +225,7 @@ with gr.Blocks(css=css) as demo:
     
     train_btn.click(
         fn = main,
-        inputs = [image, b_lora_name, instance_prompt, training_type, training_steps],
+        inputs = [image, b_lora_name, instance_prompt, class_prompt, training_type, training_steps],
         outputs = [status]
     )