Generate fun images of your dog

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

In this demo notebook, we demonstrate how to use the JumpStart APIs to fine-tune the stable diffusion model to your dog images dataset and deploy the fine-tuned model. To execute this notebook, you will need a collection of dog images. You may upload as little as five images to your local folder and run the notebook.

Note: This notebook contains a very simplified version of the features available for Stable Diffusion in JumpStart. Please refer to the Introduction to JumpStart - Text to Image notebook for a more comprehensive list of Stable Diffusion features available in JumpStart.

Note: To run this notebook, you would need ml.g4dn.2xlarge instance type for training and for inference.

1. Set Up

2. Fine-tune the pre-trained model on a custom dataset

   • Retrieve Training Artifacts
   • Set Training parameters
   • Start Training
   • Deploy and run inference on the fine-tuned model

1. Set Up

Set up credentials and create a local training directory where you will upload your training images.

Please upload images of your dog to training_images local folder and change use_local_images=True.

Upload dataset to S3

Next, we upload the dataset to S3 bucket. If the bucket does not exists, we create a new bucket.

If you have an existing bucket you would like to use, please replace the training_bucket with your bucket in the cell above and avoid executing the following cell.

Next we upload the training datasets (images and dataset_info.json) to the S3 bucket.

2. Fine-tune the pre-trained model on a custom dataset

2.1. Retrieve Training Artifacts

Here, we retrieve the training docker container, the training algorithm source, and the pre-trained base model. Note that model_version="*" fetches the latest model.

2.2. Set Training parameters

There are two kinds of parameters that need to be set for training. The first one are the parameters for the training job. These include: (i) Training data path: This is S3 folder in which the input data is stored, (ii) Output path: This the s3 folder in which the training output is stored. (iii) Training instance type: This indicates the type of machine on which to run the training. We defined the training instance type above to fetch the correct train_image_uri.

The second set of parameters are algorithm specific training hyper-parameters.

For algorithm specific hyper-parameters, we start by fetching python dictionary of the training hyper-parameters that the algorithm accepts with their default values. This can then be overridden to custom values.

2.3. Start Training

We start by creating the estimator object with all the required assets and then launch the training job. It takes less than 10 mins on the default dataset.

2.4. Deploy and run inference on the fine-tuned model

A trained model does nothing on its own. We now want to use the model to perform inference. For this example, that means predicting the bounding boxes of an image. We start by retrieving the JumpStart artifacts for deploying an endpoint. However, instead of base_predictor, we deploy the sd_estimator that we have fine-tuned.

Next, we query the finetuned model, parse the response and display the generated image. Please execute the following cells.

Next, we delete the endpoint corresponding to the finetuned model.

Conclusion

In this notebook, we saw a simple workflow on how you can fine-tune the stable diffusion text-to-image model on your dataset with a small set of images. You can adapt the notebook your dataset by uploading images of the desired subject and changing the prompts. For instance, if you would like to generate images of your cat, please upload cat images in the first step and change dog to cat in the instance_prompt before training and while inocking endpoint with fine-tuned model.

This notebook contains a barebone code to train and deploy the stable diffusion model. Please refer to the Introduction to JumpStart - Text to Image for additional features such as (i) How to deploy a pre-trained Stable Diffusion model (more than 80 available in JumpStart), (ii) How to set parameters such as num_steps, guidance scale during inference, (iii) Prompt Engineering, (iv) How to set training related parameters.

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.