notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_xl_lcm.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 - Latent Consistent Model\n",
        "\n",
        "<table align=\"left\"><tbody><tr>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_stable_diffusion_xl_lcm.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion_xl_lcm.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      <br>View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates deploying the Stable Diffusion XL [Latent Consistent Model (LCM)](https://huggingface.co/docs/diffusers/main/en/using-diffusers/inference_with_lcm) and [LCM-LoRA](https://huggingface.co/docs/diffusers/main/en/using-diffusers/inference_with_lcm_lora) based on the [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) models 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) with LCM and LCM-LoRA.\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": "b07877158ac3"
      },
      "source": [
        "## Run the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "42fbc614a652"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "import base64\n",
        "import os\n",
        "import sys\n",
        "import uuid\n",
        "from datetime import datetime\n",
        "from io import BytesIO\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "from PIL import Image\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, please change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "\n",
        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "# Set up the default SERVICE_ACCOUNT.\n",
        "SERVICE_ACCOUNT = None\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "# 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/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\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",
        "    \"\"\"Arrange and display multiple images in a grid.\"\"\"\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\n",
        "\n",
        "\n",
        "def deploy_model(model_id, task, refiner_model_id=\"none\"):\n",
        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
        "    model_name = \"stable-diffusion\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-{task}-endpoint\")\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"TASK\": task,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/predictions/diffusers_serving\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\n",
        "    )\n",
        "    machine_type = \"g2-standard-8\"\n",
        "    accelerator_type = \"NVIDIA_L4\"\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=1,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "    )\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "9d0bdfefba5e"
      },
      "outputs": [],
      "source": [
        "# @title Deploy the model to Vertex for online prediction\n",
        "\n",
        "# @markdown This section uploads the model to Model Registry and deploys it to a Vertex AI Endpoint resource. The model deployment step will take ~20 minutes to complete.\n",
        "\n",
        "# @markdown Once deployed, you can send a batch of text prompts to the endpoint to generated images.\n",
        "\n",
        "# @markdown Click \"Show Code\" to see more details.\n",
        "\n",
        "# @markdown **Select one of the two deployment options in the following section.**\n",
        "\n",
        "DEPLOY_MODE = \"SDXL-LCM\"  # @param [\"SDXL-LCM\", \"SDXL-LCM-LORA\"] {isTemplate:true}\n",
        "\n",
        "model_id = \"stabilityai/stable-diffusion-xl-base-1.0\"\n",
        "task = \"text-to-image-sdxl-lcm\"\n",
        "\n",
        "if DEPLOY_MODE == \"SDXL-LCM-LORA\":\n",
        "    task = \"text-to-image-refiner\"\n",
        "\n",
        "model, endpoint = deploy_model(model_id=model_id, task=task)\n",
        "print(\"endpoint_name:\", endpoint.name)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "4ab04da3ec9a"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
        "\n",
        "# @markdown Example:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown Prompt: A serious capybara at work, wearing a suit\n",
        "# @markdown ```\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint.name` allows us to get the endpoint name of\n",
        "#   the endpoint `endpoint` created in the cell above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "\n",
        "# You may uncomment the code below to load an existing endpoint.\n",
        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "# print(\"Using this existing endpoint from a different session: {aip_endpoint_name}\")\n",
        "\n",
        "# @markdown You may adjust the parameters below to achieve best image quality.\n",
        "\n",
        "prompt = \"A serious capybara at work, wearing a suit\"  # @param {type: \"string\"}\n",
        "height = 1024  # @param {type:\"integer\"}\n",
        "width = 1024  # @param {type:\"number\"}\n",
        "num_inference_steps = 8  # @param {type:\"number\"}\n",
        "\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
        "        \"height\": height,\n",
        "        \"width\": width,\n",
        "        \"num_inference_steps\": num_inference_steps,\n",
        "    },\n",
        "]\n",
        "response = endpoint.predict(instances=instances)\n",
        "images = [base64_to_image(image) for image in response.predictions]\n",
        "image_grid(images, rows=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# @title Clean up resources\n",
        "\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
        "\n",
        "try:\n",
        "    # Undeploy model and delete endpoint.\n",
        "    endpoint.delete(force=True)\n",
        "\n",
        "    # Delete model.\n",
        "    model.delete()\n",
        "\n",
        "except Exception as e:\n",
        "    print(e)\n",
        "\n",
        "# Delete bucket.\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_stable_diffusion_xl_lcm.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}