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.
pip install py-dreambooth- 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)- 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!
- Full documentation can be found here: https://py-dreambooth.readthedocs.io.
- 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)