Spaces:
Running
Running
File size: 2,836 Bytes
6c4222e 00ac9d6 6c4222e 45803c3 6c4222e 1025a0e 6c4222e ffd05ee 6c4222e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
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 = 10
checkpoint_steps = 10
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;}
"""
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) |