py-dreambooth

Py-Dreambooth

Py-Dreambooth is a Python package that makes it easy to create AI avatar images from photos of you, your family, friends, or pets!

• Tasks are pre-configured with the most efficient defaults, which greatly streamlines the workload. A number of helper functions are also provided.
• This is designed to be modular and extensible to many different models. Currently supported models are the Stable Diffusion Dreambooth, Stable Diffusion Dreambooth LoRA, and Stable Diffusion XL Dreambooth LoRA.
• This is designed to give you the flexibility to choose local or cloud resources to train your model and generate images.

⚙️ How to Install

pip install py-dreambooth

🚀 Quick Start

• Prepare about 10-20 high-quality solo selfie photos (jpg or png) and put them in a specific directory.
• Please run on a machine with a GPU of 16GB or more. (If you're fine-tuning SDXL, you'll need 24GB of VRAM.)

from py_dreambooth.dataset import LocalDataset
from py_dreambooth.model import SdDreamboothModel
from py_dreambooth.trainer import LocalTrainer
from py_dreambooth.utils.image_helpers import display_images
from py_dreambooth.utils.prompt_helpers import make_prompt

DATA_DIR = "data"  # The directory where you put your prepared photos
OUTPUT_DIR = "models"  

dataset = LocalDataset(DATA_DIR)
dataset = dataset.preprocess_images(detect_face=True)

SUBJECT_NAME = "<YOUR-NAME>"  
CLASS_NAME = "person"

model = SdDreamboothModel(subject_name=SUBJECT_NAME, class_name=CLASS_NAME)
trainer = LocalTrainer(output_dir=OUTPUT_DIR)

predictor = trainer.fit(model, dataset)

# Use the prompt helper to create an awesome AI avatar!
prompt = next(make_prompt(SUBJECT_NAME, CLASS_NAME))

images = predictor.predict(
    prompt, height=768, width=512, num_images_per_prompt=2,
)

display_images(images, fig_size=10)

🏃‍♀️ Tutorials

• Take a look at the 01-local-tutorial.ipynb file to learn how to get it running on your local Jupyter Notebook.
• If you're interested in running it with AWS cloud resources, take a look at the 02-aws-tutorial.ipynb file.
• Or, get started right away with the Google Colab Notebook here!

📚 Documentation

• Full documentation can be found here: https://py-dreambooth.readthedocs.io.

References

• DreamBooth: Fine-Tuning Text-to-Image Diffusion Models for Subject-Driven Generation (Paper)
• LoRA: Low-Rank Adaptation of Large Language Models (Paper)
• Fine-Tune Text-to-Image Stable Diffusion Models with Amazon SageMaker JumpStart (Blog)
• Training Stable Diffusion with Dreambooth Using 🧨 Diffusers (Blog)
• Diffusers: DreamBooth Training Example
• Diffusers: DreamBooth Training Example for Stable Diffusion XL (SDXL)