#!/usr/bin/env python
"""
Demo showcasing parameter-efficient fine-tuning of Stable Dissfusion via Dreambooth leveraging 🤗 PEFT (https://github.com/huggingface/peft)
The code in this repo is partly adapted from the following repositories:
https://huggingface.co/spaces/hysts/LoRA-SD-training
https://huggingface.co/spaces/multimodalart/dreambooth-training
"""
from __future__ import annotations
import os
import pathlib
import gradio as gr
import torch
from typing import List
from inference import InferencePipeline
from trainer import Trainer
from uploader import upload
TITLE = "# RealFill Training and Inference Demo 🎨"
DESCRIPTION = "Demo showcasing parameter-efficient fine-tuning of Stable Diffusion Inpainting via RealFill leveraging 🤗 PEFT (https://github.com/huggingface/peft)."
ORIGINAL_SPACE_ID = "thuanz123/peft-sd-realfill"
SPACE_ID = os.getenv("SPACE_ID", ORIGINAL_SPACE_ID)
SHARED_UI_WARNING = f"""# Attention - This Space doesn't work in this shared UI. You can duplicate and use it with a paid private T4 GPU.
"""
if os.getenv("SYSTEM") == "spaces" and SPACE_ID != ORIGINAL_SPACE_ID:
SETTINGS = f'Settings'
else:
SETTINGS = "Settings"
CUDA_NOT_AVAILABLE_WARNING = f"""# Attention - Running on CPU.
You can assign a GPU in the {SETTINGS} tab if you are running this on HF Spaces.
"T4 small" is sufficient to run this demo.
"""
def show_warning(warning_text: str) -> gr.Blocks:
with gr.Blocks() as demo:
with gr.Box():
gr.Markdown(warning_text)
return demo
def update_output_files() -> dict:
paths = sorted(pathlib.Path("results").glob("*.pt"))
config_paths = sorted(pathlib.Path("results").glob("*.json"))
paths = paths + config_paths
paths = [path.as_posix() for path in paths] # type: ignore
return gr.update(value=paths or None)
def create_training_demo(trainer: Trainer, pipe: InferencePipeline) -> gr.Blocks:
with gr.Blocks() as demo:
base_model = gr.Dropdown(
choices=[
"runwayml/stable-diffusion-inpainting",
"stabilityai/stable-diffusion-2-inpainting",
],
value="stabilityai/stable-diffusion-2-inpainting",
label="Base Model",
visible=True,
)
resolution = gr.Dropdown(choices=["512"], value="512", label="Resolution", visible=False)
with gr.Row():
with gr.Box():
gr.Markdown("Training Data")
ref_images = gr.Files(label="Reference images")
target_image = gr.Files(label="Target image")
target_mask = gr.Files(label="Target mask")
gr.Markdown(
"""
- Upload reference images of the scene you are planning on training on.
- For a concept prompt, use a unique, made up word to avoid collisions.
- Guidelines for getting good results:
- 1-5 images of the object from different angles
- 2000 iterations should be good enough.
- LoRA Rank for unet - 8
- LoRA Alpha for unet - 16
- lora dropout - 0.1
- LoRA Bias for unet - `none`
- Uncheck `FP16` and `8bit-Adam` only if you have VRAM at least 32GB
- Experiment with various values for lora dropouts, enabling/disabling fp16 and 8bit-Adam
"""
)
with gr.Box():
gr.Markdown("Training Parameters")
num_training_steps = gr.Number(label="Number of Training Steps", value=2000, precision=0)
unet_learning_rate = gr.Number(label="Unet Learning Rate", value=2e-4)
text_encoder_learning_rate = gr.Number(label="Text Encoder Learning Rate", value=4e-5)
gradient_checkpointing = gr.Checkbox(label="Whether to use gradient checkpointing", value=True)
lora_rank = gr.Number(label="LoRA Rank for unet", value=8, precision=0)
lora_alpha = gr.Number(
label="LoRA Alpha for unet. scaling factor = lora_alpha/lora_r", value=16, precision=0
)
lora_dropout = gr.Number(label="lora dropout", value=0.1)
lora_bias = gr.Dropdown(
choices=["none", "all", "lora_only"],
value="none",
label="LoRA Bias for unet. This enables bias params to be trainable based on the bias type",
visible=True,
)
gradient_accumulation = gr.Number(label="Number of Gradient Accumulation", value=1, precision=0)
fp16 = gr.Checkbox(label="FP16", value=True)
use_8bit_adam = gr.Checkbox(label="Use 8bit Adam", value=True)
gr.Markdown(
"""
- It will take about 40-60 minutes to train for 2000 steps with a T4 GPU.
- You may want to try a small number of steps first, like 1, to see if everything works fine in your environment.
- Note that your trained models will be deleted when the second training is started. You can upload your trained model in the "Upload" tab.
"""
)
run_button = gr.Button("Start Training")
with gr.Box():
with gr.Row():
check_status_button = gr.Button("Check Training Status")
with gr.Column():
with gr.Box():
gr.Markdown("Message")
training_status = gr.Markdown()
output_files = gr.Files(label="Trained Model Files")
run_button.click(fn=pipe.clear)
run_button.click(
fn=trainer.run,
inputs=[
base_model,
resolution,
num_training_steps,
ref_images,
target_image,
target_mask,
unet_learning_rate,
text_encoder_learning_rate,
gradient_accumulation,
fp16,
use_8bit_adam,
gradient_checkpointing,
lora_rank,
lora_alpha,
lora_bias,
lora_dropout,
],
outputs=[
training_status,
output_files,
],
queue=False,
)
check_status_button.click(fn=trainer.check_if_running, inputs=None, outputs=training_status, queue=False)
check_status_button.click(fn=update_output_files, inputs=None, outputs=output_files, queue=False)
return demo
def find_model_files() -> list[str]:
curr_dir = pathlib.Path(__file__).parent
paths = sorted(curr_dir.glob('*'))
paths = [
path for path in paths
if (path / 'model_index.json').exists()
]
return [path.relative_to(curr_dir).as_posix() for path in paths]
def reload_realfill_model_list() -> dict:
return gr.update(choices=find_model_files())
def create_inference_demo(pipe: InferencePipeline) -> gr.Blocks:
with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
reload_button = gr.Button("Reload Model List")
realfill_model = gr.Dropdown(
choices=find_model_files(), label="RealFill Model File"
)
target_image = gr.Files(label="Target image")
target_mask = gr.Files(label="Target mask")
seed = gr.Slider(label="Seed", minimum=0, maximum=100000, step=1, value=1)
with gr.Accordion("Other Parameters", open=False):
num_steps = gr.Slider(label="Number of Steps", minimum=0, maximum=1000, step=1, value=50)
guidance_scale = gr.Slider(label="CFG Scale", minimum=0, maximum=50, step=0.1, value=7)
run_button = gr.Button("Generate")
gr.Markdown(
"""
- After training, you can press "Reload Model List" button to load your trained model names.
"""
)
with gr.Column():
result = gr.Image(label="Result")
reload_button.click(fn=reload_realfill_model_list, inputs=None, outputs=realfill_model)
run_button.click(
fn=pipe.run,
inputs=[
realfill_model,
target_image,
target_mask,
seed,
num_steps,
guidance_scale,
],
outputs=result,
queue=False,
)
seed.change(
fn=pipe.run,
inputs=[
realfill_model,
target_image,
target_mask,
seed,
num_steps,
guidance_scale,
],
outputs=result,
queue=False,
)
return demo
def create_upload_demo() -> gr.Blocks:
with gr.Blocks() as demo:
model_name = gr.Textbox(label="Model Name")
hf_token = gr.Textbox(label="Hugging Face Token (with write permission)")
upload_button = gr.Button("Upload")
with gr.Box():
gr.Markdown("Message")
result = gr.Markdown()
gr.Markdown(
"""
- You can upload your trained model to your private Model repo (i.e. https://huggingface.co/{your_username}/{model_name}).
- You can find your Hugging Face token [here](https://huggingface.co/settings/tokens).
"""
)
upload_button.click(fn=upload, inputs=[model_name, hf_token], outputs=result)
return demo
pipe = InferencePipeline()
trainer = Trainer()
with gr.Blocks(css="style.css") as demo:
if os.getenv("IS_SHARED_UI"):
show_warning(SHARED_UI_WARNING)
if not torch.cuda.is_available():
show_warning(CUDA_NOT_AVAILABLE_WARNING)
gr.Markdown(TITLE)
gr.Markdown(DESCRIPTION)
with gr.Tabs():
with gr.TabItem("Train"):
create_training_demo(trainer, pipe)
with gr.TabItem("Test"):
create_inference_demo(pipe)
with gr.TabItem("Upload"):
create_upload_demo()
demo.queue(default_enabled=False).launch(share=False)