diff --git a/genai/image_generation/imggen_with_txt.py b/genai/image_generation/imggen_with_txt.py
index 41fa7377ab8..e43868779e7 100644
--- a/genai/image_generation/imggen_with_txt.py
+++ b/genai/image_generation/imggen_with_txt.py
@@ -12,35 +12,32 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from google.genai.types import Image
 
-def generate_images(output_gcs_uri: str) -> str:
+
+def generate_images(output_file: str) -> Image:
     # [START googlegenaisdk_imggen_with_txt]
     from google import genai
-    from google.genai.types import GenerateImagesConfig
 
     client = genai.Client()
 
     # TODO(developer): Update and un-comment below line
-    # output_gcs_uri = "gs://your-bucket/your-prefix"
+    # output_file = "output-image.png"
 
     image = client.models.generate_images(
-        model="imagen-3.0-generate-002",
+        model="imagen-4.0-generate-preview-05-20",
         prompt="A dog reading a newspaper",
-        config=GenerateImagesConfig(
-            aspect_ratio="1:1",
-            number_of_images=1,
-            safety_filter_level="BLOCK_MEDIUM_AND_ABOVE",
-            person_generation="DONT_ADULT",
-            output_gcs_uri=output_gcs_uri,
-        ),
     )
 
+    image.generated_images[0].image.save(output_file)
+
+    print(f"Created output image using {len(image.generated_images[0].image.image_bytes)} bytes")
     # Example response:
-    # gs://your-bucket/your-prefix
-    print(image.generated_images[0].image.gcs_uri)
+    # Created output image using 1234567 bytes
+
     # [END googlegenaisdk_imggen_with_txt]
-    return image.generated_images[0].image.gcs_uri
+    return image.generated_images[0].image
 
 
 if __name__ == "__main__":
-    generate_images(output_gcs_uri="gs://your-bucket/your-prefix")
+    generate_images(output_file="test_resources/dog_newspaper.png")
diff --git a/genai/image_generation/test_image_generation.py b/genai/image_generation/test_image_generation.py
index 764940bdaa7..1e18cf4e11e 100644
--- a/genai/image_generation/test_image_generation.py
+++ b/genai/image_generation/test_image_generation.py
@@ -41,6 +41,8 @@
 # os.environ['GOOGLE_CLOUD_PROJECT'] = "add-your-project-name"
 
 GCS_OUTPUT_BUCKET = "python-docs-samples-tests"
+RESOURCES = os.path.join(os.path.dirname(__file__), "test_resources")
+OUTPUT_FILE = os.path.join(RESOURCES, "dog_newspaper.png")
 
 
 @pytest.fixture(scope="session")
@@ -58,7 +60,7 @@ def output_gcs_uri() -> str:
 
 def test_img_generation(output_gcs_uri: str) -> None:
     response = imggen_with_txt.generate_images(
-        output_gcs_uri=output_gcs_uri
+        OUTPUT_FILE
     )
     assert response
 
diff --git a/genai/image_generation/test_resources/dog_newspaper.png b/genai/image_generation/test_resources/dog_newspaper.png
new file mode 100644
index 00000000000..0e502cdbb61
Binary files /dev/null and b/genai/image_generation/test_resources/dog_newspaper.png differ