notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TirJ-SGQseby"
      },
      "source": [
        "# Vertex AI Model Garden Keras Stable Diffusion\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dwGLvtIeECLK"
      },
      "source": [
        "**_NOTE_**: This notebook has been tested in the following environment:\n",
        "\n",
        "* Python version = 3.9\n",
        "\n",
        "You can open this notebook directly in colab, or create [google managed](https://cloud.google.com/vertex-ai/docs/workbench/managed/create-instance) or [user managed](https://cloud.google.com/vertex-ai/docs/workbench/user-managed/create-new) workbench instances."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to use [Keras Stable Diffusion](https://keras.io/api/keras_cv/models/stable_diffusion) in Vertex AI Model Garden."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0z9r_mBmDeYh"
      },
      "source": [
        "### Objective\n",
        "\n",
        "* Run local inferences for pretrained or customized models\n",
        "\n",
        "* Deploy pretrained or customized models in Google Cloud Vertex AI\n",
        "\n",
        "* Finetune models in Google Cloud Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xxo28lDtDxn-"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "We use the dataset\n",
        "[Pokémon BLIP captions](https://huggingface.co/datasets/lambdalabs/pokemon-blip-captions) to show how to finetune the stable diffusion models.\n",
        "However, we'll use a slightly different version which was derived from the original\n",
        "dataset to fit better with `tf.data`. Refer to\n",
        "[the documentation](https://huggingface.co/datasets/sayakpaul/pokemon-blip-original-version)\n",
        "for more details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AEnkHABrDijz"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z__i0w0lCAsW"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Jvqs-ehKlaYh"
      },
      "outputs": [],
      "source": [
        "if \"google.colab\" in str(get_ipython()):\n",
        "    # Configs for colab notebooks.\n",
        "    ! pip3 install --upgrade --quiet google-cloud-aiplatform\n",
        "\n",
        "    # Automatically restart kernel after installs\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)\n",
        "\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()\n",
        "\n",
        "# Configs for all notebooks.\n",
        "! pip3 install --quiet keras-cv==0.4.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KEukV6uRk_S3"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
        "1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
        "\n",
        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint.\n",
        "\n",
        "1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
        "Cloud SDK uses the right project for all the commands in this notebook.\n",
        "\n",
        "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "### Set your project, region and buckets\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)\n",
        "\n",
        "You can change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations).\n",
        "\n",
        "You can create a storage bucket to store intermediate artifacts such as datasets, trained models etc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YjNCFxq0JxlA"
      },
      "outputs": [],
      "source": [
        "# The project and bucket are for experiments below.\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "\n",
        "# The form for BUCKET_URI is gs://<bucket-name>.\n",
        "BUCKET_URI = \"\"  # @param {type:\"string\"}\n",
        "REGION = \"us-central1\"  # @param {type: \"string\"}\n",
        "\n",
        "import os\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"keras\")\n",
        "DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
        "MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uDjp76aaLZY9"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5uv7-iDKLbO0"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZZFPe_GezXg8"
      },
      "source": [
        "### Define constants and common functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XcYUGwr-AJGY"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import os\n",
        "from datetime import datetime\n",
        "from io import BytesIO\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from google.cloud import storage\n",
        "from PIL import Image\n",
        "\n",
        "GCS_URI_PREFIX = \"gs://\"\n",
        "\n",
        "# Training constants.\n",
        "TRAINING_JOB_PREFIX = \"train\"\n",
        "TRAIN_CONTAINER_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-train:latest\"\n",
        ")\n",
        "TRAIN_MACHINE_TYPE = \"a2-highgpu-1g\"\n",
        "TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_A100\"\n",
        "TRAIN_NUM_GPU = 1\n",
        "RESOLUTION = 512\n",
        "\n",
        "# Prediction constants.\n",
        "PREDICTION_CONTAINER_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/keras-serve:latest\"\n",
        ")\n",
        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
        "DEPLOY_JOB_PREFIX = \"deploy\"\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str):\n",
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "def download_data_to_gcs(tar_filepath, gcs_bucket):\n",
        "    filename_with_ext = os.path.basename(tar_filepath)\n",
        "    filename_without_ext = filename_with_ext.replace(\".tar.gz\", \"\")\n",
        "    print(\"Download files from: \", tar_filepath)\n",
        "    ! wget $tar_filepath -O $filename_with_ext\n",
        "    ! mkdir -p $filename_without_ext\n",
        "    ! tar -xvf $filename_with_ext -C .\n",
        "\n",
        "    ! gsutil -m cp -r $filename_without_ext $gcs_bucket/\n",
        "    gcs_path = os.path.join(gcs_bucket, filename_without_ext)\n",
        "    print(\"Upload files to: \", gcs_path)\n",
        "    return gcs_path\n",
        "\n",
        "\n",
        "def download_gcs_file_to_local(gcs_uri: str, local_path: str):\n",
        "    \"\"\"Download a gcs file to a local path.\n",
        "\n",
        "    Args:\n",
        "      gcs_uri: A string of file path on GCS.\n",
        "      local_path: A string of local file path.\n",
        "    \"\"\"\n",
        "    if not gcs_uri.startswith(GCS_URI_PREFIX):\n",
        "        raise ValueError(f\"{gcs_uri} is not a GCS path starting with {GCS_URI_PREFIX}.\")\n",
        "    client = storage.Client()\n",
        "    os.makedirs(os.path.dirname(local_path), exist_ok=True)\n",
        "    with open(local_path, \"wb\") as f:\n",
        "        client.download_blob_to_file(gcs_uri, f)\n",
        "\n",
        "\n",
        "def deploy_model(model_path, service_account):\n",
        "\n",
        "    deploy_model_name = get_job_name_with_datetime(DEPLOY_JOB_PREFIX)\n",
        "    print(\"The deployed job name is: \", deploy_model_name)\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": \"keras-stable-diffusion-v1-4-001\",\n",
        "        \"MODEL_PATH\": f\"{model_path}\",\n",
        "        \"IMAGE_WIDTH\": f\"{RESOLUTION}\",\n",
        "        \"IMAGE_HEIGHT\": f\"{RESOLUTION}\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{deploy_model_name}-endpoint\")\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=deploy_model_name,\n",
        "        serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_predict_route=\"/predict\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=PREDICTION_MACHINE_TYPE,\n",
        "        accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
        "        accelerator_count=1,\n",
        "        min_replica_count=1,\n",
        "        max_replica_count=1,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "def base64_to_image(image_str):\n",
        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
        "    return image\n",
        "\n",
        "\n",
        "def display_image(image):\n",
        "    _ = plt.figure(figsize=(20, 15))\n",
        "    plt.grid(False)\n",
        "    plt.imshow(image)\n",
        "\n",
        "\n",
        "def display_image_grid(imgs, rows=2, cols=2):\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n",
        "    for i, img in enumerate(imgs):\n",
        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
        "    return grid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "epo-RHXzcBBT"
      },
      "source": [
        "## Run inferences\n",
        "\n",
        "This section shows how to run inferences with Keras Stable Diffusion models.\n",
        "\n",
        "1. Run inferences locally\n",
        "2. Run inferences with serving dockers\n",
        "\n",
        "You can run inferences with pre-trained models from Keras team, or your own finetuned models.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6zsa9vnBHhvO"
      },
      "outputs": [],
      "source": [
        "# Sets the model_path to empty to load the pre-trained model from Keras team.\n",
        "# Sets the model_path to a gcs uri to load the finetuned models.\n",
        "model_path = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ld39hkcIceE2"
      },
      "source": [
        "### Run inferences locally\n",
        "Local inferences can finish in seconds with GPUs.\n",
        "\n",
        "Load models first."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G1nCKVSac3Y5"
      },
      "outputs": [],
      "source": [
        "from keras_cv.models import StableDiffusion\n",
        "\n",
        "model = StableDiffusion(img_height=RESOLUTION, img_width=RESOLUTION, jit_compile=True)\n",
        "if model_path.startswith(GCS_URI_PREFIX):\n",
        "    local_model_path = \"/tmp/saved_model.h5\"\n",
        "    download_gcs_file_to_local(model_path, local_model_path)\n",
        "    model.diffusion_model.load_weights(local_model_path)\n",
        "elif model_path:\n",
        "    model.diffusion_model.load_weights(model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ABaCSIWuP-_G"
      },
      "source": [
        "Then run inferences."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pnyeVsh8RNI5"
      },
      "outputs": [],
      "source": [
        "batch_size = 1\n",
        "img = model.text_to_image(\n",
        "    prompt=\"a squirrel in Picasso style\",\n",
        "    batch_size=batch_size,  # How many images to generate at once\n",
        "    num_steps=25,  # Number of iterations (controls image quality)\n",
        "    seed=123,  # A fixed seed guarantees the same prompt always generates the same image\n",
        ")\n",
        "for i in range(batch_size):\n",
        "    display_image(img[i])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kY87SU9Adq4o"
      },
      "source": [
        "### Serve models with dockers\n",
        "When serve models with dockers, we will deploy models in Google Cloud Vertex AI. The default setting will use 1 V100 GPU for deployment.\n",
        "\n",
        "Please create a Service Account for serving with dockers if you do not have one yet.\n",
        "\n",
        "The model deployment will take ~10 minutes to finish."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yCB9vu7RenY6"
      },
      "outputs": [],
      "source": [
        "# The service account looks like:\n",
        "# '<account_name>@<project>.iam.gserviceaccount.com'\n",
        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "model, endpoint = deploy_model(\n",
        "    model_path=model_path,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        ")\n",
        "\n",
        "endpoint_id = endpoint.name\n",
        "print(\"endpoint id is: \", endpoint_id)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "72_BW_BgfvYT"
      },
      "source": [
        "Once deployed, you can send a batch of text prompts to the endpoint to generate images.\n",
        "\n",
        "Note, the inference time for the first request for a fresh deployment will need more time to process and take ~45 seconds on one V100 GPU. The inferences for further request is ~12 seconds on one V100 GPU per image."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U_jrNcZ5eVbH"
      },
      "outputs": [],
      "source": [
        "# # Loads an existing endpoint as below.\n",
        "# endpoint_id = <An Existing Endpoint ID>\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "instances = [\n",
        "    {\"prompt\": \"a squirrel in Picasso style\"},\n",
        "    {\"prompt\": \"a dog in Picasso style\"},\n",
        "    {\"prompt\": \"a cat in Picasso style\"},\n",
        "    {\"prompt\": \"a deer in Picasso style\"},\n",
        "]\n",
        "\n",
        "parameters = {\n",
        "    \"batch_size\": 1,  # How many images to generate at once\n",
        "    \"num_steps\": 25,  # Number of iterations (controls image quality)\n",
        "    \"seed\": 123,  # A fixed seed guarantees the same prompt always generates the same image\n",
        "}\n",
        "response = endpoint.predict(instances=instances, parameters=parameters)\n",
        "# prediction['predicted_image'] will contains the prediction images in a batch.\n",
        "# The batch size in this example is 1, and the visualization only parses the\n",
        "# first predicted image.\n",
        "images = [\n",
        "    base64_to_image(prediction[\"predicted_image\"][0])\n",
        "    for prediction in response.predictions\n",
        "]\n",
        "display_image_grid(images, rows=2, cols=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LiQF7fm6f842"
      },
      "source": [
        "### Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eqJyypt-f9K6"
      },
      "outputs": [],
      "source": [
        "# Undeploys models and deletes endpoints.\n",
        "endpoint.delete(force=True)\n",
        "# Deletes models.\n",
        "model.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RB_xY9ipr7ZU"
      },
      "source": [
        "## Finetune models\n",
        "This section shows how to finetune Keras Stable diffusion models with training dockers.\n",
        "\n",
        "If you would like to use finetuned models, please go to the section `Run inferences`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OD3TtaWs5b4v"
      },
      "source": [
        "### Download data\n",
        " We download the data to GCS storage for the experiments with training dockers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2TVB8MU-5i-q"
      },
      "outputs": [],
      "source": [
        "# Skips this step if you have already downloaded the dataset.\n",
        "download_data_to_gcs(\n",
        "    \"https://huggingface.co/datasets/sayakpaul/pokemon-blip-original-version/resolve/main/pokemon_dataset.tar.gz\",\n",
        "    DATA_BUCKET,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ee7Hzq8O5jgF"
      },
      "source": [
        "### Start training jobs\n",
        "We finetune 512*512 stable diffusion models with 1 epoch and it takes ~15 minutes to finish using 1 A100 GPU with default settings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "data_csv = os.path.join(DATA_BUCKET, \"pokemon_dataset/data.csv\")\n",
        "epochs = 1\n",
        "\n",
        "train_job_name = get_job_name_with_datetime(TRAINING_JOB_PREFIX)\n",
        "model_dir = os.path.join(MODEL_BUCKET, train_job_name)\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": TRAIN_MACHINE_TYPE,\n",
        "            \"accelerator_type\": TRAIN_ACCELERATOR_TYPE,\n",
        "            \"accelerator_count\": TRAIN_NUM_GPU,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"disk_spec\": {\n",
        "            \"boot_disk_type\": \"pd-ssd\",\n",
        "            \"boot_disk_size_gb\": 500,\n",
        "        },\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAIN_CONTAINER_URI,\n",
        "            \"command\": [],\n",
        "            \"env\": [\n",
        "                {\n",
        "                    \"name\": \"RESOLUTION\",\n",
        "                    \"value\": f\"{RESOLUTION}\",\n",
        "                },\n",
        "            ],\n",
        "            \"args\": [\n",
        "                f\"--epochs={epochs}\",\n",
        "                f\"--input_csv_path={data_csv}\",\n",
        "                f\"--output_model_dir={model_dir}\",\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "train_job = aiplatform.CustomJob(\n",
        "    display_name=train_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "train_job.run()\n",
        "\n",
        "model_path = os.path.join(model_dir, \"saved_model.h5\")\n",
        "print(\"The trained model is saved as: \", model_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wBlQ6FQlJhBi"
      },
      "source": [
        "After the training finishes, you can use `model_path` and then go to the `Run inferences` section above to run predictions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkH2nrpdp4sp"
      },
      "source": [
        "### Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ax6vQVZhp9pR"
      },
      "outputs": [],
      "source": [
        "train_job.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1dijQDiZWegt"
      },
      "source": [
        "## References\n",
        "\n",
        "- [Fine-tuning Stable Diffusion](https://keras.io/examples/generative/finetune_stable_diffusion/)\n",
        "- [StableDiffusion image-generation model](https://keras.io/api/keras_cv/models/stable_diffusion/)\n",
        "- [High-performance image generation using Stable Diffusion in KerasCV](https://keras.io/guides/keras_cv/generate_images_with_stable_diffusion/)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "model_garden_keras_stable_diffusion.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}