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") is_shared_ui = True if "fffiloni/train-dreambooth-lora-sdxl" in os.environ['SPACE_ID'] else False is_gpu_associated = torch.cuda.is_available() if is_gpu_associated: gpu_info = getoutput('nvidia-smi') if("A10G" in gpu_info): which_gpu = "A10G" elif("T4" in gpu_info): which_gpu = "T4" else: which_gpu = "CPU" 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}'", f"--validation_prompt=a teddy bear in {instance_prompt} style", "--num_validation_images=1", "--validation_epochs=300", "--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): if is_shared_ui: raise gr.Error("This Space only works in duplicated instances") if not is_gpu_associated: raise gr.Error("Please associate a T4 or A10G GPU for this Space") if image_path == None: raise gr.Error("You forgot to specify an image reference") if b_lora_trained_folder == "": raise gr.Error("You forgot to specify a name for you model") if instance_prompt == "": raise gr.Error("You forgot to specify an 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}" css = """ #col-container {max-width: 780px; margin-left: auto; margin-right: auto;} div#warning-ready { background-color: #ecfdf5; padding: 0 10px 5px; margin: 20px 0; } div#warning-ready > .gr-prose > h2, div#warning-ready > .gr-prose > p { color: #057857!important; } div#warning-duplicate { background-color: #ebf5ff; padding: 0 10px 5px; margin: 20px 0; } div#warning-duplicate > .gr-prose > h2, div#warning-duplicate > .gr-prose > p { color: #0f4592!important; } div#warning-duplicate strong { color: #0f4592; } p.actions { display: flex; align-items: center; margin: 20px 0; } div#warning-duplicate .actions a { display: inline-block; margin-right: 10px; } div#warning-setgpu { background-color: #fff4eb; padding: 0 10px 5px; margin: 20px 0; } div#warning-setgpu > .gr-prose > h2, div#warning-setgpu > .gr-prose > p { color: #92220f!important; } div#warning-setgpu a, div#warning-setgpu b { color: #91230f; } div#warning-setgpu p.actions > a { display: inline-block; background: #1f1f23; border-radius: 40px; padding: 6px 24px; color: antiquewhite; text-decoration: none; font-weight: 600; font-size: 1.2em; } """ with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): if is_shared_ui: top_description = gr.HTML(f'''

Attention: this Space need to be duplicated to work

To make it work, duplicate the Space and run it on your own profile using a private GPU (T4-small or A10G-small).
A T4 costs US$0.60/h, so it should cost < US$1 to train most models.

Duplicate this Space to start training your own B-LoRa model

''', elem_id="warning-duplicate") else: if(is_gpu_associated): top_description = gr.HTML(f'''

You have successfully associated a {which_gpu} GPU to the B-LoRa Training Space 🎉

You can now train your model! You will be billed by the minute from when you activated the GPU until when it is turned off.

''', elem_id="warning-ready") else: top_description = gr.HTML(f'''

You have successfully duplicated the B-LoRa Training Space 🎉

There's only one step left before you can train your model: attribute a T4-small or A10G-small GPU to it (via the Settings tab) and run the training below. You will be billed by the minute from when you activate the GPU until when it is turned off.

🔥   Set recommended GPU

''', elem_id="warning-setgpu") gr.Markdown("# B-LoRa Training UI 💭") with gr.Row(): image = gr.Image(label="Image Reference", sources=["upload"], type="filepath") with gr.Column(): 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") status = gr.Textbox(label="status") train_btn.click( fn = main, inputs = [image, b_lora_name, instance_prompt], outputs = [status] ) demo.launch(debug=True)