diff --git a/genai/image_generation/imggen_mmflash_with_txt.py b/genai/image_generation/imggen_mmflash_with_txt.py
index 0ee371b7e84..cd6c458a757 100644
--- a/genai/image_generation/imggen_mmflash_with_txt.py
+++ b/genai/image_generation/imggen_mmflash_with_txt.py
@@ -15,10 +15,12 @@
 
 def generate_content() -> str:
     # [START googlegenaisdk_imggen_mmflash_with_txt]
+    import os
+    from io import BytesIO
+
     from google import genai
     from google.genai.types import GenerateContentConfig, Modality
     from PIL import Image
-    from io import BytesIO
 
     client = genai.Client()
 
@@ -34,7 +36,10 @@ def generate_content() -> str:
             print(part.text)
         elif part.inline_data:
             image = Image.open(BytesIO((part.inline_data.data)))
-            image.save("output_folder/example-image-eiffel-tower.png")
+            # Ensure the output directory exists
+            output_dir = "output_folder"
+            os.makedirs(output_dir, exist_ok=True)
+            image.save(os.path.join(output_dir, "example-image-eiffel-tower.png"))
 
     # [END googlegenaisdk_imggen_mmflash_with_txt]
     return True