notebooks/community/model_garden/model_garden_jax_stable_diffusion_xl.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - Stable Diffusion XL 1.0 - TPU v5e\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_jax_stable_diffusion_xl.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/model_garden/model_garden_jax_stable_diffusion_xl.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_jax_stable_diffusion_xl.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",
        "    (a Python-3 GPU notebook with preinstalled HuggingFace/transformer libraries is recommended)\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates how to deploy the [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) model on Vertex AI for online prediction.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Deploy the model to a [Vertex AI Endpoint resource](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
        "- Run online predictions for text-to-image.\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 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": "264c07757582"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "**NOTE**: \n",
        "\n",
        "*  Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.\n",
        "*  This Notebook demonstrate how to deploy the model [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) on Vertex AI prediction endpoint with a TPU v5e instance (machine type of `ct5lp-hightpu-1t`). Please ensure you have enough resource quota in region `us-west1`. If not, please follow the [instructions](https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#securing_capacity) to get quota."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ioensNKM8ned"
      },
      "source": [
        "### Setup notebook"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d73ffa0c0b83"
      },
      "source": [
        "#### Colab\n",
        "Run the following commands for Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2707b02ef5df"
      },
      "outputs": [],
      "source": [
        "if \"google.colab\" in str(get_ipython()):\n",
        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()\n",
        "\n",
        "# Restart the notebook kernel after installs.\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bb7adab99e41"
      },
      "source": [
        "### Setup Google Cloud project\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",
        "\n",
        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\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 models to Vertex AI endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6c460088b873"
      },
      "source": [
        "Set following variables for experiments environment:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "855d6b96f291"
      },
      "outputs": [],
      "source": [
        "# Cloud project id.\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The region you want to launch jobs in.\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The Cloud Storage bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
        "GCS_BUCKET = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The service account for deploying fine tuned model.\n",
        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e828eb320337"
      },
      "source": [
        "Initialize Vertex AI API:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "12cd25839741"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2cc825514deb"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "b42bd4fa2b2d"
      },
      "outputs": [],
      "source": [
        "# The pre-built serving docker image. It contains serving scripts and models.\n",
        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-diffusers-serve-tpu:20240110_1526_RC00\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c250872074f"
      },
      "source": [
        "### Define common functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "354da31189dc"
      },
      "outputs": [],
      "source": [
        "import base64\n",
        "from io import BytesIO\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "from PIL import Image\n",
        "\n",
        "\n",
        "def base64_to_image(image_str):\n",
        "    \"\"\"Convert base64 encoded string to an image.\"\"\"\n",
        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
        "    return image\n",
        "\n",
        "\n",
        "def image_grid(imgs, rows=2, cols=2):\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new(\n",
        "        mode=\"RGB\", size=(cols * w + 10 * cols, rows * h), color=(255, 255, 255)\n",
        "    )\n",
        "    for i, img in enumerate(imgs):\n",
        "        grid.paste(img, box=(i % cols * w + 10 * i, i // cols * h))\n",
        "    return grid\n",
        "\n",
        "\n",
        "def deploy_model(model_id):\n",
        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
        "    model_name = model_id + \"-tpu\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_predict_route=\"/predict\",\n",
        "        serving_container_health_route=\"/health\",\n",
        "    )\n",
        "    machine_type = \"ct5lp-hightpu-1t\"\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        enable_access_logging=True,\n",
        "        min_replica_count=1,\n",
        "        sync=True,\n",
        "    )\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bf7f82732e61"
      },
      "source": [
        "## Upload and Deploy models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1cc26e68d7b0"
      },
      "source": [
        "This section uploads the model to Model Registry and deploys it to a Vertex AI Endpoint resource.\n",
        "\n",
        "The model deployment step will take ~30 minutes to complete."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cd7b56421392"
      },
      "source": [
        "### Text-to-image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6d331b1ea337"
      },
      "source": [
        "Deploy the stable diffusion xl model for the text-to-image task.\n",
        "\n",
        "Once deployed, you can send a batch of text prompts to the endpoint to generated images.\n",
        "\n",
        "When deployed on one TPU V5e instance, the averaged inference time of one image is ~3 seconds."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bf55e38815dc"
      },
      "outputs": [],
      "source": [
        "# Set the model_id to \"stabilityai/stable-diffusion-xl-base-1.0\" to load the OSS pre-trained model.\n",
        "model, endpoint = deploy_model(\n",
        "    model_id=\"stabilityai/stable-diffusion-xl-base-1.0\",\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4ab04da3ec9a"
      },
      "outputs": [],
      "source": [
        "instances = [\n",
        "    {\n",
        "        \"prompt\": \"Photorealistic whale swimming in abyss\",\n",
        "        \"height\": 1024,\n",
        "        \"width\": 1024,\n",
        "    },\n",
        "    {\n",
        "        \"prompt\": \"Photorealistic happy dog running\",\n",
        "        \"height\": 1024,\n",
        "        \"width\": 1024,\n",
        "    },\n",
        "]\n",
        "response = endpoint.predict(instances=instances)\n",
        "\n",
        "images = [\n",
        "    base64_to_image(prediction.get(\"images\")[0]) for prediction in response.predictions\n",
        "]\n",
        "image_grid(images, rows=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af21a3cff1e0"
      },
      "source": [
        "### Clean up resources:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "model_garden_jax_stable_diffusion_xl.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}