Dreambooth migration to Keras 3

[kharshith-k]

Please find my PR for migrating Dreambooth tutorial from keras2 to keras3
Colab gist

[gemini-code-assist]

Summary of Changes

Hello @kharshith-k, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request focuses on upgrading the Dreambooth tutorial to align with the Keras 3 API. The changes involve updating imports, refactoring TensorFlow-specific operations to use Keras 3's backend-agnostic keras.ops, and adjusting the training loop and model components to reflect the latest Keras conventions. This ensures the tutorial remains functional and demonstrates best practices for users leveraging the multi-backend capabilities of Keras 3.

Highlights

• Keras 3 Migration: The Dreambooth tutorial has been comprehensively updated to be compatible with Keras 3, replacing tf.keras imports and API calls with the new keras namespace and keras.ops for backend-agnostic operations.
• API Modernization: Numerous TensorFlow-specific API calls, such as tf.concat, tf.split, tf.cast, tf.random.normal, and tf.keras.optimizers.experimental.AdamW, have been replaced with their Keras 3 equivalents like keras.ops.concatenate, keras.ops.split, keras.ops.cast, keras.random.normal, and keras.optimizers.AdamW.
• Mixed Precision and Loss Computation Refinement: The mixed precision handling has been streamlined by removing explicit get_scaled_loss and get_unscaled_gradients calls, and the loss computation logic within DreamBoothTrainer has been revised to explicitly use keras.ops.mean(keras.ops.square(...)) for MSE calculation.
• VAE Integration and Latent Sampling: The VAE (Variational Autoencoder) initialization has been simplified to directly use the ImageEncoder, with the sample_from_encoder_outputs method now handling the extraction of mean and variance internally, removing the need for manual layer slicing.
• Dependency and Utility Updates: The keras_cv installation command has been updated, tensorflow.experimental.numpy import has been removed, and utility functions like tf.keras.utils.get_file and tf.keras.callbacks.ModelCheckpoint have been migrated to their keras counterparts.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• examples/generative/dreambooth.py
  • Updated the 'Last modified' date in the file header.
  • Modified the pip install command to remove the specific keras_cv version and tensorflow installation.
  • Replaced from tensorflow import keras with import keras and updated tf.keras.utils.get_file to keras.utils.get_file.
  • Changed tf.convert_to_tensor for POS_IDS to np.array and added explicit GPU/CPU device placement logic for text embedding computation.
  • Updated keras_cv.layers.CenterCrop and keras_cv.layers.RandomFlip to keras.layers.CenterCrop and keras.layers.RandomFlip respectively.
  • Removed the import of tensorflow.experimental.numpy.
  • Changed DreamBoothTrainer inheritance from tf.keras.Model to keras.Model.
  • Replaced tf.concat with keras.ops.concatenate for image and embedded text concatenation.
  • Simplified VAE latent sampling by removing sample_from_encoder_outputs call and directly using self.vae(images, training=False).
  • Updated tnp.random.randint to tf.random.uniform for timestep sampling and tf.cast to keras.ops.cast for noisy latents.
  • Removed mixed precision scaling and unscaling logic from the train_step method.
  • Renamed compute_loss to compute_loss_for_dreambooth and implemented loss calculation using keras.ops.mean(keras.ops.square(...)).
  • Migrated tf.math.log, tf.math.exp, tf.range, tf.convert_to_tensor, tf.concat, tf.split, tf.clip_by_value, tf.exp, tf.random.normal, tf.shape to their keras.ops or keras.random equivalents.
  • Simplified save_weights and load_weights methods by removing save_format, options, and by_name arguments.
  • Updated tf.keras.mixed_precision.set_global_policy to keras.mixed_precision.set_global_policy.
  • Modified VAE initialization to directly use image_encoder instead of creating a tf.keras.Model wrapper.
  • Replaced tf.keras.optimizers.experimental.AdamW with keras.optimizers.AdamW.
  • Changed the checkpoint filename from dreambooth-unet.h5 to dreambooth-unet.weights.h5 and updated tf.keras.callbacks.ModelCheckpoint to keras.callbacks.ModelCheckpoint.
• examples/generative/ipynb/dreambooth.ipynb
  • Updated the 'Last modified' date in the notebook metadata.
  • Modified the pip install command in a code cell to remove the specific keras_cv version and tensorflow installation.
  • Replaced from tensorflow import keras with import keras in a code cell.
  • Updated tf.keras.utils.get_file to keras.utils.get_file in relevant code cells.
  • Changed tf.convert_to_tensor for POS_IDS to np.array and added explicit GPU/CPU device placement logic for text embedding computation in a code cell.
  • Updated keras_cv.layers.CenterCrop and keras_cv.layers.RandomFlip to keras.layers.CenterCrop and keras.layers.RandomFlip respectively in a code cell.
  • Removed the import of tensorflow.experimental.numpy in a code cell.
  • Changed DreamBoothTrainer inheritance from tf.keras.Model to keras.Model in a code cell.
  • Replaced tf.concat with keras.ops.concatenate for image and embedded text concatenation in a code cell.
  • Simplified VAE latent sampling by removing sample_from_encoder_outputs call and directly using self.vae(images, training=False) in a code cell.
  • Updated tnp.random.randint to tf.random.uniform for timestep sampling and tf.cast to keras.ops.cast for noisy latents in a code cell.
  • Removed mixed precision scaling and unscaling logic from the train_step method in a code cell.
  • Renamed compute_loss to compute_loss_for_dreambooth and implemented loss calculation using keras.ops.mean(keras.ops.square(...)) in a code cell.
  • Migrated tf.math.log, tf.math.exp, tf.range, tf.convert_to_tensor, tf.concat, tf.split, tf.clip_by_value, tf.exp, tf.random.normal, tf.shape to their keras.ops or keras.random equivalents in code cells.
  • Simplified save_weights and load_weights methods by removing save_format, options, and by_name arguments in a code cell.
  • Updated tf.keras.mixed_precision.set_global_policy to keras.mixed_precision.set_global_policy in a code cell.
  • Modified VAE initialization to directly use image_encoder instead of creating a tf.keras.Model wrapper in a code cell.
  • Replaced tf.keras.optimizers.experimental.AdamW with keras.optimizers.AdamW in a code cell.
  • Changed the checkpoint filename from dreambooth-unet.h5 to dreambooth-unet.weights.h5 and updated tf.keras.callbacks.ModelCheckpoint to keras.callbacks.ModelCheckpoint in a code cell.

Activity

• No human activity (comments, reviews) has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request successfully migrates the Dreambooth tutorial to Keras 3 by updating API calls to be backend-agnostic. The changes are well-aligned with the goals of Keras 3. I've included a couple of suggestions to improve code quality by refactoring a repetitive code block and removing some dead code. Overall, great work on the migration!

[sachinprasadhs]

Thanks for the tutorial and making this completely backend agnostic.

Few inline comments and few comments below for the removed part.

1. You have removed the complete description, keep the description back and modify the content according to the newly changed code.
2. Some of the short comments for code explainabiluty is removed, check what can be kept back and also add comments for the complex code part for the specific logic.