mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
Add Keras stable diffusion notebook for model garden (#1891)
* Add Keras stable diffusion notebook for model garden * Fix minor typos in Keras Stable Diffusion * Fix format after minor typo fixes * Update the objective structure * update minor comments
This commit is contained in:
genquan9
@@ -59,3 +59,4 @@
|
||||
/notebooks/community/model_garden/model_garden_pytorch_blip2.ipynb @xiangxu-google
|
||||
/notebooks/community/model_garden/model_garden_pytorch_detectron2.ipynb @lavraicse
|
||||
/notebooks/community/generative_ai/text_embedding_api_semantic_search_with_scann.ipynb @henrytansetiawan
|
||||
/notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb @genquan9
|
||||
|
||||
@@ -0,0 +1,663 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "ur8xi4C7S06n"
|
||||
},
|
||||
"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": "TirJ-SGQseby"
|
||||
},
|
||||
"source": [
|
||||
"# Vertex AI Model Garden Keras Stable Diffusion\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_keras_stable_diffusion.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",
|
||||
"\n",
|
||||
" <td>\n",
|
||||
" <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_keras_stable_diffusion.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> <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_keras_stable_diffusion.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": "dwGLvtIeECLK"
|
||||
},
|
||||
"source": [
|
||||
"**_NOTE_**: This notebook has been tested in the following environment:\n",
|
||||
"\n",
|
||||
"* Python version = 3.9"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "tvgnzT1CKxrO"
|
||||
},
|
||||
"source": [
|
||||
"## Overview\n",
|
||||
"\n",
|
||||
"This notebook demonstrates how to use [Keras Stable Diffusion](https://keras.io/api/keras_cv/models/stable_diffusion) in Vertex AI Model Garden.\n",
|
||||
"\n",
|
||||
"### Objective\n",
|
||||
"\n",
|
||||
"* Run local inferences for pretrained or customized models\n",
|
||||
"\n",
|
||||
"* Deploy pretrained or customized models in Google Cloud Vertex\n",
|
||||
"\n",
|
||||
"* Finetune models in Google Cloud Vertex\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\n",
|
||||
"pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
|
||||
"pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
|
||||
"Calculator](https://cloud.google.com/products/calculator/)\n",
|
||||
"to generate a cost estimate based on your projected usage."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KEukV6uRk_S3"
|
||||
},
|
||||
"source": [
|
||||
"## Before you begin"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "z__i0w0lCAsW"
|
||||
},
|
||||
"source": [
|
||||
"### Set up notebooks\n",
|
||||
"\n",
|
||||
"- Colab notebook\n",
|
||||
"\n",
|
||||
" You can open this as colab notebook directly.\n",
|
||||
"\n",
|
||||
"- Workbench notebook\n",
|
||||
"\n",
|
||||
" You can open this as workbench notebook with workbench instances. You can create [google managed](https://cloud.google.com/vertex-ai/docs/workbench/managed/create-instance) or [user managed](https://cloud.google.com/vertex-ai/docs/workbench/user-managed/create-new) workbench instances.\n",
|
||||
"\n",
|
||||
"Then, run the following commands to set up notebooks."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "Jvqs-ehKlaYh"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"if \"google.colab\" in str(get_ipython()):\n",
|
||||
" # Configs for colab notebooks.\n",
|
||||
" ! pip3 install --upgrade google-cloud-aiplatform\n",
|
||||
"\n",
|
||||
" # Automatically restart kernel after installs\n",
|
||||
" import IPython\n",
|
||||
"\n",
|
||||
" app = IPython.Application.instance()\n",
|
||||
" app.kernel.do_shutdown(True)\n",
|
||||
"\n",
|
||||
" from google.colab import auth as google_auth\n",
|
||||
"\n",
|
||||
" google_auth.authenticate_user()\n",
|
||||
"\n",
|
||||
"! pip3 install keras-cv==0.4.1"
|
||||
]
|
||||
},
|
||||
{
|
||||
"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",
|
||||
"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",
|
||||
"1. If you are running this notebook locally, you will need to install the [Cloud SDK](https://cloud.google.com/sdk).\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.\n",
|
||||
"\n",
|
||||
"1. Enter your project ID in the cell below. Then run the cell to make sure the\n",
|
||||
"Cloud SDK uses the right project for all the commands in this notebook.\n",
|
||||
"\n",
|
||||
"**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "9wExiMUxFk91"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"\n",
|
||||
"from google.cloud import aiplatform\n",
|
||||
"\n",
|
||||
"# The project and bucket are for experiments below.\n",
|
||||
"PROJECT_ID = \"\" # @param {type:\"string\"}\n",
|
||||
"# The form for BUCKET_URI is gs://<bucket-name>.\n",
|
||||
"BUCKET_URI = \"\" # @param {type:\"string\"}\n",
|
||||
"\n",
|
||||
"REGION = \"us-central1\"\n",
|
||||
"! gcloud config set project $PROJECT_ID\n",
|
||||
"\n",
|
||||
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
|
||||
"EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"keras\")\n",
|
||||
"DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
|
||||
"MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")\n",
|
||||
"\n",
|
||||
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
|
||||
"\n",
|
||||
"# Training constants.\n",
|
||||
"TRAINING_JOB_PREFIX = \"train\"\n",
|
||||
"TRAIN_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/keras-train:latest\"\n",
|
||||
"TRAIN_MACHINE_TYPE = \"a2-highgpu-1g\"\n",
|
||||
"TRAIN_ACCELERATOR_TYPE = \"NVIDIA_TESLA_A100\"\n",
|
||||
"TRAIN_NUM_GPU = 1\n",
|
||||
"RESOLUTION = 512\n",
|
||||
"\n",
|
||||
"# Prediction constants.\n",
|
||||
"PREDICTION_CONTAINER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/keras-serve:latest\"\n",
|
||||
"PREDICTION_ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
|
||||
"PREDICTION_MACHINE_TYPE = \"n1-standard-8\"\n",
|
||||
"DEPLOY_JOB_PREFIX = \"deploy\"\n",
|
||||
"\n",
|
||||
"# The service account for deploying fine tuned model.\n",
|
||||
"# The service account looks like:\n",
|
||||
"# '<account_name>@<project>.iam.gserviceaccount.com'\n",
|
||||
"# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
|
||||
"# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
|
||||
"SERVICE_ACCOUNT = \"\" # @param {type:\"string\"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ZZFPe_GezXg8"
|
||||
},
|
||||
"source": [
|
||||
"### Define common libraries"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "XcYUGwr-AJGY"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import base64\n",
|
||||
"from datetime import datetime\n",
|
||||
"from io import BytesIO\n",
|
||||
"\n",
|
||||
"import matplotlib.pyplot as plt\n",
|
||||
"from PIL import Image\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def get_job_name_with_datetime(prefix: str):\n",
|
||||
" return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def download_data_to_gcs(tar_filepath, gcs_bucket):\n",
|
||||
" filename_with_ext = os.path.basename(tar_filepath)\n",
|
||||
" filename_without_ext = filename_with_ext.replace(\".tar.gz\", \"\")\n",
|
||||
" print(\"Download files from: \", tar_filepath)\n",
|
||||
" ! wget $tar_filepath -O $filename_with_ext\n",
|
||||
" ! mkdir -p $filename_without_ext\n",
|
||||
" ! tar -xvf $filename_with_ext -C .\n",
|
||||
"\n",
|
||||
" ! gsutil -m cp -r $filename_without_ext $gcs_bucket/\n",
|
||||
" gcs_path = os.path.join(gcs_bucket, filename_without_ext)\n",
|
||||
" print(\"Upload files to: \", gcs_path)\n",
|
||||
" return gcs_path\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def deploy_model(model_path):\n",
|
||||
"\n",
|
||||
" deploy_model_name = get_job_name_with_datetime(DEPLOY_JOB_PREFIX)\n",
|
||||
" print(\"The deployed job name is: \", deploy_model_name)\n",
|
||||
" serving_env = {\n",
|
||||
" \"MODEL_PATH\": model_path,\n",
|
||||
" \"IMAGE_WIDTH\": RESOLUTION,\n",
|
||||
" \"IMAGE_HEIGHT\": RESOLUTION,\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" endpoint = aiplatform.Endpoint.create(display_name=f\"{deploy_model_name}-endpoint\")\n",
|
||||
" model = aiplatform.Model.upload(\n",
|
||||
" display_name=deploy_model_name,\n",
|
||||
" serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
|
||||
" serving_container_ports=[8501],\n",
|
||||
" serving_container_predict_route=\"/predict\",\n",
|
||||
" serving_container_health_route=\"/ping\",\n",
|
||||
" serving_container_environment_variables=serving_env,\n",
|
||||
" )\n",
|
||||
" model.deploy(\n",
|
||||
" endpoint=endpoint,\n",
|
||||
" machine_type=PREDICTION_MACHINE_TYPE,\n",
|
||||
" accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
|
||||
" accelerator_count=1,\n",
|
||||
" min_replica_count=1,\n",
|
||||
" max_replica_count=1,\n",
|
||||
" deploy_request_timeout=1800,\n",
|
||||
" service_account=SERVICE_ACCOUNT,\n",
|
||||
" )\n",
|
||||
" return model, endpoint\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 display_image(image):\n",
|
||||
" _ = plt.figure(figsize=(20, 15))\n",
|
||||
" plt.grid(False)\n",
|
||||
" plt.imshow(image)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def display_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"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "epo-RHXzcBBT"
|
||||
},
|
||||
"source": [
|
||||
"## Run inferences\n",
|
||||
"\n",
|
||||
"This section shows how to run inferences with Keras Stable Diffusion models.\n",
|
||||
"\n",
|
||||
"1. Run inferences locally\n",
|
||||
"2. Run inferences with serving dockers\n",
|
||||
"\n",
|
||||
"You can run inferences with pre-trained models from Keras team, or your own finetuned models.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "6zsa9vnBHhvO"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Sets the model_path to empty to load the pre-trained model from Keras team.\n",
|
||||
"# Sets the model_path to a gcs uri to load the finetuned models.\n",
|
||||
"model_path = \"\" # @param {type:\"string\"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ld39hkcIceE2"
|
||||
},
|
||||
"source": [
|
||||
"### Run inferences locally\n",
|
||||
"Local inferences can finish in seconds with GPUs.\n",
|
||||
"\n",
|
||||
"Load models first."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "G1nCKVSac3Y5"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from keras_cv.models import StableDiffusion\n",
|
||||
"\n",
|
||||
"model = StableDiffusion(img_height=RESOLUTION, img_width=RESOLUTION, jit_compile=True)\n",
|
||||
"if model_path:\n",
|
||||
" model.diffusion_model.load_weights(model_path)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ABaCSIWuP-_G"
|
||||
},
|
||||
"source": [
|
||||
"Then run inferences."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "pnyeVsh8RNI5"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"batch_size = 1\n",
|
||||
"img = model.text_to_image(\n",
|
||||
" prompt=\"a flamingo in Picasso style\",\n",
|
||||
" batch_size=batch_size, # How many images to generate at once\n",
|
||||
" num_steps=25, # Number of iterations (controls image quality)\n",
|
||||
" seed=123, # A fixed seed guarantees the same prompt always generates the same image\n",
|
||||
")\n",
|
||||
"for i in range(batch_size):\n",
|
||||
" display_image(img[i])"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "kY87SU9Adq4o"
|
||||
},
|
||||
"source": [
|
||||
"### Serve models with dockers\n",
|
||||
"When serve models with dockers, we will deploy models in Google Cloud. You can start the deployment jobs with CPUs to save costs. The model deployment will take ~10 minutes to finish."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "yCB9vu7RenY6"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"model, endpoint = deploy_model(model_path=model_path)\n",
|
||||
"\n",
|
||||
"endpoint_id = endpoint.name\n",
|
||||
"print(\"endpoint id is: \", endpoint_id)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "72_BW_BgfvYT"
|
||||
},
|
||||
"source": [
|
||||
"Once deployed, you can send a batch of text prompts to the endpoint to generated images.\n",
|
||||
"\n",
|
||||
"Note, the inference time for the first request for a fresh deployment will need more time to process and take ~45 seconds on one V100 GPU. The inferences for further request is ~12 seconds on one V100 GPU per image."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "U_jrNcZ5eVbH"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# # Loads an existing endpoint as below.\n",
|
||||
"# endpoint_id = <An Existing Endpoint ID>\n",
|
||||
"# aip_endpoint_name = (\n",
|
||||
"# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
|
||||
"# )\n",
|
||||
"# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
|
||||
"\n",
|
||||
"instances = [\n",
|
||||
" {\"prompt\": \"a squirrel in Picasso style\"},\n",
|
||||
" {\"prompt\": \"a dog in Picasso style\"},\n",
|
||||
" {\"prompt\": \"a cat in Picasso style\"},\n",
|
||||
" {\"prompt\": \"a deer in Picasso style\"},\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"parameters = {\n",
|
||||
" \"batch_size\": 1, # How many images to generate at once\n",
|
||||
" \"num_steps\": 25, # Number of iterations (controls image quality)\n",
|
||||
" \"seed\": 123, # A fixed seed guarantees the same prompt always generates the same image\n",
|
||||
"}\n",
|
||||
"response = endpoint.predict(instances=instances, parameters=parameters)\n",
|
||||
"# prediction['predicted_image'] will contains the prediction images in a batch.\n",
|
||||
"# The batch size in this example is 1, and the visualization only parses the\n",
|
||||
"# first predicted image.\n",
|
||||
"images = [\n",
|
||||
" base64_to_image(prediction[\"predicted_image\"][0])\n",
|
||||
" for prediction in response.predictions\n",
|
||||
"]\n",
|
||||
"display_image_grid(images, rows=2, cols=2)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "LiQF7fm6f842"
|
||||
},
|
||||
"source": [
|
||||
"### Clean up"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "eqJyypt-f9K6"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Undeploys models and deletes endpoints.\n",
|
||||
"endpoint.delete(force=True)\n",
|
||||
"# Deletes models.\n",
|
||||
"model.delete()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "RB_xY9ipr7ZU"
|
||||
},
|
||||
"source": [
|
||||
"## Finetune models\n",
|
||||
"This section shows how to finetune Keras Stable diffusion models with trainig dockers.\n",
|
||||
"\n",
|
||||
"If you would like to use finetuned models, please go to the section `Run inferences`."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "OD3TtaWs5b4v"
|
||||
},
|
||||
"source": [
|
||||
"### Download data\n",
|
||||
"By default, we use the dataset\n",
|
||||
"[Pokémon BLIP captions](https://huggingface.co/datasets/lambdalabs/pokemon-blip-captions).\n",
|
||||
"However, we'll use a slightly different version which was derived from the original\n",
|
||||
"dataset to fit better with `tf.data`. Refer to\n",
|
||||
"[the documentation](https://huggingface.co/datasets/sayakpaul/pokemon-blip-original-version)\n",
|
||||
"for more details. We download the data to GCS storage for the experiments with training dockers."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "2TVB8MU-5i-q"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Skips this step if you have already downloaded the dataset.\n",
|
||||
"download_data_to_gcs(\n",
|
||||
" \"https://huggingface.co/datasets/sayakpaul/pokemon-blip-original-version/resolve/main/pokemon_dataset.tar.gz\",\n",
|
||||
" DATA_BUCKET,\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Ee7Hzq8O5jgF"
|
||||
},
|
||||
"source": [
|
||||
"### Start training jobs\n",
|
||||
"We finetune models with 10 steps and it takes ~15 minutes to finish using 1 A100 GPU with default settings."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "riG_qUokg0XZ"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"DATA_CSV = os.path.join(DATA_BUCKET, \"pokemon_dataset/data.csv\")\n",
|
||||
"\n",
|
||||
"train_job_name = get_job_name_with_datetime(TRAINING_JOB_PREFIX)\n",
|
||||
"model_dir = os.path.join(MODEL_BUCKET, train_job_name)\n",
|
||||
"worker_pool_specs = [\n",
|
||||
" {\n",
|
||||
" \"machine_spec\": {\n",
|
||||
" \"machine_type\": TRAIN_MACHINE_TYPE,\n",
|
||||
" \"accelerator_type\": TRAIN_ACCELERATOR_TYPE,\n",
|
||||
" \"accelerator_count\": TRAIN_NUM_GPU,\n",
|
||||
" },\n",
|
||||
" \"replica_count\": 1,\n",
|
||||
" \"disk_spec\": {\n",
|
||||
" \"boot_disk_type\": \"pd-ssd\",\n",
|
||||
" \"boot_disk_size_gb\": 500,\n",
|
||||
" },\n",
|
||||
" \"container_spec\": {\n",
|
||||
" \"image_uri\": TRAIN_CONTAINER_URI,\n",
|
||||
" \"command\": [],\n",
|
||||
" \"env\": [\n",
|
||||
" {\n",
|
||||
" \"name\": \"RESOLUTION\",\n",
|
||||
" \"value\": RESOLUTION,\n",
|
||||
" },\n",
|
||||
" ],\n",
|
||||
" \"args\": [\n",
|
||||
" \"--epochs=10\",\n",
|
||||
" f\"--input_csv_path={DATA_CSV}\",\n",
|
||||
" f\"--output_model_dir={model_dir}\",\n",
|
||||
" ],\n",
|
||||
" },\n",
|
||||
" }\n",
|
||||
"]\n",
|
||||
"\n",
|
||||
"train_job = aiplatform.CustomJob(\n",
|
||||
" display_name=train_job_name,\n",
|
||||
" project=PROJECT_ID,\n",
|
||||
" worker_pool_specs=worker_pool_specs,\n",
|
||||
" staging_bucket=STAGING_BUCKET,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"train_job.run()\n",
|
||||
"\n",
|
||||
"model_path = os.path.join(model_dir, \"saved_model.h5\")\n",
|
||||
"print(\"The trained model is saved as: \", model_path)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wBlQ6FQlJhBi"
|
||||
},
|
||||
"source": [
|
||||
"After the training finishes, you can use `model_path` and then go to the `Run inferences` section above to run predictions."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "kkH2nrpdp4sp"
|
||||
},
|
||||
"source": [
|
||||
"### Clean up"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "Ax6vQVZhp9pR"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"train_job.delete()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "1dijQDiZWegt"
|
||||
},
|
||||
"source": [
|
||||
"## References\n",
|
||||
"\n",
|
||||
"- [Fine-tuning Stable Diffusion](https://keras.io/examples/generative/finetune_stable_diffusion/)\n",
|
||||
"- [StableDiffusion image-generation model](https://keras.io/api/keras_cv/models/stable_diffusion/)\n",
|
||||
"- [High-performance image generation using Stable Diffusion in KerasCV](https://keras.io/guides/keras_cv/generate_images_with_stable_diffusion/)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"name": "model_garden_keras_stable_diffusion.ipynb",
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
Reference in New Issue
Block a user