CINXE.COM

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/generative/dreambooth/" />  <meta property="og:title" content="Keras documentation: DreamBooth"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: DreamBooth"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>DreamBooth</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink active" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink2" href="/examples/generative/ddim/">Denoising Diffusion Implicit Models</a> <a class="nav-sublink2" href="/examples/generative/random_walks_with_stable_diffusion_3/">A walk through latent space with Stable Diffusion 3</a> <a class="nav-sublink2 active" href="/examples/generative/dreambooth/">DreamBooth</a> <a class="nav-sublink2" href="/examples/generative/ddpm/">Denoising Diffusion Probabilistic Models</a> <a class="nav-sublink2" href="/examples/generative/fine_tune_via_textual_inversion/">Teach StableDiffusion new concepts via Textual Inversion</a> <a class="nav-sublink2" href="/examples/generative/finetune_stable_diffusion/">Fine-tuning Stable Diffusion</a> <a class="nav-sublink2" href="/examples/generative/vae/">Variational AutoEncoder</a> <a class="nav-sublink2" href="/examples/generative/dcgan_overriding_train_step/">GAN overriding Model.train_step</a> <a class="nav-sublink2" href="/examples/generative/wgan_gp/">WGAN-GP overriding Model.train_step</a> <a class="nav-sublink2" href="/examples/generative/conditional_gan/">Conditional GAN</a> <a class="nav-sublink2" href="/examples/generative/cyclegan/">CycleGAN</a> <a class="nav-sublink2" href="/examples/generative/gan_ada/">Data-efficient GANs with Adaptive Discriminator Augmentation</a> <a class="nav-sublink2" href="/examples/generative/deep_dream/">Deep Dream</a> <a class="nav-sublink2" href="/examples/generative/gaugan/">GauGAN for conditional image generation</a> <a class="nav-sublink2" href="/examples/generative/pixelcnn/">PixelCNN</a> <a class="nav-sublink2" href="/examples/generative/stylegan/">Face image generation with StyleGAN</a> <a class="nav-sublink2" href="/examples/generative/vq_vae/">Vector-Quantized Variational Autoencoders</a> <a class="nav-sublink2" href="/examples/generative/random_walks_with_stable_diffusion/">A walk through latent space with Stable Diffusion</a> <a class="nav-sublink2" href="/examples/generative/neural_style_transfer/">Neural style transfer</a> <a class="nav-sublink2" href="/examples/generative/adain/">Neural Style Transfer with AdaIN</a> <a class="nav-sublink2" href="/examples/generative/gpt2_text_generation_with_keras_hub/">GPT2 Text Generation with KerasHub</a> <a class="nav-sublink2" href="/examples/generative/text_generation_gpt/">GPT text generation from scratch with KerasHub</a> <a class="nav-sublink2" href="/examples/generative/text_generation_with_miniature_gpt/">Text generation with a miniature GPT</a> <a class="nav-sublink2" href="/examples/generative/lstm_character_level_text_generation/">Character-level text generation with LSTM</a> <a class="nav-sublink2" href="/examples/generative/text_generation_fnet/">Text Generation using FNet</a> <a class="nav-sublink2" href="/examples/generative/midi_generation_with_transformer/">Music Generation with Transformer Models</a> <a class="nav-sublink2" href="/examples/generative/molecule_generation/">Drug Molecule Generation with VAE</a> <a class="nav-sublink2" href="/examples/generative/wgan-graphs/">WGAN-GP with R-GCN for the generation of small molecular graphs</a> <a class="nav-sublink2" href="/examples/generative/real_nvp/">Density estimation using Real NVP</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparam Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/generative/'>Generative Deep Learning</a> / DreamBooth </div> <div class='k-content'> <h1 id="dreambooth">DreamBooth</h1> Author: <a href="https://twitter.com/RisingSayak">Sayak Paul</a>, <a href="https://twitter.com/algo_diver">Chansung Park</a> Date created: 2023/02/01 Last modified: 2023/02/05 Description: Implementing DreamBooth. <div class='example_version_banner keras_2'>ⓘ This example uses Keras 2</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/generative/ipynb/dreambooth.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/generative/dreambooth.py">GitHub source</a> <hr /> <h2 id="introduction">Introduction</h2> In this example, we implement DreamBooth, a fine-tuning technique to teach new visual concepts to text-conditioned Diffusion models with just 3 - 5 images. DreamBooth was proposed in <a href="https://arxiv.org/abs/2208.12242">DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation</a> by Ruiz et al. DreamBooth, in a sense, is similar to the <a href="https://keras.io/examples/generative/finetune_stable_diffusion/">traditional way of fine-tuning a text-conditioned Diffusion model except</a> for a few gotchas. This example assumes that you have basic familiarity with Diffusion models and how to fine-tune them. Here are some reference examples that might help you to get familiarized quickly: <ul> <li><a href="https://keras.io/guides/keras_cv/generate_images_with_stable_diffusion/">High-performance image generation using Stable Diffusion in KerasCV</a></li> <li><a href="https://keras.io/examples/generative/fine_tune_via_textual_inversion/">Teach StableDiffusion new concepts via Textual Inversion</a></li> <li><a href="https://keras.io/examples/generative/finetune_stable_diffusion/">Fine-tuning Stable Diffusion</a></li> </ul> First, let's install the latest versions of KerasCV and TensorFlow. <div class="codehilite"><pre><code>!pip install -q -U keras_cv==0.6.0 !pip install -q -U tensorflow </code></pre></div> If you're running the code, please ensure you're using a GPU with at least 24 GBs of VRAM. <hr /> <h2 id="initial-imports">Initial imports</h2> <div class="codehilite"><pre><code>import math import keras_cv import matplotlib.pyplot as plt import numpy as np import tensorflow as tf from imutils import paths from tensorflow import keras </code></pre></div> <hr /> <h2 id="usage-of-dreambooth">Usage of DreamBooth</h2> ... is very versatile. By teaching Stable Diffusion about your favorite visual concepts, you can <ul> <li> Recontextualize objects in interesting ways: <img alt="" src="https://i.imgur.com/4Da9ozw.png" /> </li> </ul> <ul> <li> Generate artistic renderings of the underlying visual concept: <img alt="" src="https://i.imgur.com/nI2N8bI.png" /> </li> </ul> And many other applications. We welcome you to check out the original <a href="https://arxiv.org/abs/2208.12242">DreamBooth paper</a> in this regard. <hr /> <h2 id="download-the-instance-and-class-images">Download the instance and class images</h2> DreamBooth uses a technique called "prior preservation" to meaningfully guide the training procedure such that the fine-tuned models can still preserve some of the prior semantics of the visual concept you're introducing. To know more about the idea of "prior preservation" refer to <a href="https://dreambooth.github.io/">this document</a>. Here, we need to introduce a few key terms specific to DreamBooth: <ul> <li>Unique class: Examples include "dog", "person", etc. In this example, we use "dog".</li> <li>Unique identifier: A unique identifier that is prepended to the unique class while forming the "instance prompts". In this example, we use "sks" as this unique identifier.</li> <li>Instance prompt: Denotes a prompt that best describes the "instance images". An example prompt could be - "f"a photo of {unique_id} {unique_class}". So, for our example, this becomes - "a photo of sks dog".</li> <li>Class prompt: Denotes a prompt without the unique identifier. This prompt is used for generating "class images" for prior preservation. For our example, this prompt is - "a photo of dog".</li> <li>Instance images: Denote the images that represent the visual concept you're trying to teach aka the "instance prompt". This number is typically just 3 - 5. We typically gather these images ourselves.</li> <li>Class images: Denote the images generated using the "class prompt" for using prior preservation in DreamBooth training. We leverage the pre-trained model before fine-tuning it to generate these class images. Typically, 200 - 300 class images are enough.</li> </ul> In code, this generation process looks quite simply: <div class="codehilite"><pre><code>from tqdm import tqdm import numpy as np import hashlib import keras_cv import PIL import os class_images_dir = "class-images" os.makedirs(class_images_dir, exist_ok=True) model = keras_cv.models.StableDiffusion(img_width=512, img_height=512, jit_compile=True) class_prompt = "a photo of dog" num_imgs_to_generate = 200 for i in tqdm(range(num_imgs_to_generate)): images = model.text_to_image( class_prompt, batch_size=3, ) idx = np.random.choice(len(images)) selected_image = PIL.Image.fromarray(images[idx]) hash_image = hashlib.sha1(selected_image.tobytes()).hexdigest() image_filename = os.path.join(class_images_dir, f"{hash_image}.jpg") selected_image.save(image_filename) </code></pre></div> To keep the runtime of this example short, the authors of this example have gone ahead and generated some class images using <a href="https://colab.research.google.com/gist/sayakpaul/6b5de345d29cf5860f84b6d04d958692/generate_class_priors.ipynb">this notebook</a>. Note that prior preservation is an optional technique used in DreamBooth, but it almost always helps in improving the quality of the generated images. <div class="codehilite"><pre><code>instance_images_root = tf.keras.utils.get_file( origin="https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/instance-images.tar.gz", untar=True, ) class_images_root = tf.keras.utils.get_file( origin="https://huggingface.co/datasets/sayakpaul/sample-datasets/resolve/main/class-images.tar.gz", untar=True, ) </code></pre></div> <hr /> <h2 id="visualize-images">Visualize images</h2> First, let's load the image paths. <div class="codehilite"><pre><code>instance_image_paths = list(paths.list_images(instance_images_root)) class_image_paths = list(paths.list_images(class_images_root)) </code></pre></div> Then we load the images from the paths. <div class="codehilite"><pre><code>def load_images(image_paths): images = [np.array(keras.utils.load_img(path)) for path in image_paths] return images </code></pre></div> And then we make use a utility function to plot the loaded images. <div class="codehilite"><pre><code>def plot_images(images, title=None): plt.figure(figsize=(20, 20)) for i in range(len(images)): ax = plt.subplot(1, len(images), i + 1) if title is not None: plt.title(title) plt.imshow(images[i]) plt.axis("off") </code></pre></div> Instance images: <div class="codehilite"><pre><code>plot_images(load_images(instance_image_paths[:5])) </code></pre></div> <img alt="png" src="/img/examples/generative/dreambooth/dreambooth_16_0.png" /> Class images: <div class="codehilite"><pre><code>plot_images(load_images(class_image_paths[:5])) </code></pre></div> <img alt="png" src="/img/examples/generative/dreambooth/dreambooth_18_0.png" /> <hr /> <h2 id="prepare-datasets">Prepare datasets</h2> Dataset preparation includes two stages: (1): preparing the captions, (2) processing the images. <h3 id="prepare-the-captions">Prepare the captions</h3> <div class="codehilite"><pre><code># Since we're using prior preservation, we need to match the number # of instance images we're using. We just repeat the instance image paths # to do so. new_instance_image_paths = [] for index in range(len(class_image_paths)): instance_image = instance_image_paths[index % len(instance_image_paths)] new_instance_image_paths.append(instance_image) # We just repeat the prompts / captions per images. unique_id = "sks" class_label = "dog" instance_prompt = f"a photo of {unique_id} {class_label}" instance_prompts = [instance_prompt] * len(new_instance_image_paths) class_prompt = f"a photo of {class_label}" class_prompts = [class_prompt] * len(class_image_paths) </code></pre></div> Next, we embed the prompts to save some compute. <div class="codehilite"><pre><code>import itertools # The padding token and maximum prompt length are specific to the text encoder. # If you're using a different text encoder be sure to change them accordingly. padding_token = 49407 max_prompt_length = 77 # Load the tokenizer. tokenizer = keras_cv.models.stable_diffusion.SimpleTokenizer() # Method to tokenize and pad the tokens. def process_text(caption): tokens = tokenizer.encode(caption) tokens = tokens + [padding_token] * (max_prompt_length - len(tokens)) return np.array(tokens) # Collate the tokenized captions into an array. tokenized_texts = np.empty( (len(instance_prompts) + len(class_prompts), max_prompt_length) ) for i, caption in enumerate(itertools.chain(instance_prompts, class_prompts)): tokenized_texts[i] = process_text(caption) # We also pre-compute the text embeddings to save some memory during training. POS_IDS = tf.convert_to_tensor([list(range(max_prompt_length))], dtype=tf.int32) text_encoder = keras_cv.models.stable_diffusion.TextEncoder(max_prompt_length) gpus = tf.config.list_logical_devices("GPU") # Ensure the computation takes place on a GPU. # Note that it's done automatically when there's a GPU present. # This example just attempts at showing how you can do it # more explicitly. with tf.device(gpus[0].name): embedded_text = text_encoder( [tf.convert_to_tensor(tokenized_texts), POS_IDS], training=False ).numpy() # To ensure text_encoder doesn't occupy any GPU space. del text_encoder </code></pre></div> <hr /> <h2 id="prepare-the-images">Prepare the images</h2> <div class="codehilite"><pre><code>resolution = 512 auto = tf.data.AUTOTUNE augmenter = keras.Sequential( layers=[ keras_cv.layers.CenterCrop(resolution, resolution), keras_cv.layers.RandomFlip(), keras.layers.Rescaling(scale=1.0 / 127.5, offset=-1), ] ) def process_image(image_path, tokenized_text): image = tf.io.read_file(image_path) image = tf.io.decode_png(image, 3) image = tf.image.resize(image, (resolution, resolution)) return image, tokenized_text def apply_augmentation(image_batch, embedded_tokens): return augmenter(image_batch), embedded_tokens def prepare_dict(instance_only=True): def fn(image_batch, embedded_tokens): if instance_only: batch_dict = { "instance_images": image_batch, "instance_embedded_texts": embedded_tokens, } return batch_dict else: batch_dict = { "class_images": image_batch, "class_embedded_texts": embedded_tokens, } return batch_dict return fn def assemble_dataset(image_paths, embedded_texts, instance_only=True, batch_size=1): dataset = tf.data.Dataset.from_tensor_slices((image_paths, embedded_texts)) dataset = dataset.map(process_image, num_parallel_calls=auto) dataset = dataset.shuffle(5, reshuffle_each_iteration=True) dataset = dataset.batch(batch_size) dataset = dataset.map(apply_augmentation, num_parallel_calls=auto) prepare_dict_fn = prepare_dict(instance_only=instance_only) dataset = dataset.map(prepare_dict_fn, num_parallel_calls=auto) return dataset </code></pre></div> <hr /> <h2 id="assemble-dataset">Assemble dataset</h2> <div class="codehilite"><pre><code>instance_dataset = assemble_dataset( new_instance_image_paths, embedded_text[: len(new_instance_image_paths)], ) class_dataset = assemble_dataset( class_image_paths, embedded_text[len(new_instance_image_paths) :], instance_only=False, ) train_dataset = tf.data.Dataset.zip((instance_dataset, class_dataset)) </code></pre></div> <hr /> <h2 id="check-shapes">Check shapes</h2> Now that the dataset has been prepared, let's quickly check what's inside it. <div class="codehilite"><pre><code>sample_batch = next(iter(train_dataset)) print(sample_batch[0].keys(), sample_batch[1].keys()) for k in sample_batch[0]: print(k, sample_batch[0][k].shape) for k in sample_batch[1]: print(k, sample_batch[1][k].shape) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>dict_keys(['instance_images', 'instance_embedded_texts']) dict_keys(['class_images', 'class_embedded_texts']) instance_images (1, 512, 512, 3) instance_embedded_texts (1, 77, 768) class_images (1, 512, 512, 3) class_embedded_texts (1, 77, 768) </code></pre></div> </div> During training, we make use of these keys to gather the images and text embeddings and concat them accordingly. <hr /> <h2 id="dreambooth-training-loop">DreamBooth training loop</h2> Our DreamBooth training loop is very much inspired by <a href="https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py">this script</a> provided by the Diffusers team at Hugging Face. However, there is an important difference to note. We only fine-tune the UNet (the model responsible for predicting noise) and don't fine-tune the text encoder in this example. If you're looking for an implementation that also performs the additional fine-tuning of the text encoder, refer to <a href="https://github.com/sayakpaul/dreambooth-keras/">this repository</a>. <div class="codehilite"><pre><code>import tensorflow.experimental.numpy as tnp class DreamBoothTrainer(tf.keras.Model): # Reference: # https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py def __init__( self, diffusion_model, vae, noise_scheduler, use_mixed_precision=False, prior_loss_weight=1.0, max_grad_norm=1.0, **kwargs, ): super().__init__(**kwargs) self.diffusion_model = diffusion_model self.vae = vae self.noise_scheduler = noise_scheduler self.prior_loss_weight = prior_loss_weight self.max_grad_norm = max_grad_norm self.use_mixed_precision = use_mixed_precision self.vae.trainable = False def train_step(self, inputs): instance_batch = inputs[0] class_batch = inputs[1] instance_images = instance_batch["instance_images"] instance_embedded_text = instance_batch["instance_embedded_texts"] class_images = class_batch["class_images"] class_embedded_text = class_batch["class_embedded_texts"] images = tf.concat([instance_images, class_images], 0) embedded_texts = tf.concat([instance_embedded_text, class_embedded_text], 0) batch_size = tf.shape(images)[0] with tf.GradientTape() as tape: # Project image into the latent space and sample from it. latents = self.sample_from_encoder_outputs(self.vae(images, training=False)) # Know more about the magic number here: # https://keras.io/examples/generative/fine_tune_via_textual_inversion/ latents = latents * 0.18215 # Sample noise that we'll add to the latents. noise = tf.random.normal(tf.shape(latents)) # Sample a random timestep for each image. timesteps = tnp.random.randint( 0, self.noise_scheduler.train_timesteps, (batch_size,) ) # Add noise to the latents according to the noise magnitude at each timestep # (this is the forward diffusion process). noisy_latents = self.noise_scheduler.add_noise( tf.cast(latents, noise.dtype), noise, timesteps ) # Get the target for loss depending on the prediction type # just the sampled noise for now. target = noise # noise_schedule.predict_epsilon == True # Predict the noise residual and compute loss. timestep_embedding = tf.map_fn( lambda t: self.get_timestep_embedding(t), timesteps, dtype=tf.float32 ) model_pred = self.diffusion_model( [noisy_latents, timestep_embedding, embedded_texts], training=True ) loss = self.compute_loss(target, model_pred) if self.use_mixed_precision: loss = self.optimizer.get_scaled_loss(loss) # Update parameters of the diffusion model. trainable_vars = self.diffusion_model.trainable_variables gradients = tape.gradient(loss, trainable_vars) if self.use_mixed_precision: gradients = self.optimizer.get_unscaled_gradients(gradients) gradients = [tf.clip_by_norm(g, self.max_grad_norm) for g in gradients] self.optimizer.apply_gradients(zip(gradients, trainable_vars)) return {m.name: m.result() for m in self.metrics} def get_timestep_embedding(self, timestep, dim=320, max_period=10000): half = dim // 2 log_max_period = tf.math.log(tf.cast(max_period, tf.float32)) freqs = tf.math.exp( -log_max_period * tf.range(0, half, dtype=tf.float32) / half ) args = tf.convert_to_tensor([timestep], dtype=tf.float32) * freqs embedding = tf.concat([tf.math.cos(args), tf.math.sin(args)], 0) return embedding def sample_from_encoder_outputs(self, outputs): mean, logvar = tf.split(outputs, 2, axis=-1) logvar = tf.clip_by_value(logvar, -30.0, 20.0) std = tf.exp(0.5 * logvar) sample = tf.random.normal(tf.shape(mean), dtype=mean.dtype) return mean + std * sample def compute_loss(self, target, model_pred): # Chunk the noise and model_pred into two parts and compute the loss # on each part separately. # Since the first half of the inputs has instance samples and the second half # has class samples, we do the chunking accordingly. model_pred, model_pred_prior = tf.split( model_pred, num_or_size_splits=2, axis=0 ) target, target_prior = tf.split(target, num_or_size_splits=2, axis=0) # Compute instance loss. loss = self.compiled_loss(target, model_pred) # Compute prior loss. prior_loss = self.compiled_loss(target_prior, model_pred_prior) # Add the prior loss to the instance loss. loss = loss + self.prior_loss_weight * prior_loss return loss def save_weights(self, filepath, overwrite=True, save_format=None, options=None): # Overriding this method will allow us to use the `ModelCheckpoint` # callback directly with this trainer class. In this case, it will # only checkpoint the `diffusion_model` since that's what we're training # during fine-tuning. self.diffusion_model.save_weights( filepath=filepath, overwrite=overwrite, save_format=save_format, options=options, ) def load_weights(self, filepath, by_name=False, skip_mismatch=False, options=None): # Similarly override `load_weights()` so that we can directly call it on # the trainer class object. self.diffusion_model.load_weights( filepath=filepath, by_name=by_name, skip_mismatch=skip_mismatch, options=options, ) </code></pre></div> <hr /> <h2 id="trainer-initialization">Trainer initialization</h2> <div class="codehilite"><pre><code># Comment it if you are not using a GPU having tensor cores. tf.keras.mixed_precision.set_global_policy("mixed_float16") use_mp = True # Set it to False if you're not using a GPU with tensor cores. image_encoder = keras_cv.models.stable_diffusion.ImageEncoder() dreambooth_trainer = DreamBoothTrainer( diffusion_model=keras_cv.models.stable_diffusion.DiffusionModel( resolution, resolution, max_prompt_length ), # Remove the top layer from the encoder, which cuts off the variance and only # returns the mean. vae=tf.keras.Model( image_encoder.input, image_encoder.layers[-2].output, ), noise_scheduler=keras_cv.models.stable_diffusion.NoiseScheduler(), use_mixed_precision=use_mp, ) # These hyperparameters come from this tutorial by Hugging Face: # https://github.com/huggingface/diffusers/tree/main/examples/dreambooth learning_rate = 5e-6 beta_1, beta_2 = 0.9, 0.999 weight_decay = (1e-2,) epsilon = 1e-08 optimizer = tf.keras.optimizers.experimental.AdamW( learning_rate=learning_rate, weight_decay=weight_decay, beta_1=beta_1, beta_2=beta_2, epsilon=epsilon, ) dreambooth_trainer.compile(optimizer=optimizer, loss="mse") </code></pre></div> <hr /> <h2 id="train">Train!</h2> We first calculate the number of epochs, we need to train for. <div class="codehilite"><pre><code>num_update_steps_per_epoch = train_dataset.cardinality() max_train_steps = 800 epochs = math.ceil(max_train_steps / num_update_steps_per_epoch) print(f"Training for {epochs} epochs.") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Training for 4 epochs. </code></pre></div> </div> And then we start training! <div class="codehilite"><pre><code>ckpt_path = "dreambooth-unet.h5" ckpt_callback = tf.keras.callbacks.ModelCheckpoint( ckpt_path, save_weights_only=True, monitor="loss", mode="min", ) dreambooth_trainer.fit(train_dataset, epochs=epochs, callbacks=[ckpt_callback]) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Epoch 1/4 200/200 [==============================] - 301s 462ms/step - loss: 0.1203 Epoch 2/4 200/200 [==============================] - 94s 469ms/step - loss: 0.1139 Epoch 3/4 200/200 [==============================] - 94s 469ms/step - loss: 0.1016 Epoch 4/4 200/200 [==============================] - 94s 469ms/step - loss: 0.1231 <keras.callbacks.History at 0x7f19726600a0> </code></pre></div> </div> <hr /> <h2 id="experiments-and-inference">Experiments and inference</h2> We ran various experiments with a slightly modified version of this example. Our experiments are based on <a href="https://github.com/sayakpaul/dreambooth-keras/">this repository</a> and are inspired by <a href="https://huggingface.co/blog/dreambooth">this blog post</a> from Hugging Face. First, let's see how we can use the fine-tuned checkpoint for running inference. <div class="codehilite"><pre><code># Initialize a new Stable Diffusion model. dreambooth_model = keras_cv.models.StableDiffusion( img_width=resolution, img_height=resolution, jit_compile=True ) dreambooth_model.diffusion_model.load_weights(ckpt_path) # Note how the unique identifier and the class have been used in the prompt. prompt = f"A photo of {unique_id} {class_label} in a bucket" num_imgs_to_gen = 3 images_dreamboothed = dreambooth_model.text_to_image(prompt, batch_size=num_imgs_to_gen) plot_images(images_dreamboothed, prompt) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>By using this model checkpoint, you acknowledge that its usage is subject to the terms of the CreativeML Open RAIL-M license at https://raw.githubusercontent.com/CompVis/stable-diffusion/main/LICENSE 50/50 [==============================] - 42s 160ms/step </code></pre></div> </div> <img alt="png" src="/img/examples/generative/dreambooth/dreambooth_40_1.png" /> Now, let's load checkpoints from a different experiment we conducted where we also fine-tuned the text encoder along with the UNet: <div class="codehilite"><pre><code>unet_weights = tf.keras.utils.get_file( origin="https://huggingface.co/chansung/dreambooth-dog/resolve/main/lr%409e-06-max_train_steps%40200-train_text_encoder%40True-unet.h5" ) text_encoder_weights = tf.keras.utils.get_file( origin="https://huggingface.co/chansung/dreambooth-dog/resolve/main/lr%409e-06-max_train_steps%40200-train_text_encoder%40True-text_encoder.h5" ) dreambooth_model.diffusion_model.load_weights(unet_weights) dreambooth_model.text_encoder.load_weights(text_encoder_weights) images_dreamboothed = dreambooth_model.text_to_image(prompt, batch_size=num_imgs_to_gen) plot_images(images_dreamboothed, prompt) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>Downloading data from https://huggingface.co/chansung/dreambooth-dog/resolve/main/lr%409e-06-max_train_steps%40200-train_text_encoder%40True-unet.h5 3439088208/3439088208 [==============================] - 67s 0us/step Downloading data from https://huggingface.co/chansung/dreambooth-dog/resolve/main/lr%409e-06-max_train_steps%40200-train_text_encoder%40True-text_encoder.h5 492466760/492466760 [==============================] - 9s 0us/step 50/50 [==============================] - 8s 159ms/step </code></pre></div> </div> <img alt="png" src="/img/examples/generative/dreambooth/dreambooth_42_1.png" /> The default number of steps for generating an image in <code>text_to_image()</code> <a href="https://github.com/keras-team/keras-cv/blob/3575bc3b944564fe15b46b917e6555aa6a9d7be0/keras_cv/models/stable_diffusion/stable_diffusion.py#L73">is 50</a>. Let's increase it to 100. <div class="codehilite"><pre><code>images_dreamboothed = dreambooth_model.text_to_image( prompt, batch_size=num_imgs_to_gen, num_steps=100 ) plot_images(images_dreamboothed, prompt) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>100/100 [==============================] - 16s 159ms/step </code></pre></div> </div> <img alt="png" src="/img/examples/generative/dreambooth/dreambooth_44_1.png" /> Feel free to experiment with different prompts (don't forget to add the unique identifier and the class label!) to see how the results change. We welcome you to check out our codebase and more experimental results <a href="https://github.com/sayakpaul/dreambooth-keras#results">here</a>. You can also read <a href="https://huggingface.co/blog/dreambooth">this blog post</a> to get more ideas. <hr /> <h2 id="acknowledgements">Acknowledgements</h2> <ul> <li>Thanks to the <a href="https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth.py">DreamBooth example script</a> provided by Hugging Face which helped us a lot in getting the initial implementation ready quickly.</li> <li>Getting DreamBooth to work on human faces can be challenging. We have compiled some general recommendations <a href="https://github.com/sayakpaul/dreambooth-keras#notes-on-preparing-data-for-dreambooth-training-of-faces">here</a>. Thanks to <a href="https://no.linkedin.com/in/abhi1thakur">Abhishek Thakur</a> for helping with these.</li> </ul> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#dreambooth'>DreamBooth</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#initial-imports'>Initial imports</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#usage-of-dreambooth'>Usage of DreamBooth</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#download-the-instance-and-class-images'>Download the instance and class images</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#visualize-images'>Visualize images</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#prepare-datasets'>Prepare datasets</a> </div> <div class='k-outline-depth-3'> <a href='#prepare-the-captions'>Prepare the captions</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#prepare-the-images'>Prepare the images</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#assemble-dataset'>Assemble dataset</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#check-shapes'>Check shapes</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#dreambooth-training-loop'>DreamBooth training loop</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#trainer-initialization'>Trainer initialization</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#train'>Train!</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#experiments-and-inference'>Experiments and inference</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#acknowledgements'>Acknowledgements</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>