diff --git a/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb b/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb
index 431ecdf11..bd95f43e3 100644
--- a/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb
+++ b/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb
@@ -54,7 +54,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "3de7470326a2"
@@ -128,13 +127,11 @@
         "\n",
         "    google_auth.authenticate_user()\n",
         "    ! pip3 install --upgrade pip\n",
-        "    ! pip3 install torchvision==0.14.1\n",
+        "    ! pip3 install torch==2.0.1+cu118\n",
         "    ! pip3 install transformers==4.27.1\n",
         "    ! pip3 install diffusers==0.20.1\n",
         "    ! pip3 install datasets==2.9.0\n",
         "    ! pip3 install accelerate==0.21.0\n",
-        "    ! pip3 install triton==2.0.0.dev20221120\n",
-        "    ! pip3 install xformers==0.0.16\n",
         "    # Install gdown for downloading example training images.\n",
         "    ! pip3 install gdown\n",
         "    # Remove wrong cublas version.\n",
@@ -154,7 +151,7 @@
       },
       "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_inpainting.ipynb) to deploy the notebook to a Vertex AI Workbench Instance.\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.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",
@@ -258,10 +255,8 @@
         ")\n",
         "\n",
         "# The pre-built serving docker images. They contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve\"\n",
-        "PEFT_SERVE_DOCKER_URI = (\n",
-        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve\"\n",
-        ")"
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve:sd-optimized\"\n",
+        "PEFT_SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve:sd-optimized\""
       ]
     },
     {
@@ -434,15 +429,22 @@
       "outputs": [],
       "source": [
         "import torch\n",
-        "from diffusers import StableDiffusionPipeline\n",
+        "from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline\n",
+        "from diffusers.models.attention_processor import AttnProcessor2_0\n",
         "\n",
         "model_id = \"runwayml/stable-diffusion-v1-5\"\n",
         "pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)\n",
         "pipe = pipe.to(\"cuda\")\n",
         "\n",
+        "\n",
+        "pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)\n",
+        "# We enable the memory-efficient attention implementation in PyTorch 2.0 which\n",
+        "# automatically enables several optimizations depending on the inputs and the GPU type.\n",
+        "pipe.unet.set_attn_processor(AttnProcessor2_0())\n",
+        "\n",
         "prompt = \"a photo of an astronaut riding a horse on mars\"\n",
         "\n",
-        "results = pipe(prompt=prompt, guidance_scale=7.5)\n",
+        "results = pipe(prompt=prompt, guidance_scale=7.5, num_inference_steps=25)\n",
         "images = results.images\n",
         "nsfw_detects = results.nsfw_content_detected\n",
         "display(images[0])\n",
@@ -751,7 +753,6 @@
       ]
     },
     {
-      "attachments": {},
       "cell_type": "markdown",
       "metadata": {
         "id": "EY-iQTub1UJg"
@@ -761,7 +762,7 @@
         "\n",
         "This section uses [LoRA](https://arxiv.org/abs/2106.09685) to finetune the [stable-diffusion-v1.5](https://huggingface.co/runwayml/stable-diffusion-v1-5) model with [lambdalabs/pokemon-blip-captions](https://huggingface.co/datasets/lambdalabs/pokemon-blip-captions).\n",
         "\n",
-        "Finetuning with LoRA \n",
+        "Finetuning with LoRA\n",
         "\n",
         "The LoRA weights will be saved after the finetuning job finishes and it can be loaded by the [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/text2img) to run inference."
       ]