part_stable_diffusion_2_1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "7d9bbf86da5e"
   },
   "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": "99c1c3fc2ca5"
   },
   "source": [
    "# Vertex AI Model Garden - Stable Diffusion V2.1\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_pytorch_stable_diffusion_2_1.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_pytorch_stable_diffusion_2_1.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_pytorch_stable_diffusion_2_1.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 running local inference for [stabilityai/stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) on either [Colab](https://colab.research.google.com) or [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench). This notebook also demonstrates finetuning stabilityai/stable-diffusion-2-1 with [Dreambooth](https://huggingface.co/docs/diffusers/training/dreambooth) and deploying it on Vertex AI for online prediction.\n",
    "\n",
    "### Objective\n",
    "\n",
    "- Run local predictions for text-to-image and text-guided-image-to-image with serving dockers.\n",
    "- Finetune the stabilityai/stable-diffusion-2-1 model with [Dreambooth](https://huggingface.co/docs/diffusers/training/dreambooth).\n",
    "- Upload the model to [Vertex AI Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\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 and text-guided-image-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",
    "tags": []
   },
   "source": [
    "## Before you begin\n",
    "\n",
    "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ioensNKM8ned"
   },
   "source": [
    "### Setup notebook"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fb671e75ca7b"
   },
   "source": [
    "#### Workbench only\n",
    "1.   Follow [this link](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_pytorch_stable_diffusion_2_1.ipynb) to deploy the notebook to a Vertex AI Workbench Instance.\n",
    "2.   Select `Create a new Notebook`.\n",
    "3.   Click `Advanced Options`.\n",
    "4.   In the **Environment** tab, select `Debian 10` for **Operating System** and select `Custom Container` for **Environment**.\n",
    "5.   Set `Docker container image` to `us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/transformers-notebook`.\n",
    "6.   Under **Machine configuration**, select 1 `T4` GPU and select `Install NVIDIA GPU driver automatically for me`.\n",
    "7.   Click `Create` to create the Vertex AI Workbench instance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bb7adab99e41",
    "tags": []
   },
   "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 fine tuned model to Vertex AI endpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6c460088b873"
   },
   "source": [
    "Fill following variables for experiments environment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "855d6b96f291",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Creating gs://my-project-0004-346516-pytorch112kagglewbi/...\n",
      "\u001b[33mWARNING: Ignoring invalid distribution -ryptography (/opt/conda/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Ignoring invalid distribution -ryptography (/opt/conda/lib/python3.10/site-packages)\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.1.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import socket\n",
    "import re\n",
    "\n",
    "UNIQUE_PREFIX = socket.gethostname()\n",
    "UNIQUE_PREFIX = re.sub('[^A-Za-z0-9]+', '', UNIQUE_PREFIX)\n",
    "\n",
    "# Cloud project id.\n",
    "PROJECT_IDS = !(gcloud config get-value core/project)\n",
    "PROJECT_ID = PROJECT_IDS[0]  # @param {type:\"string\"}\n",
    "\n",
    "# The region you want to launch jobs in.\n",
    "PREFIX_NUM_ONLY=int(str(re.search(r'\\d+', UNIQUE_PREFIX).group()))\n",
    "REGION_ALLOCATE=PREFIX_NUM_ONLY%3\n",
    "if REGION_ALLOCATE == 0:\n",
    "    REGION = \"asia-southeast1\"\n",
    "elif REGION_ALLOCATE == 1:\n",
    "    REGION = \"us-central1\"\n",
    "else:\n",
    "    REGION = \"europe-west4\"\n",
    "\n",
    "# The Cloud Storage bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
    "GCS_BUCKET = f\"{PROJECT_ID}-{UNIQUE_PREFIX}\"  # @param {type:\"string\"} \n",
    "BUCKET_URI = f\"gs://{GCS_BUCKET}\"  # @param {type:\"string\"}\n",
    "\n",
    "# The service account for deploying fine tuned model.\n",
    "SERVICE_ACCOUNT = !(gcloud config get-value core/account)  # @param {type:\"string\"}\n",
    "SERVICE_ACCOUNT = SERVICE_ACCOUNT[0]\n",
    "\n",
    "! gcloud storage buckets create {BUCKET_URI} --project={PROJECT_ID} --location={REGION}\n",
    "! pip install -q gdown"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e828eb320337"
   },
   "source": [
    "Initialize Vertex-AI API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "12cd25839741",
    "scrolled": true,
    "tags": []
   },
   "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",
    "tags": []
   },
   "source": [
    "### Define constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "b42bd4fa2b2d",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The pre-built training docker image. It contains training scripts and models.\n",
    "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-train:20240401_0802_RC00\"\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:20240331_0836_RC00\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0c250872074f",
    "tags": []
   },
   "source": [
    "### Define common functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "354da31189dc",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import base64\n",
    "import glob\n",
    "import os\n",
    "from datetime import datetime\n",
    "from io import BytesIO\n",
    "\n",
    "import requests\n",
    "from google.cloud import aiplatform, storage\n",
    "from PIL import Image\n",
    "\n",
    "def create_job_name(prefix):\n",
    "    user = os.environ.get(\"USER\")\n",
    "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "    job_name = f\"{prefix}-{user}-{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 base64_to_jpg(image_str):\n",
    "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
    "    with BytesIO() as output_buffer:\n",
    "        jpg_data = output_buffer.getvalue()\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))\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):\n",
    "    model_name = model_id\n",
    "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-{task}-endpoint\")\n",
    "    serving_env = {\n",
    "        \"MODEL_ID\": model_id,\n",
    "        \"TASK\": task,\n",
    "    }\n",
    "    model = aiplatform.Model.upload(\n",
    "        display_name=f\"{model_name}-{task}-model\",\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",
    "        service_account=SERVICE_ACCOUNT,\n",
    "    )\n",
    "    return model, endpoint\n",
    "\n",
    "\n",
    "def get_bucket_and_blob_name(filepath):\n",
    "    # The gcs path is of the form gs://<bucket-name>/<blob-name>\n",
    "    gs_suffix = filepath.split(\"gs://\", 1)[1]\n",
    "    return tuple(gs_suffix.split(\"/\", 1))\n",
    "\n",
    "\n",
    "def upload_local_dir_to_gcs(local_dir_path, gcs_dir_path):\n",
    "    \"\"\"Uploads files in a local directory to a GCS directory.\"\"\"\n",
    "    client = storage.Client()\n",
    "    bucket_name = gcs_dir_path.split(\"/\")[2]\n",
    "    bucket = client.get_bucket(bucket_name)\n",
    "    for local_file in glob.glob(local_dir_path + \"/**\"):\n",
    "        if not os.path.isfile(local_file):\n",
    "            continue\n",
    "        filename = local_file[1 + len(local_dir_path) :]\n",
    "        gcs_file_path = os.path.join(gcs_dir_path, filename)\n",
    "        _, blob_name = get_bucket_and_blob_name(gcs_file_path)\n",
    "        blob = bucket.blob(blob_name)\n",
    "        blob.upload_from_filename(local_file)\n",
    "        print(\"Copied {} to {}.\".format(local_file, gcs_file_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e70e3519ff8b",
    "tags": []
   },
   "source": [
    "## Finetune with Dreambooth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0dc65d8f0689"
   },
   "source": [
    "This section uses [dreambooth](https://dreambooth.github.io/) to finetune the [stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) model with [5 dog images](https://drive.google.com/drive/folders/1BO_dyz-p65qhBRRMRA4TbZ8qW4rB99JZ) to personalize the text-to-image model.\n",
    "\n",
    "It finetunes both text encoder and unet of the stable diffusion model up to 800 steps. The whole finetuning job takes 45 minutes to finish using 1 A100 GPU.\n",
    "\n",
    "The full model will be saved after the finetuning job finishs and it can be loaded by the [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) to run inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "34048707df5c"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Retrieving folder contents\n",
      "Processing file 1aIWXIj3iB-7CfRFLROHq9U09EdrpUH-T alvan-nee-9M0tSjb-cpA-unsplash.jpeg\n",
      "Processing file 15IDLUW4Jpwot_xVpWWVyCONWpiBBpQIi alvan-nee-bQaAJCbNq3g-unsplash.jpeg\n",
      "Processing file 1FYFAp1sFfYh7I2eDcLK_BP7diKgGQib_ alvan-nee-brFsZ7qszSY-unsplash.jpeg\n",
      "Processing file 1zonwTmg7Ggl_w2ELQyTDVNkjtQTlsDq4 alvan-nee-eoqnr8ikwFE-unsplash.jpeg\n",
      "Processing file 1vI0QwqfGyWNdloy0L9N4Gz4qN3_BduwL alvan-nee-Id1DBHv4fbg-unsplash.jpeg\n",
      "Retrieving folder contents completed\n",
      "Building directory structure\n",
      "Building directory structure completed\n",
      "Downloading...\n",
      "From: https://drive.google.com/uc?id=1aIWXIj3iB-7CfRFLROHq9U09EdrpUH-T\n",
      "To: /home/jupyter/GenAI5/Course/dog/alvan-nee-9M0tSjb-cpA-unsplash.jpeg\n",
      "100%|█████████████████████████████████████████| 677k/677k [00:00<00:00, 114MB/s]\n",
      "Downloading...\n",
      "From: https://drive.google.com/uc?id=15IDLUW4Jpwot_xVpWWVyCONWpiBBpQIi\n",
      "To: /home/jupyter/GenAI5/Course/dog/alvan-nee-bQaAJCbNq3g-unsplash.jpeg\n",
      "100%|███████████████████████████████████████| 1.40M/1.40M [00:00<00:00, 164MB/s]\n",
      "Downloading...\n",
      "From: https://drive.google.com/uc?id=1FYFAp1sFfYh7I2eDcLK_BP7diKgGQib_\n",
      "To: /home/jupyter/GenAI5/Course/dog/alvan-nee-brFsZ7qszSY-unsplash.jpeg\n",
      "100%|██████████████████████████████████████| 1.19M/1.19M [00:00<00:00, 36.9MB/s]\n",
      "Downloading...\n",
      "From: https://drive.google.com/uc?id=1zonwTmg7Ggl_w2ELQyTDVNkjtQTlsDq4\n",
      "To: /home/jupyter/GenAI5/Course/dog/alvan-nee-eoqnr8ikwFE-unsplash.jpeg\n",
      "100%|███████████████████████████████████████| 1.17M/1.17M [00:00<00:00, 114MB/s]\n",
      "Downloading...\n",
      "From: https://drive.google.com/uc?id=1vI0QwqfGyWNdloy0L9N4Gz4qN3_BduwL\n",
      "To: /home/jupyter/GenAI5/Course/dog/alvan-nee-Id1DBHv4fbg-unsplash.jpeg\n",
      "100%|███████████████████████████████████████| 1.16M/1.16M [00:00<00:00, 154MB/s]\n",
      "Download completed\n",
      "Copied dog/alvan-nee-eoqnr8ikwFE-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog/alvan-nee-eoqnr8ikwFE-unsplash.jpeg.\n",
      "Copied dog/alvan-nee-bQaAJCbNq3g-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog/alvan-nee-bQaAJCbNq3g-unsplash.jpeg.\n",
      "Copied dog/alvan-nee-9M0tSjb-cpA-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog/alvan-nee-9M0tSjb-cpA-unsplash.jpeg.\n",
      "Copied dog/alvan-nee-brFsZ7qszSY-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog/alvan-nee-brFsZ7qszSY-unsplash.jpeg.\n",
      "Copied dog/alvan-nee-Id1DBHv4fbg-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog/alvan-nee-Id1DBHv4fbg-unsplash.jpeg.\n",
      "Copied dog/alvan-nee-eoqnr8ikwFE-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog_class/alvan-nee-eoqnr8ikwFE-unsplash.jpeg.\n",
      "Copied dog/alvan-nee-bQaAJCbNq3g-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog_class/alvan-nee-bQaAJCbNq3g-unsplash.jpeg.\n",
      "Copied dog/alvan-nee-9M0tSjb-cpA-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog_class/alvan-nee-9M0tSjb-cpA-unsplash.jpeg.\n",
      "Copied dog/alvan-nee-brFsZ7qszSY-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog_class/alvan-nee-brFsZ7qszSY-unsplash.jpeg.\n",
      "Copied dog/alvan-nee-Id1DBHv4fbg-unsplash.jpeg to gs://my-project-0004-346516-pytorch112kagglewbi/dreambooth/dog_class/alvan-nee-Id1DBHv4fbg-unsplash.jpeg.\n"
     ]
    }
   ],
   "source": [
    "# Download example training images.\n",
    "!gdown --folder https://drive.google.com/drive/folders/1BO_dyz-p65qhBRRMRA4TbZ8qW4rB99JZ\n",
    "\n",
    "# Upload data to Cloud Storage bucket.\n",
    "upload_local_dir_to_gcs(\"dog\", f\"gs://{GCS_BUCKET}/dreambooth/dog\")\n",
    "upload_local_dir_to_gcs(\"dog\", f\"gs://{GCS_BUCKET}/dreambooth/dog_class\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "969cfeb79317"
   },
   "source": [
    "**NOTE**: If the upload step fails due to lacking of permission, you need to [grant the Storage Object Admin role](https://cloud.google.com/storage/docs/access-control/using-iam-permissions) for the Cloud account of the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "65467b361315",
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Output directory:\n",
      "gs://my-project-0004-346516-pytorch112kagglewbi/aiplatform-custom-training-2024-08-15-03:25:57.529 \n",
      "View Training:\n",
      "https://console.cloud.google.com/ai/platform/locations/us-central1/training/2420374084597579776?project=255766800726\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "View backing custom job:\n",
      "https://console.cloud.google.com/ai/platform/locations/us-central1/training/3077763290751827968?project=255766800726\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "Error: Training failed with:\n",
      "code: 13\n",
      "message: \"Internal error occurred. Please retry in a few minutes. If you still experience errors, contact Cloud AI Platform.\"\n",
      "\n",
      "Deleting CustomContainerTrainingJob : projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776\n",
      "CustomContainerTrainingJob deleted. . Resource name: projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776\n",
      "Deleting CustomContainerTrainingJob resource: projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776\n",
      "Delete CustomContainerTrainingJob backing LRO: projects/255766800726/locations/us-central1/operations/5886165302353330176\n",
      "CustomContainerTrainingJob resource projects/255766800726/locations/us-central1/trainingPipelines/2420374084597579776 deleted.\n",
      "Error may becaused due to machine unavailable - A100s - trying L4s \n",
      "Training Output directory:\n",
      "gs://my-project-0004-346516-pytorch112kagglewbi/aiplatform-custom-training-2024-08-15-05:08:27.421 \n",
      "View Training:\n",
      "https://console.cloud.google.com/ai/platform/locations/us-central1/training/8295935525013946368?project=255766800726\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/8295935525013946368 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "View backing custom job:\n",
      "https://console.cloud.google.com/ai/platform/locations/us-central1/training/7240316401418764288?project=255766800726\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/8295935525013946368 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/8295935525013946368 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/8295935525013946368 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/8295935525013946368 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n",
      "CustomContainerTrainingJob projects/255766800726/locations/us-central1/trainingPipelines/8295935525013946368 current state:\n",
      "PipelineState.PIPELINE_STATE_RUNNING\n"
     ]
    }
   ],
   "source": [
    "# The pre-trained model to be loaded.\n",
    "model_id = \"stabilityai/stable-diffusion-2-1\"\n",
    "\n",
    "# Input and output path.\n",
    "instance_dir = f\"/gcs/{GCS_BUCKET}/dreambooth/dog\"\n",
    "class_dir = f\"/gcs/{GCS_BUCKET}/dreambooth/dog_class\"\n",
    "output_dir = f\"/gcs/{GCS_BUCKET}/dreambooth/output\"\n",
    "\n",
    "# # Worker pool spec.\n",
    "# machine_type = \"a2-highgpu-1g\"\n",
    "# num_nodes = 1\n",
    "# gpu_type = \"NVIDIA_TESLA_A100\"\n",
    "# num_gpus = 1\n",
    "\n",
    "# Pass training arguments and launch job.\n",
    "# See https://github.com/huggingface/diffusers/blob/v0.14.0/examples/dreambooth/train_dreambooth.py#L75\n",
    "# for a full list of training arguments.\n",
    "\n",
    "\n",
    "def train_model_with_config(num_nodes, machine_type, gpu_type, num_gpus ):\n",
    "\n",
    "    # Setup training job.\n",
    "    job_name = create_job_name(\"dreambooth-stable-diffusion\")\n",
    "    job = aiplatform.CustomContainerTrainingJob(\n",
    "        display_name=job_name,\n",
    "        container_uri=TRAIN_DOCKER_URI,\n",
    "    )\n",
    "    try:\n",
    "        model = job.run(\n",
    "        args=[\n",
    "            \"dreambooth/train_dreambooth.py\",\n",
    "            f\"--pretrained_model_name_or_path={model_id}\",\n",
    "            \"--train_text_encoder\",\n",
    "            f\"--instance_data_dir={instance_dir}\",\n",
    "            f\"--class_data_dir={class_dir}\",\n",
    "            f\"--output_dir={output_dir}\",\n",
    "            \"--mixed_precision=fp16\",\n",
    "            \"--with_prior_preservation\",\n",
    "            \"--prior_loss_weight=1.0\",\n",
    "            \"--instance_prompt='a photo of cute dog'\",\n",
    "            \"--class_prompt='a photo of dog'\",\n",
    "            \"--resolution=768\",\n",
    "            \"--train_batch_size=1\",\n",
    "            \"--sample_batch_size=1\",\n",
    "            \"--gradient_accumulation_steps=1\",\n",
    "            \"--gradient_checkpointing\",\n",
    "            \"--learning_rate=5e-6\",\n",
    "            \"--lr_scheduler=constant\",\n",
    "            \"--lr_warmup_steps=0\",\n",
    "            \"--num_class_images=200\",\n",
    "            \"--max_train_steps=200\",\n",
    "            \"--checkpointing_steps=100\",\n",
    "        ],\n",
    "        replica_count=num_nodes,\n",
    "        machine_type=machine_type,\n",
    "        accelerator_type=gpu_type,\n",
    "        accelerator_count=num_gpus,\n",
    "        )\n",
    "    except Exception as e:\n",
    "        print(f\"Error: {e}\")\n",
    "        job.delete()\n",
    "    return(model)\n",
    "\n",
    "\n",
    "\n",
    "try:\n",
    "    # Code that might potentially cause an error\n",
    "    machine_type = 'g2-standard-96' #\"a2-highgpu-1g\"\n",
    "    num_nodes = 1\n",
    "    gpu_type = \"NVIDIA_L4\"\n",
    "    num_gpus = 8\n",
    "    model = train_model_with_config(num_nodes, machine_type, gpu_type, num_gpus)\n",
    "except :  # Replace 'ErrorType' with the specific error you want to catch \n",
    "    try : \n",
    "        print(\"Error may becaused due to machine unavailable - smaller - trying L4s \")\n",
    "        machine_type = \"g2-standard-96\"\n",
    "        num_nodes = 1\n",
    "        gpu_type = \"NVIDIA_L4\"\n",
    "        num_gpus = 8\n",
    "        model = train_model_with_config(num_nodes, machine_type, gpu_type, num_gpus)\n",
    "    except :\n",
    "        print(\"Error may becaused due to machine unavailable - L4s also - trying A100s again \")\n",
    "        machine_type = \"a2-highgpu-1g\"\n",
    "        num_nodes = 1\n",
    "        gpu_type = \"NVIDIA_TESLA_A100\"\n",
    "        num_gpus = 1\n",
    "        model = train_model_with_config(num_nodes, machine_type, gpu_type, num_gpus)\n",
    "else:\n",
    "    # Code to execute if there's no error in the 'try' block\n",
    "    print(\"Error may becaused due to machine unavailable of any type \")\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bf7f82732e61",
    "tags": []
   },
   "source": [
    "## Upload and Deploy models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1cc26e68d7b0"
   },
   "source": [
    "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
    "\n",
    "The model deployment step will take ~15 minutes to complete. \n",
    "\n",
    "The first request requires some additional time for model compilation (up to a\n",
    "few minutes), but the future requests should be processed much faster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cd7b56421392"
   },
   "source": [
    "### Text-to-image Deployment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6d331b1ea337"
   },
   "source": [
    "Deploy the stable diffusion 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 `NVIDIA_L4` GPU, the averaged inference time of a request is ~3-4 seconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bf55e38815dc",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "# Set the model_id to \"stabilityai/stable-diffusion-2-1\" to load the OSS pre-trained model.\n",
    "# text_to_image_model, text_to_image_endpoint = deploy_model(\n",
    "#     model_id=f\"gs://{GCS_BUCKET}/dreambooth/output\", task=\"text-to-image\"\n",
    "# )\n",
    "\n",
    "text_to_image_model, text_to_image_endpoint = deploy_model(\n",
    "    model_id=\"stabilityai/stable-diffusion-2-1\", task=\"text-to-image\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "80b3fd2ace09"
   },
   "source": [
    "NOTE: The model weights will be downloaded after the deployment succeeds. Thus upto 10 minutes of additional waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c1e51f764a60"
   },
   "source": [
    "### Text-guided image-to-image Deployment\n",
    "Deploy the stable diffusion model for the text-guided image-to-image task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "65e32356fbd1",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "# Set the model_id to a GCS path, like \"gs://GCS_BUCKET/dreambooth/output\", to load the dreambooth finetuned model above.\n",
    "image_to_image_model, image_to_image_endpoint = deploy_model(\n",
    "    model_id=\"stabilityai/stable-diffusion-2-1\", task=\"image-to-image\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Images with Stable Diffusion\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Generate an Image from a prompt using text to image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enter a prompt to generate an image using Stable Diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4ab04da3ec9a",
    "tags": []
   },
   "outputs": [],
   "source": [
    "prompt = \"A peaceful park\"\n",
    "\n",
    "instances = [\n",
    "    {\"prompt\": prompt}\n",
    "]\n",
    "response = text_to_image_endpoint.predict(instances=instances)\n",
    "generated_image = [base64_to_image(image) for image in response.predictions]\n",
    "display(generated_image[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate an image from a reference image and prompt using image to image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enter a new prompt on how you want to edit the image generated above. \\\n",
    "The strength refers to the amount of deviation from the original image in a range from 0.0 to 1.0 with 0 being no change"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "83a50fd4a1ed",
    "tags": []
   },
   "outputs": [],
   "source": [
    "prompt = \"A peaceful park in a van gogh style\"\n",
    "strength = 0.4\n",
    "\n",
    "init_image = image_to_base64(generated_image[0])\n",
    "\n",
    "instances = [\n",
    "    {\n",
    "        \"prompt\": prompt,\n",
    "        \"image\": init_image,\n",
    "        \"strength\": strength\n",
    "    },\n",
    "]\n",
    "\n",
    "response = image_to_image_endpoint.predict(instances=instances)\n",
    "image = [base64_to_image(image) for image in response.predictions]\n",
    "display(image[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Cleanup"
   ]
  },
  {
   "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",
    "#text_to_image_endpoint.delete(force=True)\n",
    "\n",
    "# Delete models.\n",
    "#text_to_image_model.delete()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ed3795d474b9"
   },
   "source": [
    "Clean up resources:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "b53b883257b4",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Undeploy model and delete endpoint.\n",
    "#image_to_image_endpoint.delete(force=True)\n",
    "\n",
    "# Delete models.\n",
    "#image_to_image_model.delete()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Inline Model Testing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "e8a42fa49305",
    "tags": []
   },
   "source": [
    "## Run inferences locally"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1d5ebc91c786",
    "tags": []
   },
   "source": [
    "### Text-to-image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d39ed8c97cc5",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import torch\n",
    "# from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline\n",
    "# from diffusers.models.attention_processor import AttnProcessor2_0\n",
    "\n",
    "# device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "# model_id = \"stabilityai/stable-diffusion-2-1\"\n",
    "# pipe = StableDiffusionPipeline.from_pretrained(\n",
    "#     model_id, torch_dtype=torch.float16, use_safetensors=True\n",
    "# )\n",
    "# pipe.to(device)\n",
    "\n",
    "# pipe.scheduler = DPMSolverMultistepScheduler.from_config(\n",
    "#     pipe.scheduler.config, algorithm_type=\"dpmsolver++\"\n",
    "# )\n",
    "# pipe.unet.set_attn_processor(AttnProcessor2_0())\n",
    "# pipe.unet.to(memory_format=torch.channels_last)\n",
    "\n",
    "\n",
    "# prompt = \"A DSLR photo of Mt. Fuji with cherry blossom in the background\"\n",
    "\n",
    "# results = pipe(prompt=prompt, height=768, width=768, num_inference_steps=25)\n",
    "# images = results.images\n",
    "# display(images[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5aed5ed7b6f6",
    "tags": []
   },
   "source": [
    "### Text-guided image-to-image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "faadabb7728f",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# import torch\n",
    "# from diffusers import (DPMSolverMultistepScheduler,\n",
    "#                        StableDiffusionImg2ImgPipeline)\n",
    "# from diffusers.models.attention_processor import AttnProcessor2_0\n",
    "\n",
    "# device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "# model_id = \"stabilityai/stable-diffusion-2-1\"\n",
    "# pipe = StableDiffusionImg2ImgPipeline.from_pretrained(\n",
    "#     model_id, torch_dtype=torch.float16\n",
    "# )\n",
    "# pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)\n",
    "# pipe = pipe.to(device)\n",
    "\n",
    "# pipe.scheduler = DPMSolverMultistepScheduler.from_config(\n",
    "#     pipe.scheduler.config, algorithm_type=\"dpmsolver++\"\n",
    "# )\n",
    "# pipe.unet.set_attn_processor(AttnProcessor2_0())\n",
    "# pipe.unet.to(memory_format=torch.channels_last)\n",
    "\n",
    "# # url = \"https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg\"\n",
    "\n",
    "# # response = requests.get(url)\n",
    "# # init_image = Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
    "# # init_image = init_image.resize((768, 512))\n",
    "\n",
    "# init_image = image_to_base64(generated_image[0])\n",
    "# init_image = base64_to_jpg(init_image)\n",
    "\n",
    "# prompt = \"A peaceful park in a van gogh style\"\n",
    "\n",
    "# results = pipe(prompt=prompt, image=init_image, strength=0.4)\n",
    "# images = results.images\n",
    "# display(init_image)\n",
    "# display(images[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "model_garden_pytorch_stable_diffusion_2_1.ipynb",
   "toc_visible": false
  },
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m113",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m113"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-root-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}