notebooks/community/vertex_endpoints/torchserve/dreambooth_stablediffusion.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "18ebbd838e32"
      },
      "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": "219f1b1fe8fe"
      },
      "source": [
        "# Deploy and host a Stable Diffusion model on Vertex AI\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vertex_endpoints/torchserve/dreambooth_stablediffusion.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",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vertex_endpoints/torchserve/dreambooth_stablediffusion.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/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/vertex_endpoints/torchserve/dreambooth_stablediffusion.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": "fce05a8186d6"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to deploy and host a fine-tuned [Stable Diffusion 1.5](https://huggingface.co/runwayml/stable-diffusion-v1-5) model on Vertex AI. For hosting, you use the PyTorch 3 container built for Vertex AI with [TorchServe](https://pytorch.org/serve/index.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c76216b03fec"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this tutorial, you learn how to host and deploy a Stable Diffusion 1.5 model on Vertex AI.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services:\n",
        "\n",
        "+ Vertex AI `Model` resource\n",
        "+ Vertex AI `Endpoint` resource\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "+ Create a `torchserve` handler for responding to prediction requests.\n",
        "+ Upload a Stable Diffusion 1.5 model on a prebuilt PyTorch container in Vertex AI.\n",
        "+ Deploy a model to a Vertex AI Endpoint.\n",
        "+ Send requests to the endpoint and parse the responses using Vertex AI Prediction service."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c6deba5a8557"
      },
      "source": [
        "### Model\n",
        "\n",
        "This notebook uses a collection of model artifacts fine-tuned to generate images of a small dog. These are the same images used in the original [DreamBooth paper](https://dreambooth.github.io/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "911dc651ea9c"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI models\n",
        "* Vertex AI endpoints\n",
        "* Vertex AI prediction\n",
        "* Cloud Storage\n",
        "* (Optionally) Vertex AI Workbench\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0d36bd3d53fa"
      },
      "source": [
        "## Hardware requirements\n",
        "\n",
        "This notebook requires that you use a GPU with a sufficient amount of VRAM available. It was tested on a `NVIDIA Tesla A100 GPU` with 85 GB of VRAM. Run the following cell to ensure that you have the correct hardware configuration."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3dd4022552e5"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi --query-gpu=name,memory.total,memory.free --format=\"csv,noheader\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a782627e5f73"
      },
      "source": [
        "### Create a user-managed notebook on Vertex AI\n",
        "\n",
        "If you are using Vertex AI Workbench, you can create a notebook with the correct configuration by doing the following:\n",
        "\n",
        "+ Go to [Vertex AI Workbench](https://console.cloud.google.com/vertex-ai/workbench/user-managed) in the Google Cloud Console.\n",
        "+ Click **New Notebook** and then click **PyTorch 1.13** > **With 1 NVIDIA T4**.\n",
        "+ In the **New notebook** dialog box, click **Advanced Options**. The **Create a user-managed notebook** page opens up.\n",
        "+ In the **Create a user-managed notebook** page, do the following:\n",
        "  * In the **Notebook name** box, type a name for your notebook, for example \"my-stablediffusion-nb\".\n",
        "  * In the **Machine type** drop-down, select **A2 highgpu** > **a2-highgpu-1g**.\n",
        "  * In the **GPU type** drop-down, select **NVIDIA Tesla A100**.\n",
        "  * Check the box next to **Install NVIDIA GPU driver automatically for me**\n",
        "  * Expand **Disk(s)** and do the following:\n",
        "    - Under **Boot disk type**, select **SSD Persistent Disk**.\n",
        "    - Under **Data disk type**, select **SSD Persistent Disk**.\n",
        "  * Click **Create**."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FAPoU8Sm5E6e"
      },
      "source": [
        "<div style=\"background:#feefe3; padding:5px; color:#aa0000\">\n",
        "<strong>Caution:</strong> Using a Vertex AI Workbench notebook with the above configuration can increase your costs significantly. You can estimate your costs using the <a href=\"https://cloud.google.com/products/calculator\"><u>costs calculator</u></a>.</div>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0bb4201cc99a"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook.\n",
        "\n",
        "**Note**: You might need to change the version of PyTorch (`torch`) installed by `pip`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9c769df171a6"
      },
      "outputs": [],
      "source": [
        "%%writefile requirements.txt\n",
        "diffusers\n",
        "ftfy\n",
        "google-cloud-aiplatform\n",
        "gradio\n",
        "ninja\n",
        "tensorboard==1.15.0\n",
        "torch\n",
        "torchaudio\n",
        "torchvision\n",
        "torchserve\n",
        "torch-model-archiver\n",
        "torch-workflow-archiver\n",
        "transformers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e46804ac90d8"
      },
      "outputs": [],
      "source": [
        "%pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58707a750154"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "77c11549298a"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "294df346a918"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1j6f9HApxa"
      },
      "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). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WReHDGG5g0XY"
      },
      "source": [
        "#### Set your 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": {
        "id": "oM1iC_MfAts1"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-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": [
        "#### 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": {
        "id": "region"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBCra4QMA2wR"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74ccc9e52986"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "de775a3773ba"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "254614fa0c46"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ef21552ccea8"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "603adbbf0532"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f6b2ccc891ed"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cb5c4ca3e851"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7348591eda51"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "import math\n",
        "\n",
        "import torch\n",
        "from diffusers import StableDiffusionPipeline\n",
        "from google.cloud import aiplatform\n",
        "from IPython import display\n",
        "from PIL import Image\n",
        "from torch import autocast"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "697566b5f660"
      },
      "source": [
        "## Optional: View model inferences\n",
        "\n",
        "Before uploading the model to Vertex AI, you can review the expected output from the model. The model used in this notebook is available for your use and can be downloaded from Cloud Storage. This download may take a few minutes to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d63df8d91215"
      },
      "outputs": [],
      "source": [
        "!gsutil -m cp gs://cloud-samples-data/vertex-ai/model-deployment/models/stable-diffusion/model_artifacts.zip \\\n",
        "   ."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "89d613ae9573"
      },
      "outputs": [],
      "source": [
        "!unzip model_artifacts.zip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9192ea4f3b57"
      },
      "source": [
        "### Create new images\n",
        "\n",
        "With everything in place, you can now generate new images from the Stable Diffusion model. First you must load your model into a `StableDiffusionPipeline`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cd1f30223b79"
      },
      "outputs": [],
      "source": [
        "model_path = \"model_artifacts\"\n",
        "\n",
        "pipe = StableDiffusionPipeline.from_pretrained(\n",
        "    model_path, torch_dtype=torch.float16\n",
        ").to(\"cuda\")\n",
        "\n",
        "g_cuda = None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dd91520f58cf"
      },
      "outputs": [],
      "source": [
        "g_cuda = torch.Generator(device=\"cuda\")\n",
        "seed = 52362\n",
        "g_cuda.manual_seed(seed)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1e9c5d734b9f"
      },
      "source": [
        "With the model loaded into a `StableDiffusionPipeline`, you can now generate results (inferences) from the model. Each set of inference requires an input (called a [prompt](https://learnprompting.org/)) that specifies what the model should create.\n",
        "\n",
        "You can also vary other inputs into the model, as shown in the following cell."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a28b73de55ce"
      },
      "outputs": [],
      "source": [
        "prompt = \"photo of examplePup dog in a Monet style\"\n",
        "\n",
        "num_samples = 4\n",
        "num_batches = 1\n",
        "num_columns = 2\n",
        "guidance_scale = 10\n",
        "num_inference_steps = 50\n",
        "height = 512\n",
        "width = 512"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "11a7bb79de59"
      },
      "outputs": [],
      "source": [
        "def image_grid(imgs, cols):\n",
        "    total = len(imgs)\n",
        "    rows = math.ceil(total / cols)\n",
        "\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h))\n",
        "    grid_w, grid_h = grid.size\n",
        "\n",
        "    for i, img in enumerate(imgs):\n",
        "        grid.paste(img, box=(i % cols * w, i // cols * h))\n",
        "    return grid\n",
        "\n",
        "\n",
        "all_images = []\n",
        "for _ in range(num_batches):\n",
        "    with autocast(\"cuda\"):\n",
        "        images = pipe(\n",
        "            [prompt] * num_samples,\n",
        "            height=height,\n",
        "            width=width,\n",
        "            num_inference_steps=num_inference_steps,\n",
        "            guidance_scale=guidance_scale,\n",
        "        ).images\n",
        "        all_images.extend(images)\n",
        "\n",
        "\n",
        "grid = image_grid(all_images, num_columns)\n",
        "grid"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbc72963ff88"
      },
      "source": [
        "## Deploy the model to Vertex AI\n",
        "\n",
        "You can host your Stable Diffusion 1.5 model on a Vertex AI endpoint where you can get inferences from it online. Uploading your model is a four step process: \n",
        "\n",
        "1. Create a custom TorchServe handler.\n",
        "1. Upload the model artifacts onto Cloud Storage.\n",
        "2. Create a Vertex AI model with the model artifacts and a prebuilt PyTorch container image.\n",
        "3. Deploy the Vertex AI model onto an endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eafbb0e0-40e6-43a0-a38e-edc54323da51"
      },
      "source": [
        "### Create the custom TorchServe handler\n",
        "\n",
        "The model deployed to Vertex AI uses [TorchServe](https://pytorch.org/serve/) to handle requests and return responses from the model. You must create a custom TorchServe handler to include in with the model artifacts uploaded to Vertex AI.\n",
        "\n",
        "The handler file should be included in the directory with the other model artifacts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "94567a87-9d74-4c87-a749-306ddaf01b61"
      },
      "outputs": [],
      "source": [
        "%%writefile model_artifacts/handler.py\n",
        "\n",
        "\"\"\"Customized handler for Stable Diffusion 1.5.\"\"\"\n",
        "import base64\n",
        "import logging\n",
        "from io import BytesIO\n",
        "\n",
        "import torch\n",
        "from diffusers import EulerDiscreteScheduler\n",
        "from diffusers import StableDiffusionPipeline\n",
        "from ts.torch_handler.base_handler import BaseHandler\n",
        "\n",
        "logger = logging.getLogger(__name__)\n",
        "model_id = 'runwayml/stable-diffusion-v1-5'\n",
        "\n",
        "\n",
        "class ModelHandler(BaseHandler):\n",
        "\n",
        "  def __init__(self):\n",
        "    self.initialized = False\n",
        "    self.map_location = None\n",
        "    self.device = None\n",
        "    self.use_gpu = True\n",
        "    self.store_avg = True\n",
        "    self.pipe = None\n",
        "\n",
        "  def initialize(self, context):\n",
        "    \"\"\"Initializes the pipe.\"\"\"\n",
        "    properties = context.system_properties\n",
        "    gpu_id = properties.get('gpu_id')\n",
        "\n",
        "    self.map_location, self.device, self.use_gpu = \\\n",
        "      ('cuda', torch.device('cuda:' + str(gpu_id)),\n",
        "       True) if torch.cuda.is_available() else \\\n",
        "        ('cpu', torch.device('cpu'), False)\n",
        "\n",
        "    # Use the Euler scheduler here instead\n",
        "    scheduler = EulerDiscreteScheduler.from_pretrained(model_id,\n",
        "                                                       subfolder='scheduler')\n",
        "    pipe = StableDiffusionPipeline.from_pretrained(model_id,\n",
        "                                                   scheduler=scheduler,\n",
        "                                                   torch_dtype=torch.float16)\n",
        "    pipe = pipe.to('cuda')\n",
        "    # Uncomment the following line to reduce the GPU memory usage.\n",
        "    # pipe.enable_attention_slicing()\n",
        "    self.pipe = pipe\n",
        "\n",
        "    self.initialized = True\n",
        "\n",
        "  def preprocess(self, requests):\n",
        "    \"\"\"Noting to do here.\"\"\"\n",
        "    logger.info('requests: %s', requests)\n",
        "    return requests\n",
        "\n",
        "  def inference(self, preprocessed_data, *args, **kwargs):\n",
        "    \"\"\"Run the inference.\"\"\"\n",
        "    images = []\n",
        "    for pd in preprocessed_data:\n",
        "      prompt = pd['prompt']\n",
        "      images.extend(self.pipe(prompt).images)\n",
        "    return images\n",
        "\n",
        "  def postprocess(self, output_batch):\n",
        "    \"\"\"Converts the images to base64 string.\"\"\"\n",
        "    postprocessed_data = []\n",
        "    for op in output_batch:\n",
        "      fp = BytesIO()\n",
        "      op.save(fp, format='JPEG')\n",
        "      postprocessed_data.append(base64.b64encode(fp.getvalue()).decode('utf-8'))\n",
        "      fp.close()\n",
        "    return postprocessed_data\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ace1dac0af0"
      },
      "source": [
        "After creating the handler file, you must package the handler as a model archiver (MAR) file. The output file must be named 'model.mar'."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "67707f95d440"
      },
      "outputs": [],
      "source": [
        "!torch-model-archiver \\\n",
        "  -f \\\n",
        "  --model-name model \\\n",
        "  --version 1.0 \\\n",
        "  --handler model_artifacts/handler.py \\\n",
        "  --export-path model_artifacts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ffab030f4bc8"
      },
      "source": [
        "### Upload the model artifacts to Cloud Storage\n",
        "\n",
        "Create a new folder in your Cloud Storage bucket to hold the model artifacts"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MzGDU7TWdts_"
      },
      "outputs": [],
      "source": [
        "BUCKET_NAME = \"your-bucket-name-unique\"  # @param {type:\"string\"}\n",
        "BUCKET_URI = f\"gs://{BUCKET_NAME}/\"\n",
        "FULL_GCS_PATH = f\"{BUCKET_URI}model_artifacts\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EcIXiGsCePi"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NIq7R4HZCfIc"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "971232e28657"
      },
      "source": [
        "Next, upload the model archive file and your trained Stable Diffusion 1.5 model to the folder on Cloud Storage."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ef6baf44c808"
      },
      "outputs": [],
      "source": [
        "!gsutil cp -r model_artifacts $BUCKET_URI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "402370ca9396"
      },
      "source": [
        "### Create the Vertex AI model\n",
        "\n",
        "Once you've uploaded the model artifacts into a Cloud Storage bucket, you can create a new Vertex AI model. This notebook uses the [Vertex AI SDK](https://cloud.google.com/vertex-ai/docs/start/use-vertex-ai-python-sdk) to create the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b6c58a74a0fd"
      },
      "outputs": [],
      "source": [
        "PYTORCH_PREDICTION_IMAGE_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai/prediction/pytorch-gpu.1-12:latest\"\n",
        ")\n",
        "APP_NAME = \"my-stable-diffusion\"\n",
        "VERSION = 1\n",
        "MODEL_DISPLAY_NAME = \"stable_diffusion_1_5-unique\"\n",
        "MODEL_DESCRIPTION = \"stable_diffusion_1_5 container\"\n",
        "ENDPOINT_DISPLAY_NAME = f\"{APP_NAME}-endpoint\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "07c3503a1a2e"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_NAME)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FAPoU8Sm5E6e"
      },
      "source": [
        "<div style=\"background:#e3effe; padding:5px; color:#0000aa\">\n",
        "<strong>Note:</strong> The next cell fails if you haven't <a href=\"https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com\"><u>enabled the Vertex API</u></a>.</div>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a776324dd16f"
      },
      "outputs": [],
      "source": [
        "model = aiplatform.Model.upload(\n",
        "    display_name=MODEL_DISPLAY_NAME,\n",
        "    description=MODEL_DESCRIPTION,\n",
        "    serving_container_image_uri=PYTORCH_PREDICTION_IMAGE_URI,\n",
        "    artifact_uri=FULL_GCS_PATH,\n",
        ")\n",
        "\n",
        "model.wait()\n",
        "\n",
        "print(model.display_name)\n",
        "print(model.resource_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0fbc3d371574"
      },
      "source": [
        "### Deploy the model to an endpoint\n",
        "\n",
        "To get online preductions from your Stable Diffusion 2.0 model, you must [deploy it to a Vertex AI endpoint](https://cloud.google.com/vertex-ai/docs/predictions/overview). You can again use the Vertex AI SDK to create the endpoint and deploy your model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ab29f0a770cb"
      },
      "outputs": [],
      "source": [
        "endpoint = aiplatform.Endpoint.create(display_name=ENDPOINT_DISPLAY_NAME)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "25f703df88c7"
      },
      "outputs": [],
      "source": [
        "model.deploy(\n",
        "    endpoint=endpoint,\n",
        "    deployed_model_display_name=MODEL_DISPLAY_NAME,\n",
        "    machine_type=\"n1-standard-8\",\n",
        "    accelerator_type=\"NVIDIA_TESLA_P100\",\n",
        "    accelerator_count=1,\n",
        "    traffic_percentage=100,\n",
        "    deploy_request_timeout=1200,\n",
        "    sync=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c9fc560df0a9"
      },
      "source": [
        "The previous cell, which deploys your model to the endpoint, can take a while to complete. If the previous cell times out before returning, your endpoint might still be successfully deployed to an endpoint. Check the [Cloud Console](https://console.cloud.google.com/vertex-ai/endpoints) to verify the results.\n",
        "\n",
        "You can also extend the time to wait for deployment by changing the `deploy_request_timeout` argument passed to `model.deploy()`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "88a5304dfdc9"
      },
      "source": [
        "## Get online predictions\n",
        "\n",
        "Finally, with your Stable Diffusion 1.5 model deployed to a Vertex AI endpoint, you can now get online predictions from it. Using the Vertex AI SDK, you only need a few lines of code to get an inference."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0d6bc4aa34d6"
      },
      "outputs": [],
      "source": [
        "instances = [{\"prompt\": \"An examplePup dog with a baseball jersey.\"}]\n",
        "response = endpoint.predict(instances=instances)\n",
        "\n",
        "with open(\"img5.jpg\", \"wb\") as g:\n",
        "    g.write(base64.b64decode(response.predictions[0]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "65bafefda60c"
      },
      "outputs": [],
      "source": [
        "display.Image(\"img5.jpg\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TpV-iwP9qw9c"
      },
      "source": [
        "## Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sx_vKniMq9ZX"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Delete endpoint resource\n",
        "endpoint.undeploy_all()\n",
        "endpoint.delete()\n",
        "\n",
        "# Delete model resource\n",
        "model.delete()\n",
        "\n",
        "# Delete Cloud Storage objects that were created\n",
        "delete_bucket = False\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "dreambooth_stablediffusion.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}