Create app.py

app.py

@@ -0,0 +1,82 @@
+import gradio as gr
+import torch
+import os
+import shutil
+import requests
+import subprocess
+from subprocess import getoutput
+from huggingface_hub import snapshot_download, HfApi, create_repo
+
+api = HfApi()
+
+hf_token = os.environ.get("HF_TOKEN_WITH_WRITE_PERMISSION")
+
+def train_dreambooth_blora_sdxl(instance_data_dir, b_lora_trained_folder, instance_prompt, max_train_steps, checkpoint_steps):
+    
+    script_filename = "train_dreambooth_b-lora_sdxl.py"  # Assuming it's in the same folder
+
+    command = [
+        "accelerate",
+        "launch",
+        script_filename,  # Use the local script
+        "--pretrained_model_name_or_path=stabilityai/stable-diffusion-xl-base-1.0",
+        f"--instance_data_dir={instance_data_dir}",
+        f"--output_dir={b_lora_trained_folder}",
+        f"--instance_prompt={instance_prompt}",
+        "--resolution=1024",
+        "--rank=64",
+        "--train_batch_size=1",
+        "--learning_rate=5e-5",
+        "--lr_scheduler=constant",
+        "--lr_warmup_steps=0",
+        f"--max_train_steps={max_train_steps}",
+        f"--checkpointing_steps={checkpoint_steps}",
+        "--seed=0",
+        "--gradient_checkpointing",
+        "--use_8bit_adam",
+        "--mixed_precision=fp16",
+        "--push_to_hub",
+        f"--hub_token={hf_token}"
+    ]
+
+    try:
+        subprocess.run(command, check=True)
+        print("Training is finished!")
+    
+    except subprocess.CalledProcessError as e:
+        print(f"An error occurred: {e}")
+
+def main(image_path, b_lora_trained_folder, instance_prompt):
+    
+    local_dir = "image_to_train"
+    # Check if the directory exists and create it if necessary
+    if not os.path.exists(local_dir):
+        os.makedirs(local_dir)
+
+    shutil.copy(image_path, local_dir)
+    print(f"source image has been copied in {local_dir} directory")
+    
+    max_train_steps = 1000
+    checkpoint_steps = 500
+    
+    train_dreambooth_blora_sdxl(local_dir, b_lora_trained_folder, instance_prompt, max_train_steps, checkpoint_steps)
+    
+    your_username = api.whoami(token=hf_token)["name"]
+    
+    return f"Done, your trained model has been stored in your models library: {your_username}/{b_lora_trained_folder}"
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        image = gr.Image(sources=[upload], type="filepath")
+        b_lora_name = gr.Textbox(label="b_lora_name", placeholder="b_lora_trained_folder")
+        instance_prompt = gr.Textbox(label="instance prompt", placeholder="[v42]")
+        train_btn = gr.Button("Train B-LoRa")
+        status = gr.Textbox(label="status")
+    
+    train_btn.click(
+        fn = main,
+        inputs = [image, b_lora_name, instance_prompt],
+        outputs = [status]
+    )
+
+demo.launch(debug=True)