Dedicated Endpoint Support for SDXL Dreambooth LoRA Finetuning

PiperOrigin-RevId: 770624167

Author: vertex-mg-bot

notebooks/community/model_garden/model_garden_pytorch_sd_xl_finetuning_dreambooth_lora.ipynb

@@ -63,6 +63,11 @@
         "- Deploy the model to a [Vertex AI Endpoint resource](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
         "- Run online predictions for text-to-image.\n",
         "\n",
+        "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -122,13 +127,14 @@
         "from google.cloud import aiplatform, storage\n",
         "from huggingface_hub import snapshot_download\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
-        "models, endpoints = {}, {}\n",
         "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
@@ -188,7 +194,14 @@
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
         "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\""
+        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
       ]
     },
     {
@@ -200,7 +213,7 @@
       },
       "outputs": [],
       "source": [
-        "# @title Start Dreambooth LoRA finetune\n",
+        "# @title Set up the Dreambooth LoRA finetune parameters\n",
         "\n",
         "# @markdown This section uses [Dreambooth LoRA](https://dreambooth.github.io/) to finetune\n",
         "# @markdown the [stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) model\n",
@@ -421,17 +434,19 @@
         "\n",
         "\n",
         "def deploy_model(\n",
-        "    model_id,\n",
-        "    lora_id,\n",
-        "    task,\n",
-        "    accelerator_type,\n",
-        "    machine_type,\n",
-        "    accelerator_count,\n",
+        "    model_id: str,\n",
+        "    lora_id: str,\n",
+        "    task: str,\n",
+        "    accelerator_type: str = \"g2-standard-8\",\n",
+        "    machine_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
         "):\n",
         "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
         "    model_name = model_id\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
-        "        display_name=common_util.get_job_name_with_datetime(model_name)\n",
+        "        display_name=common_util.get_job_name_with_datetime(model_name),\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
         "    )\n",
         "    serving_env = {\n",
         "        \"MODEL_ID\": model_id,\n",
@@ -466,30 +481,21 @@
         "    return model, endpoint\n",
         "\n",
         "\n",
+        "LABEL = \"sd_xl\"\n",
+        "\n",
         "# Set the model_id to \"stabilityai/stable-diffusion-xl-base-1.0\" to load the OSS pre-trained model.\n",
-        "models[\"sd_xl\"], endpoints[\"sd_xl\"] = deploy_model(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
         "    model_id=model_id,\n",
         "    lora_id=lora_id,\n",
         "    task=\"text-to-image-sdxl\",\n",
         "    accelerator_type=serve_accelerator_type,\n",
         "    machine_type=serve_machine_type,\n",
         "    accelerator_count=serve_accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")\n",
-        "print(\"endpoint_name:\", endpoints[\"sd_xl\"].name)\n",
         "\n",
-        "# Loads an existing endpoint instance using the endpoint name:\n",
-        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
-        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
-        "#   above.\n",
-        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
-        "#   an existing endpoint with the ID 1234567890123456789.\n",
-        "# You may uncomment the code below to load an existing endpoint.\n",
-        "\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)"
+        "model = models[LABEL]\n",
+        "endpoint = endpoints[LABEL]"
       ]
     },
     {
@@ -511,6 +517,20 @@
         "# @markdown the `concept_prompt` of your instance in the prompt.\n",
         "# @markdown You may adjust the parameters below to achieve best image quality.\n",
         "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\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",
+        "\n",
         "prompt = \"A picture of a sks dog in a house\"  # @param {type: \"string\"}\n",
         "negative_prompt = \"\"  # @param {type: \"string\"}\n",
         "height = 1024  # @param {type:\"integer\"}\n",
@@ -526,7 +546,11 @@
         "    \"num_inference_steps\": num_inference_steps,\n",
         "    \"guidance_scale\": guidance_scale,\n",
         "}\n",
-        "response = endpoints[\"sd_xl\"].predict(instances=instances, parameters=parameters)\n",
+        "response = endpoints[\"sd_xl\"].predict(\n",
+        "    instances=instances,\n",
+        "    parameters=parameters,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
         "\n",
         "images = [\n",
         "    common_util.base64_to_image(prediction.get(\"output\"))\n",