diff --git a/generative_ai/imagen/edit_image_mask.py b/generative_ai/imagen/edit_image_mask.py
index f8049620d06..12797eaf137 100644
--- a/generative_ai/imagen/edit_image_mask.py
+++ b/generative_ai/imagen/edit_image_mask.py
@@ -21,6 +21,7 @@
 """
 
 # [START aiplatform_imagen_edit_image_mask]
+# [START generativeaionvertexai_imagen_edit_image_mask]
 
 import argparse
 
@@ -73,6 +74,7 @@ def edit_image_mask(
     return images
 
 
+# [END generativeaionvertexai_imagen_edit_image_mask]
 # [END aiplatform_imagen_edit_image_mask]
 
 if __name__ == "__main__":
diff --git a/generative_ai/imagen/edit_image_mask_free.py b/generative_ai/imagen/edit_image_mask_free.py
index a0506eaaa1d..8e46f29f571 100644
--- a/generative_ai/imagen/edit_image_mask_free.py
+++ b/generative_ai/imagen/edit_image_mask_free.py
@@ -20,6 +20,7 @@
 """
 
 # [START aiplatform_imagen_edit_image_mask_free]
+# [START generativeaionvertexai_imagen_edit_image_mask_free]
 
 import argparse
 
@@ -64,6 +65,7 @@ def edit_image_mask_free(
     return images
 
 
+# [END generativeaionvertexai_imagen_edit_image_mask_free]
 # [END aiplatform_imagen_edit_image_mask_free]
 
 if __name__ == "__main__":
diff --git a/generative_ai/imagen/get_short_form_image_captions.py b/generative_ai/imagen/get_short_form_image_captions.py
index b5a858a7486..fa1cb49679d 100644
--- a/generative_ai/imagen/get_short_form_image_captions.py
+++ b/generative_ai/imagen/get_short_form_image_captions.py
@@ -19,6 +19,7 @@
 """
 
 # [START aiplatform_imagen_get_short_form_image_captions]
+# [START generativeaionvertexai_imagen_get_short_form_image_captions]
 
 import argparse
 
@@ -52,6 +53,7 @@ def get_short_form_image_captions(
     return captions
 
 
+# [END generativeaionvertexai_imagen_get_short_form_image_captions]
 # [END aiplatform_imagen_get_short_form_image_captions]
 
 if __name__ == "__main__":
diff --git a/generative_ai/imagen/get_short_form_image_responses.py b/generative_ai/imagen/get_short_form_image_responses.py
index c5d632fa63a..40040657b8c 100644
--- a/generative_ai/imagen/get_short_form_image_responses.py
+++ b/generative_ai/imagen/get_short_form_image_responses.py
@@ -20,6 +20,7 @@
 """
 
 # [START aiplatform_imagen_get_short_form_image_responses]
+# [START generativeaionvertexai_imagen_get_short_form_image_responses]
 
 import argparse
 
@@ -54,6 +55,7 @@ def get_short_form_image_responses(
     return answers
 
 
+# [END generativeaionvertexai_imagen_get_short_form_image_responses]
 # [END aiplatform_imagen_get_short_form_image_responses]
 
 if __name__ == "__main__":