notebooks/community/model_garden/model_garden_mediapipe_image_generation.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "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 - Image Generation with MediaPipe\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_mediapipe_image_generation.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_mediapipe_image_generation.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_mediapipe_image_generation.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",
        "**NOTE**: The checkpoint and the dataset linked in this Colab are not owned or distributed by Google, and are made available by third parties. Please review the terms and conditions made available by the third parties before using the checkpoint and data.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how you can customize a [MediaPipe Image Generator](https://developers.google.com/mediapipe/solutions/vision/image_generator), a text-to-image generator, by adding Low-Rank Adaptation ([LoRA](https://arxiv.org/abs/2106.09685)) weights to generate images of specific people, objects, and styles.\n",
        "\n",
        "Using Vertex AI's Model Garden, we will retrain a standard diffusion model on specialized dataset of specific concepts, which are identified by unique tokens. With the new LoRA weights after training, the new model is able to generate images of the new concept when the token is specified in the text prompt.\n",
        "\n",
        "Once the model is customized with LoRA weights, it should only be used to generate images of the tokenized concept. It is no longer useful as a generalized image generation model. For more on customizing a MediaPipe Image Generator with LoRA weights, see the [MediaPipe documentation](https://developers.google.com/mediapipe/solutions/vision/image_generator#lora).\n",
        "\n",
        "NOTE: If you are creating LoRa weights to generate images of specific people and faces, only use this solution on your face or faces of people who have given you permission to do so."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Da-1fQ6gmWCG"
      },
      "source": [
        "### Objective\n",
        "\n",
        "* Set up a Google Cloud project with Vertex AI.\n",
        "* Train a text-to-image difussion model on a specialized dataset to create [LoRA](https://arxiv.org/abs/2106.09685) weights.\n",
        "* Customize a general image generator into a specialized generator that can inject specific objects, people, and styles into generated images.\n",
        "* Configure the newly trained Image Generator.\n",
        "* Download, upload, and deploy the new model\n",
        "\n",
        "### 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": "KEukV6uRk_S3"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K3_5Rg1jOlkz"
      },
      "source": [
        "### 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).\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. 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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z__i0w0lCAsW"
      },
      "source": [
        "### Authenticate for Colab\n",
        "**Note**: Skip this step if you are not using [Colab](https://colab.google/)\n",
        "\n",
        "Run the following commands to install dependencies and authenticate with Google Cloud on Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Jvqs-ehKlaYh"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade pip\n",
        "\n",
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    ! pip3 install --upgrade 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()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "source": [
        "### Set your project ID (`PROJECT_ID`)\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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6d87ae73686c"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "region"
      },
      "source": [
        "### Set the storage location (`REGION`)\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tTy1gX11kCJY"
      },
      "outputs": [],
      "source": [
        "REGION = \"\"  # @param {type: \"string\"}\n",
        "REGION_PREFIX = REGION.split(\"-\")[0]\n",
        "assert REGION_PREFIX in (\n",
        "    \"us\",\n",
        "    \"europe\",\n",
        "    \"asia\",\n",
        "), f'{REGION} is not supported. It must be prefixed by \"us\", \"asia\", or \"europe\".'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zgPO1eR3CYjk"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets and trained models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "If your bucket doesn't already exist, create your Cloud Storage bucket.\n",
        "\n",
        "**NOTE**: Only run the following cell if you do not already have a bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "960505627ddf"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PyQmSRbKA8r-"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "from datetime import datetime\n",
        "\n",
        "from google.cloud import aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "init_aip:mbsdk,all"
      },
      "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": {
        "cellView": "form",
        "id": "9wExiMUxFk91"
      },
      "outputs": [],
      "source": [
        "now = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temp/%s\" % now)\n",
        "\n",
        "MODEL_EXPORT_PATH = os.path.join(STAGING_BUCKET, \"model\")\n",
        "\n",
        "IMAGE_EXPORT_PATH = os.path.join(STAGING_BUCKET, \"image\")\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n6IFz75WGCam"
      },
      "source": [
        "### Define training and serving constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "riG_qUokg0XZ"
      },
      "outputs": [],
      "source": [
        "TRAINING_JOB_DISPLAY_NAME = \"mediapipe_stable_diffusion_%s\" % now\n",
        "TRAINING_CONTAINER = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/mediapipe-stable-diffusion-train\"\n",
        "TRAINING_MACHINE_TYPE = \"a2-highgpu-1g\"\n",
        "TRAINING_ACCELERATOR_TYPE = \"NVIDIA_TESLA_A100\"\n",
        "TRAINING_ACCELERATOR_COUNT = 1\n",
        "\n",
        "PREDICTION_CONTAINER_URI = f\"{REGION_PREFIX}-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240403_0836_RC00\"\n",
        "PREDICTION_PORT = 7080\n",
        "PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
        "UPLOAD_MODEL_NAME = \"mediapipe_stable_diffusion_model_%s\" % now"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XDq9TiRUc7dV"
      },
      "source": [
        "## Train a customized Image Generator\n",
        "\n",
        "In this section, we will customize the Image Generator by training the model on images of [teapots](https://github.com/google/dreambooth/tree/main/dataset/teapot) from the [DreamBooth dataset](https://github.com/google/dreambooth/tree/main). Using the LoRA weights created through training, the new model will be able to inject teapots into generated images.\n",
        "\n",
        "This is a simple example implementation. You can modify the following cells to further customize the notebook."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bNn4SYZftLbV"
      },
      "source": [
        "### Choose the pre-trained model to download\n",
        "\n",
        "The MediaPipe Image Generator task requires you to download a trained model that matches the `runwayml/stable-diffusion-v1-5 EMA-only` model format, based on the following model: [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "edfE2BGVtPSi"
      },
      "outputs": [],
      "source": [
        "unet_url = \"https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/unet/diffusion_pytorch_model.bin\"  # @param {type:\"string\"}\n",
        "vae_url = \"https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/vae/diffusion_pytorch_model.bin\"  # @param {type:\"string\"}\n",
        "text_encoder_url = \"https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/text_encoder/pytorch_model.bin\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Oj4XAhsn78as"
      },
      "source": [
        "### Prepare input data for training\n",
        "\n",
        "Customizing a model for image generation requires a dataset that contains sample pictures of the concept instance that you want to use in the generation. The concept can be a person, object, or style.\n",
        "\n",
        "**Object**\n",
        "![](https://storage.googleapis.com/mediapipe-assets/documentation/object_lora.png)\n",
        "\n",
        "**Person**\n",
        "![](https://storage.googleapis.com/mediapipe-assets/documentation/person_lora.png)\n",
        "\n",
        "**Style**\n",
        "![](https://storage.googleapis.com/mediapipe-assets/documentation/style_lora.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZkMYDhDqef9M"
      },
      "source": [
        "You must also assign a unique token to the new concept. The prompt should include the token, which is \"monadikos\" in this case, followed by a word that describes the concept to generate. In this example, we are using \"A monadikos teapot\". The images from the [teapots](https://github.com/google/dreambooth/tree/main/dataset/teapot) dataset can be downloaded from Google Cloud Storage.\n",
        "\n",
        "The customized model will recognize the term \"monadikos teapot\", and inject an image of a teapot into the generated images."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "IndQ_m6ddUEM"
      },
      "outputs": [],
      "source": [
        "# Path to the training data folder.\n",
        "training_data_path = \"gs://mediapipe-tasks/image_generator/teapot\"  # @param {type:\"string\"}\n",
        "# An instance description of the training data.\n",
        "training_data_prompt = \"A monadikos teapot\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aaff6f5be7f6"
      },
      "source": [
        "### Set training options\n",
        "\n",
        "The Image Generator comes with a set of pre-defined hyperparameter (`HParams`) settings that work best for specific situations. You should select a template that best matches your use case."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lYF-4ydq0k0u"
      },
      "source": [
        "You can further customize hyperparameters like the learning rate and the number of training steps (epochs). For more information on these hyperparameters, see the [Google Machine Learning glossary](https://developers.google.com/machine-learning/glossary)\n",
        "\n",
        "To set custom training parameters, adjust the values for the following hyperparameters:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "um_XKbmpTaHx"
      },
      "outputs": [],
      "source": [
        "# Parameters about training configuration\n",
        "# The learning rate to use for gradient descent training.\n",
        "learning_rate: float = 0.00001  # @param {type:\"number\"}\n",
        "# Number of training steps. If set to 0, uses the default value.\n",
        "num_train_steps: int = 0  # @param {type:\"integer\"}\n",
        "# Save the checkpoint in every n steps.\n",
        "save_checkpoints_every_n: int = 100  # @param {type:\"integer\"}\n",
        "# Batch size for training.\n",
        "batch_size: int = 1  # @param {type:\"integer\"}\n",
        "\n",
        "# Dataset-related parameters\n",
        "# Whether to use random horizontal flip on data.\n",
        "random_flip: bool = False  # @param {type:\"boolean\"}\n",
        "# Whether to use random largest square crop.\n",
        "random_crop: bool = False  # @param {type:\"boolean\"}\n",
        "# Whether to distort the color of the image (jittering order is random).\n",
        "random_color_jitter: bool = False  # @param {type:\"boolean\"}\n",
        "\n",
        "# Hyperparameters for LoRA tuning\n",
        "# The rank in the low-rank matrices. If set to 0, uses the default value.\n",
        "lora_rank: int = 0  # @param {type:\"integer\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VbVUkoZda3Y4"
      },
      "source": [
        "Alternatively, you can also use one of our pre-trained models for these templates. These templates are already customized and already contain LoRA weights:\n",
        "* [Object (berry bowls)](https://storage.googleapis.com/mediapipe-tasks/image_generator/object/pytorch_lora_weights.bin)\n",
        "* [Face](https://storage.googleapis.com/mediapipe-tasks/image_generator/face/pytorch_lora_weights.bin)\n",
        "* [Style](https://storage.googleapis.com/mediapipe-tasks/image_generator/style/pytorch_lora_weights.bin)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bj31hD4W02Ui"
      },
      "outputs": [],
      "source": [
        "template = \"\"  # @param [\"\", \"face\", \"object\", \"style\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BuMz_X97F6yr"
      },
      "source": [
        "## Test the customized Image Generator model\n",
        "\n",
        "After training the custom model, we will generate images to examine the quality of the customized model. You can provide a text prompt below and configure options for generating the test images."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5NwgyEvYGOTt"
      },
      "source": [
        "### Define the test generation prompt\n",
        "\n",
        "Specify the prompt to use to test the customized model. Note that a variation of the token, \"monadikos teapots\", is included in the prompt. If you are customizing this notebook with another dataset, set a token to describe the object, person, or style depicted in the training data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2dSmBAauGEWw"
      },
      "outputs": [],
      "source": [
        "prompt: str = \"Two monadikos teapots on a table\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nvhSeCt1GbUy"
      },
      "source": [
        "### Configure the parameters to generate test images\n",
        "\n",
        "Set configuration options to run image generation with the customized model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iwmsxMzDGhOO"
      },
      "outputs": [],
      "source": [
        "# Number of steps to run inference.\n",
        "number_inference_steps: int = 50  # @param {type:\"integer\"}\n",
        "# Classifier-free guidance weight to use during inference. Weight must be is >= 1.0.\n",
        "guidance_scale: float = 7.5  # @param {type:\"number\"}\n",
        "#  Number of generated images per prompt.\n",
        "number_generated_images: int = 8  # @param {type:\"integer\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HwcCjwlBTQIz"
      },
      "source": [
        "### Tune the image generator with LoRA\n",
        "Tune the Image Generator with LoRA and generate new images based on your prompt. This can take up to 10 minutes on Vertex AI with a A100 GPU.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "aec22792ee84"
      },
      "outputs": [],
      "source": [
        "model_export_path = MODEL_EXPORT_PATH\n",
        "image_export_path = IMAGE_EXPORT_PATH\n",
        "\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": TRAINING_MACHINE_TYPE,\n",
        "            \"accelerator_type\": TRAINING_ACCELERATOR_TYPE,\n",
        "            \"accelerator_count\": TRAINING_ACCELERATOR_COUNT,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAINING_CONTAINER,\n",
        "            \"command\": [],\n",
        "            \"args\": [\n",
        "                \"--task_name=stable_diffusion\",\n",
        "                \"--model_export_path=%s\" % model_export_path,\n",
        "                \"--image_export_path=%s\" % image_export_path,\n",
        "                \"--training_data_path=%s\" % training_data_path,\n",
        "                \"--training_data_prompt='%s'\" % training_data_prompt,\n",
        "                \"--prompt='%s'\" % prompt,\n",
        "                \"--hparams_template=%s\" % template,\n",
        "                \"--hparams=%s\"\n",
        "                % json.dumps(\n",
        "                    {\n",
        "                        \"learning_rate\": learning_rate,\n",
        "                        \"num_train_steps\": num_train_steps,\n",
        "                        \"save_checkpoints_every_n\": save_checkpoints_every_n,\n",
        "                        \"batch_size\": batch_size,\n",
        "                        \"random_flip\": random_flip,\n",
        "                        \"random_crop\": random_crop,\n",
        "                        \"random_color_jitter\": random_color_jitter,\n",
        "                        \"lora_rank\": lora_rank,\n",
        "                        \"torch_vae\": vae_url,\n",
        "                        \"torch_unet\": unet_url,\n",
        "                        \"torch_text_encoder\": text_encoder_url,\n",
        "                    }\n",
        "                ),\n",
        "                \"--generator_hparams=%s\"\n",
        "                % json.dumps(\n",
        "                    {\n",
        "                        \"number_inference_steps\": number_inference_steps,\n",
        "                        \"guidance_scale\": guidance_scale,\n",
        "                        \"number_generated_images\": number_generated_images,\n",
        "                    }\n",
        "                ),\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "training_job = aiplatform.CustomJob(\n",
        "    display_name=TRAINING_JOB_DISPLAY_NAME,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "training_job.run()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hBhI0hYnCCCI"
      },
      "source": [
        "## Download images and model\n",
        "\n",
        "After training and testing the new model, you can download the generated images and the new customized model. The LoRA weights from training can also be used with the MediaPipe Tasks ImageGenerator API for on-device applications."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jdm1dOgsMH9O"
      },
      "source": [
        "### Download generated images\n",
        "\n",
        "Download and preview the generated images at different checkpoints.\n",
        "Inspecting the generated images helps to determine the best checkpoint and avoid underfitting or overfitting."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xsrMHpRvMErS"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "def copy_image(images_source, images_dest):\n",
        "    os.makedirs(images_dest, exist_ok=True)\n",
        "    ! gsutil cp -r {images_source}/* {images_dest}\n",
        "\n",
        "\n",
        "local_image_path = \"./images/\"\n",
        "copy_image(IMAGE_EXPORT_PATH, local_image_path)\n",
        "\n",
        "steps_samples = {}\n",
        "for filename in os.listdir(local_image_path):\n",
        "    absolute_path = os.path.join(local_image_path, filename)\n",
        "    if os.path.isfile(absolute_path):\n",
        "        parsed_name = filename.split(\"_\")\n",
        "        step = int(parsed_name[1])\n",
        "        if step not in steps_samples:\n",
        "            steps_samples[step] = []\n",
        "        image = plt.imread(absolute_path)\n",
        "        steps_samples[step].append(image)\n",
        "\n",
        "for step in sorted(steps_samples.keys()):\n",
        "    print(f\"\\nGenerated image with training steps {step}:\")\n",
        "    for image in steps_samples[step]:\n",
        "        plt.figure(figsize=(20, 10), dpi=150)\n",
        "        plt.axis(\"off\")\n",
        "        plt.imshow(image)\n",
        "        plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LPCvcX88b_sO"
      },
      "source": [
        "By default, the last checkpoint is used for deployment. However, we can customize that here based on the above visual inspection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ONtgDzGncCc3"
      },
      "outputs": [],
      "source": [
        "deployed_checkpoint: int = -1  # @param {type:\"integer\"}\n",
        "if deployed_checkpoint == -1:\n",
        "    deployed_checkpoint = num_train_steps\n",
        "valid_checkpoints = list(\n",
        "    range(save_checkpoints_every_n, num_train_steps + 1, save_checkpoints_every_n)\n",
        ")\n",
        "if deployed_checkpoint not in valid_checkpoints:\n",
        "    raise ValueError(\"Invalid checkpoint chosen for deployment.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zcKzIa5QeIIU"
      },
      "source": [
        "### Download model\n",
        "\n",
        "After fine-tuning and evaluating the model, you can download the model and checkpoints."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NYuQowyZEtxK"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "\n",
        "def copy_model(model_source, model_dest):\n",
        "    os.makedirs(model_dest, exist_ok=True)\n",
        "    ! gsutil -m cp -r {model_source}/* {model_dest}\n",
        "\n",
        "\n",
        "local_model_path = \"/models\"\n",
        "copy_model(MODEL_EXPORT_PATH, local_model_path)\n",
        "\n",
        "! tar czf models.tar.gz {local_model_path}/*\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import files\n",
        "\n",
        "    files.download(\"models.tar.gz\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-gqs4Kc4u6U3"
      },
      "source": [
        "## Upload and deploy to Vertex AI\n",
        "\n",
        "This section shows the way to test with trained models.\n",
        "1. Upload and deploy models to the [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction)\n",
        "2. Get [online predictions](https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions) from the deployed model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y_HV9vMbvF_6"
      },
      "source": [
        "### Upload model to Vertex AI Model Registry"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OL6wdXf-1rta"
      },
      "outputs": [],
      "source": [
        "lora_id = os.path.join(MODEL_EXPORT_PATH, f\"checkpoint_{deployed_checkpoint}\")\n",
        "\n",
        "serving_env = {\n",
        "    \"TASK\": \"text-to-image\",\n",
        "    \"MODEL_ID\": \"runwayml/stable-diffusion-v1-5\",\n",
        "    \"LORA_ID\": lora_id,\n",
        "    \"DEPLOY_SOURCE\": \"notebook\",\n",
        "}\n",
        "\n",
        "model = aiplatform.Model.upload(\n",
        "    display_name=UPLOAD_MODEL_NAME,\n",
        "    serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
        "    serving_container_ports=[PREDICTION_PORT],\n",
        "    serving_container_predict_route=\"/predictions/diffusers_serving\",\n",
        "    serving_container_health_route=\"/ping\",\n",
        "    serving_container_environment_variables=serving_env,\n",
        ")\n",
        "\n",
        "model.wait()\n",
        "\n",
        "print(\"The uploaded model name is: \", UPLOAD_MODEL_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NjXg3QHw2GCT"
      },
      "source": [
        "### Deploy the uploaded model\n",
        "\n",
        "You 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 around 1 minute to finish."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "BLTA7nRcw57D"
      },
      "outputs": [],
      "source": [
        "# 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",
        "endpoint = aiplatform.Endpoint.create(display_name=f\"{UPLOAD_MODEL_NAME}-endpoint\")\n",
        "model.deploy(\n",
        "    endpoint=endpoint,\n",
        "    machine_type=PREDICTION_MACHINE_TYPE,\n",
        "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
        "    accelerator_count=1,\n",
        "    deploy_request_timeout=1800,\n",
        "    service_account=service_account,\n",
        "    system_labels={\n",
        "        \"NOTEBOOK_NAME\": \"model_garden_mediapipe_image_generation.ipynb\"\n",
        "    },\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tMfwrH32_uIM"
      },
      "source": [
        "The docker container still needs to download and load the model after the endpoint is created. Therefore, we recommend waiting for 3 extra minutes before proceeding to the next cell.\n",
        "\n",
        "Once deployed, you can send a batch of text prompts to the endpoint to generate images."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wr7wg9JQ_tjM"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "from io import BytesIO\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image\n",
        "\n",
        "instances = [\n",
        "    {\"prompt\": \"Two monadikos teapots on a table\"},\n",
        "    {\"prompt\": \"Two monadikos teapots on the floor\"},\n",
        "]\n",
        "response = endpoint.predict(instances=instances)\n",
        "\n",
        "plt.figure()\n",
        "_, grid = plt.subplots(1, len(instances))\n",
        "for cell, prediction in zip(grid, response.predictions):\n",
        "    image = Image.open(BytesIO(base64.b64decode(prediction)))\n",
        "    cell.imshow(image)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kkH2nrpdp4sp"
      },
      "source": [
        "## Clean up\n",
        "After the export is complete, you can delete your training job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Ax6vQVZhp9pR"
      },
      "outputs": [],
      "source": [
        "if training_job.list(filter=f'display_name=\"{TRAINING_JOB_DISPLAY_NAME}\"'):\n",
        "    training_job.delete()\n",
        "# Undeploys models and deletes endpoints.\n",
        "endpoint.delete(force=True)\n",
        "model.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CkRKXILVJc7Q"
      },
      "source": [
        "You can also remove the output data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vdIZsAstJTcA"
      },
      "outputs": [],
      "source": [
        "!gsutil rm -r {STAGING_BUCKET}"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_mediapipe_image_generation.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}