diff --git a/vision/getting-started/imagen3_editing.ipynb b/vision/getting-started/imagen3_editing.ipynb
index 8302b2c55fc..d41c5c19897 100644
--- a/vision/getting-started/imagen3_editing.ipynb
+++ b/vision/getting-started/imagen3_editing.ipynb
@@ -543,10 +543,10 @@
         "id": "68909c926952"
       },
       "source": [
-        "### Product background editing via bgswap mode\n",
+        "### Product background editing via background swap mode\n",
         "\n",
         "\n",
-        "You can also use Imagen 3 for product image editing. By setting `edit_mode` to \"bgswap\", you can maintain the product content while modifying the image background.\n",
+        "You can also use Imagen 3 for product image editing. By setting `edit_mode` to \"background-swap\", you can maintain the product content while modifying the image background.\n",
         "\n",
         "For this example, start with an image stored in a Google Cloud Storage bucket, and provide a prompt describing the new background scene.  "
       ]
@@ -563,11 +563,13 @@
         "    gcs_uri=\"gs://cloud-samples-data/generative-ai/image/suitcase.png\"\n",
         ")\n",
         "raw_ref_image = RawReferenceImage(image=product_image, reference_id=0)\n",
-        "prompt = \"a light blue suitcase[0] in front of a window in an airport, lots of bright, natural lighting coming in from the windows, planes taking off in the distance\"\n",
+        "mask_ref_image = MaskReferenceImage(reference_id=1, image=None, mask_mode=\"background\")\n",
+        "\n",
+        "prompt = \"a light blue suitcase in front of a window in an airport, lots of bright, natural lighting coming in from the windows, planes taking off in the distance\"\n",
         "edited_image = edit_model.edit_image(\n",
         "    prompt=prompt,\n",
-        "    edit_mode=\"bgswap\",\n",
-        "    reference_images=[raw_ref_image],\n",
+        "    edit_mode=\"background-swap\",\n",
+        "    reference_images=[raw_ref_image, mask_ref_image],\n",
         "    seed=1,\n",
         "    number_of_images=1,\n",
         "    safety_filter_level=\"block_some\",\n",