Stable Diffusion XL DremBooth Training¶

You can also check configs/stable_diffusion_xl_dreambooth/README.md file.

Configs¶

All configuration files are placed under the configs/stable_diffusion_xl_dreambooth folder.

Following is the example config fixed from the stable_diffusion_xl_dreambooth_lora_dog config file in configs/stable_diffusion_xl_dreambooth/stable_diffusion_xl_dreambooth_lora_dog.py:

from mmengine.config import read_base

with read_base():
    from .._base_.datasets.dog_dreambooth_xl import *
    from .._base_.default_runtime import *
    from .._base_.models.stable_diffusion_xl_lora import *
    from .._base_.schedules.stable_diffusion_500 import *

Finetuning the text encoder and UNet¶

The script also allows you to finetune the text_encoder along with the unet.

from mmengine.config import read_base

with read_base():
    from .._base_.datasets.dog_dreambooth_xl import *
    from .._base_.default_runtime import *
    from .._base_.models.stable_diffusion_xl_lora import *
    from .._base_.schedules.stable_diffusion_500 import *


model.update(finetune_text_encoder=True)  # fine tune text encoder

Finetuning with Full Parameters (without LoRA)¶

The script also allows you to finetune full parameters.

from mmengine.config import read_base

with read_base():
    from .._base_.datasets.dog_dreambooth_xl import *
    from .._base_.default_runtime import *
    from .._base_.models.stable_diffusion_xl import *
    from .._base_.schedules.stable_diffusion_500 import *

Finetuning with prior-preserving loss¶

The script also allows you to finetune with prior-preserving loss.

from mmengine.config import read_base

with read_base():
    from .._base_.datasets.dog_dreambooth_xl import *
    from .._base_.default_runtime import *
    from .._base_.models.stable_diffusion_xl_lora import *
    from .._base_.schedules.stable_diffusion_500 import *


train_dataloader.update(
    dataset=dict(
        class_image_config=dict(model=model.model),
        class_prompt='a photo of dog'),  # class_prompt=str means training with prior-preserving loss
)

Run DreamBooth training¶

Run DreamBooth train

# single gpu
$ diffengine train ${CONFIG_FILE}
# Example
$ diffengine train stable_diffusion_xl_dreambooth_lora_dog

# multi gpus
$ NPROC_PER_NODE=${GPU_NUM} diffengine train ${CONFIG_FILE}

Inference with diffusers¶

Once you have trained a model, specify the path to the saved model and utilize it for inference using the diffusers.pipeline module.

from pathlib import Path

import torch
from diffusers import DiffusionPipeline, AutoencoderKL
from peft import PeftModel

checkpoint = Path('work_dirs/stable_diffusion_xl_dreambooth_lora_dog/step499')
prompt = 'A photo of sks dog in a bucket'

vae = AutoencoderKL.from_pretrained(
    'madebyollin/sdxl-vae-fp16-fix',
    torch_dtype=torch.float16,
)
pipe = DiffusionPipeline.from_pretrained(
    'stabilityai/stable-diffusion-xl-base-1.0', vae=vae, torch_dtype=torch.float16)
pipe.to('cuda')
pipe.unet = PeftModel.from_pretrained(pipe.unet, checkpoint / "unet", adapter_name="default")
if (checkpoint / "text_encoder_one").exists():
    pipe.text_encoder_one = PeftModel.from_pretrained(
        pipe.text_encoder_one, checkpoint / "text_encoder_one", adapter_name="default"
    )
if (checkpoint / "text_encoder_two").exists():
    pipe.text_encoder_one = PeftModel.from_pretrained(
        pipe.text_encoder_two, checkpoint / "text_encoder_two", adapter_name="default"
    )

image = pipe(
    prompt,
    num_inference_steps=50,
    height=1024,
    width=1024,
).images[0]
image.save('demo.png')

You can check inference docs for inferencing other settings like Full Parameter Training without LoRA.

Results Example¶

stable_diffusion_xl_dreambooth_lora_dog¶

exampledog

You can check configs/stable_diffusion_xl_dreambooth/README.md for more details.