notebooks/community/model_garden/model_garden_pytorch_instructpix2pix.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": "2bd716bf3e39"
      },
      "source": [
        "# Vertex AI Model Garden - InstructPix2Pix\n",
        "\n",
        "<table align=\"left\">\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_instructpix2pix.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_instructpix2pix.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8cd12648da4"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates deploying the [InstructPix2Pix](https://huggingface.co/timbrooks/instruct-pix2pix) model on Vertex AI for online prediction.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
        "- Run online predictions for text-guided image editing.\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": [
        "## Run the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2707b02ef5df"
      },
      "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",
        "import requests\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",
        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\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",
        "    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",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\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",
        "# Define common functions.\n",
        "def create_job_name(prefix):\n",
        "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
        "    job_name = f\"{prefix}-{now}\"\n",
        "    return job_name\n",
        "\n",
        "\n",
        "def download_image(url):\n",
        "    response = requests.get(url)\n",
        "    return Image.open(BytesIO(response.content))\n",
        "\n",
        "\n",
        "def image_to_base64(image, format=\"JPEG\"):\n",
        "    buffer = BytesIO()\n",
        "    image.save(buffer, format=format)\n",
        "    image_str = base64.b64encode(buffer.getvalue()).decode(\"utf-8\")\n",
        "    return image_str\n",
        "\n",
        "\n",
        "def base64_to_image(image_str):\n",
        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
        "    return image\n",
        "\n",
        "\n",
        "def image_grid(imgs, rows=2, cols=2):\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h), color=(255, 255, 255))\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, task):\n",
        "    model_name = \"instruct-pix2pix\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"TASK\": task,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\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",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=\"g2-standard-8\",\n",
        "        accelerator_type=\"NVIDIA_L4\",\n",
        "        accelerator_count=1,\n",
        "        deploy_request_timeout=1800,\n",
        "    )\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "b4b46c28d8b1"
      },
      "outputs": [],
      "source": [
        "# @title Upload and deploy model\n",
        "\n",
        "# @markdown This section deploys the InstructPix2Pix model for the text-guided image-to-image task.\n",
        "\n",
        "# @markdown The model deployment step will take ~15 minutes to complete.\n",
        "\n",
        "model, endpoint = deploy_model(\n",
        "    model_id=\"timbrooks/instruct-pix2pix\", task=\"instruct-pix2pix\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6be655247cb1"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send text prompt and image to the endpoint.\n",
        "\n",
        "# @markdown When deployed on one L4 GPU (default), the averaged inference time of a request is ~3 seconds.\n",
        "\n",
        "prompt = \"Add fire to the mountain\"  # @param {type: \"string\"}\n",
        "image = \"https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/mountain.png\"  # @param {type: \"string\"}\n",
        "num_inference_steps = 25  # @param {type:\"number\"}\n",
        "\n",
        "init_image = download_image(image)\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
        "        \"image\": image_to_base64(init_image),\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([init_image, images[0]], rows=1, cols=2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2ccf3714dbe9"
      },
      "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",
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete model.\n",
        "model.delete()\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_instructpix2pix.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}