diff --git a/packages/firebase_ai/firebase_ai/example/lib/pages/imagen_page.dart b/packages/firebase_ai/firebase_ai/example/lib/pages/imagen_page.dart
index c957f207278e..b5551d402ca4 100644
--- a/packages/firebase_ai/firebase_ai/example/lib/pages/imagen_page.dart
+++ b/packages/firebase_ai/firebase_ai/example/lib/pages/imagen_page.dart
@@ -12,8 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import 'package:flutter/material.dart';
+import 'dart:typed_data';
+
+import 'package:image_picker/image_picker.dart';
 import 'package:firebase_ai/firebase_ai.dart';
+
+import 'package:flutter/material.dart';
 //import 'package:firebase_storage/firebase_storage.dart';
 import '../widgets/message_widget.dart';
 
@@ -38,6 +42,10 @@ class _ImagenPageState extends State<ImagenPage> {
   final List<MessageData> _generatedContent = <MessageData>[];
   bool _loading = false;
 
+  // For image picking
+  ImagenInlineImage? _sourceImage;
+  ImagenInlineImage? _maskImageForEditing;
+
   void _scrollDown() {
     WidgetsBinding.instance.addPostFrameCallback(
       (_) => _scrollController.animateTo(
@@ -80,45 +88,89 @@ class _ImagenPageState extends State<ImagenPage> {
                 vertical: 25,
                 horizontal: 15,
               ),
-              child: Row(
+              child: Column(
                 children: [
-                  Expanded(
-                    child: TextField(
-                      autofocus: true,
-                      focusNode: _textFieldFocus,
-                      controller: _textController,
-                    ),
-                  ),
-                  const SizedBox.square(
-                    dimension: 15,
-                  ),
-                  if (!_loading)
-                    IconButton(
-                      onPressed: () async {
-                        await _testImagen(_textController.text);
-                      },
-                      icon: Icon(
-                        Icons.image_search,
-                        color: Theme.of(context).colorScheme.primary,
+                  // Generate Image Row
+                  Row(
+                    children: [
+                      Expanded(
+                        child: TextField(
+                          autofocus: true,
+                          focusNode: _textFieldFocus,
+                          decoration: const InputDecoration(
+                            hintText: 'Enter a prompt...',
+                          ),
+                          controller: _textController,
+                        ),
                       ),
-                      tooltip: 'Imagen raw data',
-                    )
-                  else
-                    const CircularProgressIndicator(),
-                  // NOTE: Keep this API private until future release.
-                  // if (!_loading)
-                  //   IconButton(
-                  //     onPressed: () async {
-                  //       await _testImagenGCS(_textController.text);
-                  //     },
-                  //     icon: Icon(
-                  //       Icons.imagesearch_roller,
-                  //       color: Theme.of(context).colorScheme.primary,
-                  //     ),
-                  //     tooltip: 'Imagen GCS',
-                  //   )
-                  // else
-                  //   const CircularProgressIndicator(),
+                      const SizedBox.square(dimension: 15),
+                      IconButton(
+                        onPressed: () async {
+                          await _pickSourceImage();
+                        },
+                        icon: Icon(
+                          Icons.add_a_photo,
+                          color: Theme.of(context).colorScheme.primary,
+                        ),
+                        tooltip: 'Pick Source Image',
+                      ),
+                      IconButton(
+                        onPressed: () async {
+                          await _pickMaskImage();
+                        },
+                        icon: Icon(
+                          Icons.add_to_photos,
+                          color: Theme.of(context).colorScheme.primary,
+                        ),
+                        tooltip: 'Pick mask',
+                      ),
+                      IconButton(
+                        onPressed: () async {
+                          await _editImageMaskFree();
+                        },
+                        icon: Icon(
+                          Icons.edit,
+                          color: Theme.of(context).colorScheme.primary,
+                        ),
+                        tooltip: 'Edit Image Mask Free',
+                      ),
+                      IconButton(
+                        onPressed: () async {
+                          await _editImageInpaintOutpaint();
+                        },
+                        icon: Icon(
+                          Icons.masks,
+                          color: Theme.of(context).colorScheme.primary,
+                        ),
+                        tooltip: 'Mask Inpaint Outpaint',
+                      ),
+                      IconButton(
+                        onPressed: () async {
+                          await _upscaleImage();
+                        },
+                        icon: Icon(
+                          Icons.plus_one,
+                          color: Theme.of(context).colorScheme.primary,
+                        ),
+                        tooltip: 'Upscale',
+                      ),
+                      if (!_loading)
+                        IconButton(
+                          onPressed: () async {
+                            await _generateImageFromPrompt(
+                              _textController.text,
+                            );
+                          },
+                          icon: Icon(
+                            Icons.image_search,
+                            color: Theme.of(context).colorScheme.primary,
+                          ),
+                          tooltip: 'Generate Image',
+                        )
+                      else
+                        const CircularProgressIndicator(),
+                    ],
+                  ),
                 ],
               ),
             ),
@@ -128,7 +180,206 @@ class _ImagenPageState extends State<ImagenPage> {
     );
   }
 
-  Future<void> _testImagen(String prompt) async {
+  Future<ImagenInlineImage?> _pickImage() async {
+    final ImagePicker picker = ImagePicker();
+    try {
+      final XFile? imageFile =
+          await picker.pickImage(source: ImageSource.gallery);
+      if (imageFile != null) {
+        // Attempt to get mimeType, default if null.
+        // Note: imageFile.mimeType might be null on some platforms or for some files.
+        final String mimeType = imageFile.mimeType ?? 'image/jpeg';
+        final Uint8List imageBytes = await imageFile.readAsBytes();
+        return ImagenInlineImage(
+            bytesBase64Encoded: imageBytes, mimeType: mimeType);
+      }
+    } catch (e) {
+      _showError('Error picking image: $e');
+    }
+    return null;
+  }
+
+  Future<void> _pickSourceImage() async {
+    final pickedImage = await _pickImage();
+    if (pickedImage != null) {
+      setState(() {
+        _sourceImage = pickedImage;
+      });
+    }
+  }
+
+  Future<void> _pickMaskImage() async {
+    final pickedImage = await _pickImage();
+    if (pickedImage != null) {
+      setState(() {
+        _maskImageForEditing = pickedImage;
+      });
+    }
+  }
+
+  Future<void> _upscaleImage() async {
+    if (_sourceImage == null) {
+      _showError('Please pick a source image for upscaling.');
+      return;
+    }
+    setState(() {
+      _loading = true;
+    });
+
+    setState(() {
+      _generatedContent.add(
+        MessageData(
+          image: Image.memory(_sourceImage!.bytesBase64Encoded),
+          text:
+              'Try to Upscaled image (Factor: ${ImagenUpscaleFactor.x2.name})',
+          fromUser: true,
+        ),
+      );
+      _scrollDown();
+    });
+
+    try {
+      final response = await widget.model.upscaleImage(
+        image: _sourceImage!,
+        upscaleFactor: ImagenUpscaleFactor.x2,
+      );
+      if (response.images.isNotEmpty) {
+        final upscaledImage = response.images[0];
+        setState(() {
+          _generatedContent.add(
+            MessageData(
+              image: Image.memory(upscaledImage.bytesBase64Encoded),
+              text: 'Upscaled image (Factor: ${ImagenUpscaleFactor.x2.name})',
+              fromUser: false,
+            ),
+          );
+          _scrollDown();
+        });
+      } else {
+        _showError('No image was returned from upscaling.');
+      }
+    } catch (e) {
+      _showError('Error upscaling image: $e');
+    }
+
+    setState(() {
+      _loading = false;
+    });
+  }
+
+  Future<void> _editImageInpaintOutpaint() async {
+    if (_sourceImage == null || _maskImageForEditing == null) {
+      _showError(
+          'Please pick a source image and a mask image for inpainting/outpainting.');
+      return;
+    }
+    setState(() {
+      _loading = true;
+    });
+
+    final String prompt = _textController.text;
+
+    setState(() {
+      _generatedContent.add(
+        MessageData(
+          image: Image.memory(_sourceImage!.bytesBase64Encoded),
+          text: prompt,
+          fromUser: true,
+        ),
+      );
+      _scrollDown();
+    });
+
+    final editConfig = ImagenEditingConfig(
+      image: _sourceImage!,
+      mask: _maskImageForEditing,
+      maskDilation: 0.01,
+      editSteps: 50,
+    );
+
+    try {
+      final response = await widget.model.editImage(
+        prompt,
+        config: editConfig,
+      );
+      if (response.images.isNotEmpty) {
+        final editedImage = response.images[0];
+        setState(() {
+          _generatedContent.add(
+            MessageData(
+              image: Image.memory(editedImage.bytesBase64Encoded),
+              text: 'Edited image (Inpaint/Outpaint): $prompt',
+              fromUser: false,
+            ),
+          );
+          _scrollDown();
+        });
+      } else {
+        _showError('No image was returned from editing.');
+      }
+    } catch (e) {
+      _showError('Error editing image: $e');
+    }
+    setState(() {
+      _loading = false;
+    });
+  }
+
+  Future<void> _editImageMaskFree() async {
+    if (_sourceImage == null) {
+      _showError('Please pick a source image for mask-free editing.');
+      return;
+    }
+    setState(() {
+      _loading = true;
+    });
+
+    final String prompt = _textController.text;
+
+    setState(() {
+      _generatedContent.add(
+        MessageData(
+          image: Image.memory(_sourceImage!.bytesBase64Encoded),
+          text: prompt,
+          fromUser: true,
+        ),
+      );
+      _scrollDown();
+    });
+    final editConfig = ImagenEditingConfig.maskFree(
+      image: _sourceImage!,
+      // numberOfImages: 1, // Default in model or could be added to UI
+    );
+
+    try {
+      final response = await widget.model.editImage(
+        prompt,
+        config: editConfig,
+      );
+      if (response.images.isNotEmpty) {
+        final editedImage = response.images[0];
+        setState(() {
+          _generatedContent.add(
+            MessageData(
+              image: Image.memory(editedImage.bytesBase64Encoded),
+              text: 'Edited image (Mask-Free): $prompt',
+              fromUser: false,
+            ),
+          );
+          _scrollDown();
+        });
+      } else {
+        _showError('No image was returned from mask-free editing.');
+      }
+    } catch (e) {
+      _showError('Error performing mask-free edit: $e');
+    }
+    setState(() {
+      _loading = false;
+    });
+  }
+
+  Future<void> _generateImageFromPrompt(String prompt) async {
     setState(() {
       _loading = true;
     });
diff --git a/packages/firebase_ai/firebase_ai/example/macos/Runner/DebugProfile.entitlements b/packages/firebase_ai/firebase_ai/example/macos/Runner/DebugProfile.entitlements
index b4bd9ee174a1..8560da29b687 100644
--- a/packages/firebase_ai/firebase_ai/example/macos/Runner/DebugProfile.entitlements
+++ b/packages/firebase_ai/firebase_ai/example/macos/Runner/DebugProfile.entitlements
@@ -14,5 +14,7 @@
 	<true/>
   <key>com.apple.security.device.audio-input</key>
   <true/>
+  <key>com.apple.security.files.user-selected.read-only</key>
+  <true/>
 </dict>
 </plist>
diff --git a/packages/firebase_ai/firebase_ai/example/macos/Runner/Info.plist b/packages/firebase_ai/firebase_ai/example/macos/Runner/Info.plist
index a81b3fd0d617..d4369e6253fa 100644
--- a/packages/firebase_ai/firebase_ai/example/macos/Runner/Info.plist
+++ b/packages/firebase_ai/firebase_ai/example/macos/Runner/Info.plist
@@ -30,5 +30,7 @@
 	<string>NSApplication</string>
   <key>NSMicrophoneUsageDescription</key>
   <string>Permission to Record audio</string>
+  <key>NSPhotoLibraryUsageDescription</key>
+  <string>This app needs access to your photo library to let you select a profile picture.</string>
 </dict>
 </plist>
diff --git a/packages/firebase_ai/firebase_ai/example/pubspec.yaml b/packages/firebase_ai/firebase_ai/example/pubspec.yaml
index 4868f106d648..475907298c36 100644
--- a/packages/firebase_ai/firebase_ai/example/pubspec.yaml
+++ b/packages/firebase_ai/firebase_ai/example/pubspec.yaml
@@ -27,6 +27,7 @@ dependencies:
     sdk: flutter
   flutter_markdown: ^0.6.20
   flutter_soloud: ^3.1.6
+  image_picker: ^1.1.2
   path_provider: ^2.1.5
   record: ^5.2.1
 
diff --git a/packages/firebase_ai/firebase_ai/lib/firebase_ai.dart b/packages/firebase_ai/firebase_ai/lib/firebase_ai.dart
index dbc95e1bca24..e7af8ecc96a0 100644
--- a/packages/firebase_ai/firebase_ai/lib/firebase_ai.dart
+++ b/packages/firebase_ai/firebase_ai/lib/firebase_ai.dart
@@ -65,8 +65,11 @@ export 'src/imagen_api.dart'
         ImagenSafetyFilterLevel,
         ImagenPersonFilterLevel,
         ImagenGenerationConfig,
-        ImagenAspectRatio;
-export 'src/imagen_content.dart' show ImagenInlineImage;
+        ImagenAspectRatio,
+        ImagenEditingConfig,
+        ImagenEditMode,
+        ImagenUpscaleFactor;
+export 'src/imagen_content.dart' show ImagenInlineImage, ImagenGenerationResponse;
 export 'src/live_api.dart'
     show
         LiveGenerationConfig,
diff --git a/packages/firebase_ai/firebase_ai/lib/src/imagen_api.dart b/packages/firebase_ai/firebase_ai/lib/src/imagen_api.dart
index 1579b6740a92..1a5b02d17841 100644
--- a/packages/firebase_ai/firebase_ai/lib/src/imagen_api.dart
+++ b/packages/firebase_ai/firebase_ai/lib/src/imagen_api.dart
@@ -15,6 +15,8 @@ import 'dart:developer';
 
 import 'package:meta/meta.dart';
 
+import 'imagen_content.dart';
+
 /// Specifies the level of safety filtering for image generation.
 ///
 /// If not specified, default will be "block_medium_and_above".
@@ -232,3 +234,86 @@ final class ImagenFormat {
           'compressionQuality': compressionQuality,
       };
 }
+
+/// Enum representing the mode for image editing.
+@experimental
+enum ImagenEditMode {
+  /// Inpaint mode for image editing.
+  inpaint,
+
+  /// Outpaint mode for image editing.
+  outpaint,
+}
+
+/// Enum representing the upscale factor for image upscaling.
+@experimental
+enum ImagenUpscaleFactor {
+  /// Upscale factor of 2x.
+  x2('x2'),
+
+  /// Upscale factor of 4x.
+  x4('x4');
+
+  const ImagenUpscaleFactor(this._jsonString);
+
+  final String _jsonString;
+
+  // ignore: public_member_api_docs
+  String toJson() => _jsonString;
+}
+
+/// Configuration for Imagen image editing.
+@experimental
+final class ImagenEditingConfig {
+  /// Source image for editing.
+  final ImagenInlineImage image;
+
+  /// Mask image for editing, optional for mask-free editing.
+  final ImagenInlineImage? mask;
+
+  /// Mask dilation factor.
+  final double? maskDilation;
+
+  /// Number of editing steps.
+  final int? editSteps;
+
+  /// Number of images to generate.
+  final int? numberOfImages;
+
+  /// Editing mode.
+  final ImagenEditMode? editMode;
+
+  // ignore: public_member_api_docs
+  ImagenEditingConfig({
+    required this.image,
+    this.mask,
+    this.maskDilation,
+    this.editSteps,
+    this.numberOfImages,
+    this.editMode,
+  });
+
+  /// Factory constructor for mask-free image editing.
+  /// Takes numberOfImages as an optional parameter.
+  factory ImagenEditingConfig.maskFree({
+    required ImagenInlineImage image,
+    int? numberOfImages,
+  }) {
+    return ImagenEditingConfig(
+      image: image,
+      numberOfImages: numberOfImages,
+      // Mask and editMode related to masking are left null for mask-free.
+      // Other fields like maskDilation, editSteps are also left null.
+    );
+  }
+
+  // ignore: public_member_api_docs
+  Map<String, dynamic> toJson() => {
+        'image': image.toJson(),
+        if (mask != null) 'mask': mask!.toJson(),
+        if (maskDilation != null) 'maskDilation': maskDilation,
+        if (editSteps != null) 'editSteps': editSteps,
+        if (numberOfImages != null) 'numberOfImages': numberOfImages,
+        if (editMode != null) 'editMode': editMode!.name,
+      };
+}
diff --git a/packages/firebase_ai/firebase_ai/lib/src/imagen_model.dart b/packages/firebase_ai/firebase_ai/lib/src/imagen_model.dart
index bf4731a3b264..752942e104d1 100644
--- a/packages/firebase_ai/firebase_ai/lib/src/imagen_model.dart
+++ b/packages/firebase_ai/firebase_ai/lib/src/imagen_model.dart
@@ -57,7 +57,7 @@ final class ImagenModel extends BaseApiClientModel {
       if (gcsUri != null) 'storageUri': gcsUri,
       'sampleCount': _generationConfig?.numberOfImages ?? 1,
       if (_generationConfig?.aspectRatio case final aspectRatio?)
-        'aspectRatio': aspectRatio,
+        'aspectRatio': aspectRatio.toJson(),
       if (_generationConfig?.negativePrompt case final negativePrompt?)
         'negativePrompt': negativePrompt,
       if (_generationConfig?.addWatermark case final addWatermark?)
@@ -110,6 +110,97 @@ final class ImagenModel extends BaseApiClientModel {
         (jsonObject) =>
             parseImagenGenerationResponse<ImagenGCSImage>(jsonObject),
       );
+
+  /// Edits an image based on a prompt and configuration.
+  @experimental
+  Future<ImagenGenerationResponse<ImagenInlineImage>> editImage(
+    String prompt, {
+    required ImagenEditingConfig config,
+  }) async {
+    // Construct the request payload.
+    final payload = {
+      'instances': [
+        {
+          'prompt': prompt,
+          'image': config.image.toJson(),
+          if (config.mask != null) 'mask': config.mask!.toJson(),
+        }
+      ],
+      'parameters': {
+        if (config.editMode != null) 'editMode': config.editMode!.name,
+        if (config.maskDilation != null) 'maskDilation': config.maskDilation,
+        if (config.editSteps != null) 'editSteps': config.editSteps,
+        'sampleCount':
+            config.numberOfImages ?? _generationConfig?.numberOfImages ?? 1,
+
+        // Parameters from model-level _generationConfig and _safetySettings
+        if (_generationConfig?.aspectRatio case final aspectRatio?)
+          'aspectRatio': aspectRatio.toJson(),
+        if (_generationConfig?.negativePrompt case final negativePrompt?)
+          'negativePrompt': negativePrompt,
+        if (_generationConfig?.addWatermark case final addWatermark?)
+          'addWatermark': addWatermark,
+        if (_generationConfig?.imageFormat case final imageFormat?)
+          'outputOption': imageFormat.toJson(),
+        if (_safetySettings?.personFilterLevel case final personFilterLevel?)
+          'personGeneration': personFilterLevel.toJson(),
+        if (_safetySettings?.safetyFilterLevel case final safetyFilterLevel?)
+          'safetySetting': safetyFilterLevel.toJson(),
+      },
+    };
+
+    return makeRequest(
+      Task.predict,
+      payload,
+      (jsonObject) =>
+          parseImagenGenerationResponse<ImagenInlineImage>(jsonObject),
+    );
+  }
+
+  /// Upscales an image.
+  @experimental
+  Future<ImagenGenerationResponse<ImagenInlineImage>> upscaleImage({
+    required ImagenInlineImage image,
+    required ImagenUpscaleFactor upscaleFactor,
+    ImagenSafetySettings? safetySettings,
+    ImagenGenerationConfig? generationConfig,
+  }) async {
+    // Construct the request payload for upscaling.
+    final payload = {
+      'instances': [
+        {
+          'image': image.toJson(),
+        }
+      ],
+      'parameters': {
+        'upscaleFactor': upscaleFactor.toJson(),
+        if (generationConfig?.aspectRatio ?? _generationConfig?.aspectRatio
+            case final aspectRatio?)
+          'aspectRatio': aspectRatio.toJson(),
+        if (generationConfig?.addWatermark ?? _generationConfig?.addWatermark
+            case final addWatermark?)
+          'addWatermark': addWatermark,
+        if (generationConfig?.imageFormat ?? _generationConfig?.imageFormat
+            case final imageFormat?)
+          'outputOption': imageFormat.toJson(),
+        if (safetySettings?.personFilterLevel ??
+                _safetySettings?.personFilterLevel
+            case final personFilterLevel?)
+          'personGeneration': personFilterLevel.toJson(),
+        if (safetySettings?.safetyFilterLevel ??
+                _safetySettings?.safetyFilterLevel
+            case final safetyFilterLevel?)
+          'safetySetting': safetyFilterLevel.toJson(),
+      },
+    };
+
+    return makeRequest(
+      Task.predict,
+      payload,
+      (jsonObject) =>
+          parseImagenGenerationResponse<ImagenInlineImage>(jsonObject),
+    );
+  }
 }
 
 /// Returns a [ImagenModel] using it's private constructor.
diff --git a/packages/firebase_ai/firebase_ai/test/chat_test.dart b/packages/firebase_ai/firebase_ai/test/chat_test.dart
index ab5819f0f12a..76f302e1162b 100644
--- a/packages/firebase_ai/firebase_ai/test/chat_test.dart
+++ b/packages/firebase_ai/firebase_ai/test/chat_test.dart
@@ -17,12 +17,12 @@ import 'package:firebase_ai/src/base_model.dart';
 import 'package:firebase_core/firebase_core.dart';
 import 'package:flutter_test/flutter_test.dart';
 
-import 'mock.dart';
+import 'mocks/mock_core.dart';
 import 'utils/matchers.dart';
 import 'utils/stub_client.dart';
 
 void main() {
-  setupFirebaseVertexAIMocks();
+  setupFirebaseAIMocks();
   // ignore: unused_local_variable
   late FirebaseApp app;
 
diff --git a/packages/firebase_ai/firebase_ai/test/firebase_vertexai_test.dart b/packages/firebase_ai/firebase_ai/test/firebase_vertexai_test.dart
index 8edb2d0f8480..a5a0e6dd2518 100644
--- a/packages/firebase_ai/firebase_ai/test/firebase_vertexai_test.dart
+++ b/packages/firebase_ai/firebase_ai/test/firebase_vertexai_test.dart
@@ -17,10 +17,10 @@ import 'package:firebase_app_check/firebase_app_check.dart';
 import 'package:firebase_core/firebase_core.dart';
 import 'package:flutter_test/flutter_test.dart';
 
-import 'mock.dart';
+import 'mocks/mock_core.dart';
 
 void main() {
-  setupFirebaseVertexAIMocks();
+  setupFirebaseAIMocks();
   // ignore: unused_local_variable
   late FirebaseApp app;
   // ignore: unused_local_variable
diff --git a/packages/firebase_ai/firebase_ai/test/google_ai_generative_model_test.dart b/packages/firebase_ai/firebase_ai/test/google_ai_generative_model_test.dart
index 9883102c2729..ef67d2505b34 100644
--- a/packages/firebase_ai/firebase_ai/test/google_ai_generative_model_test.dart
+++ b/packages/firebase_ai/firebase_ai/test/google_ai_generative_model_test.dart
@@ -17,12 +17,12 @@ import 'package:firebase_ai/src/base_model.dart';
 import 'package:firebase_core/firebase_core.dart';
 import 'package:flutter_test/flutter_test.dart';
 
-import 'mock.dart';
+import 'mocks/mock_core.dart';
 import 'utils/matchers.dart';
 import 'utils/stub_client.dart';
 
 void main() {
-  setupFirebaseVertexAIMocks();
+  setupFirebaseAIMocks();
   late FirebaseApp app;
   setUpAll(() async {
     // Initialize Firebase
diff --git a/packages/firebase_ai/firebase_ai/test/imagen_test.dart b/packages/firebase_ai/firebase_ai/test/imagen_test.dart
index 4bd7ae5b763a..d56a8d6cd8a3 100644
--- a/packages/firebase_ai/firebase_ai/test/imagen_test.dart
+++ b/packages/firebase_ai/firebase_ai/test/imagen_test.dart
@@ -15,11 +15,30 @@
 import 'dart:convert';
 import 'dart:typed_data';
 
+import 'package:firebase_ai/firebase_ai.dart';
 import 'package:firebase_ai/src/error.dart';
 import 'package:firebase_ai/src/imagen_content.dart';
+import 'package:firebase_core/firebase_core.dart';
 import 'package:flutter_test/flutter_test.dart';
 
+import 'mocks/mock_core.dart';
+
+// Mock HttpApiClient
 void main() {
+  setupFirebaseAIMocks();
+
+  setUpAll(() async {
+    await Firebase.initializeApp(
+      name: 'testApp',
+      options: const FirebaseOptions(
+        apiKey: 'test-api-key',
+        appId: 'test-app-id',
+        messagingSenderId: 'test-sender-id',
+        projectId: 'test-project-id',
+      ),
+    );
+  });
+
   group('ImagenInlineImage', () {
     test('fromJson with valid base64', () {
       final json = {
diff --git a/packages/firebase_ai/firebase_ai/test/mock.dart b/packages/firebase_ai/firebase_ai/test/mocks/mock_core.dart
similarity index 89%
rename from packages/firebase_ai/firebase_ai/test/mock.dart
rename to packages/firebase_ai/firebase_ai/test/mocks/mock_core.dart
index ed883d924371..0a78e72888a7 100644
--- a/packages/firebase_ai/firebase_ai/test/mock.dart
+++ b/packages/firebase_ai/firebase_ai/test/mocks/mock_core.dart
@@ -18,7 +18,7 @@ import 'package:flutter_test/flutter_test.dart';
 import 'package:mockito/mockito.dart';
 import 'package:plugin_platform_interface/plugin_platform_interface.dart';
 
-class MockFirebaseAppVertexAI implements TestFirebaseCoreHostApi {
+class MockFirebaseAppAI implements TestFirebaseCoreHostApi {
   @override
   Future<PigeonInitializeResponse> initializeApp(
     String appName,
@@ -58,16 +58,16 @@ class MockFirebaseAppVertexAI implements TestFirebaseCoreHostApi {
   }
 }
 
-void setupFirebaseVertexAIMocks() {
+void setupFirebaseAIMocks() {
   TestWidgetsFlutterBinding.ensureInitialized();
 
-  TestFirebaseCoreHostApi.setup(MockFirebaseAppVertexAI());
+  TestFirebaseCoreHostApi.setup(MockFirebaseAppAI());
 }
 
 // FirebaseVertexAIPlatform Mock
-class MockFirebaseVertexAI extends Mock
+class MockFirebaseAI extends Mock
     with
         // ignore: prefer_mixin, plugin_platform_interface needs to migrate to use `mixin`
         MockPlatformInterfaceMixin {
-  MockFirebaseVertexAI();
+  MockFirebaseAI();
 }
diff --git a/packages/firebase_ai/firebase_ai/test/model_test.dart b/packages/firebase_ai/firebase_ai/test/model_test.dart
index 2ddf4d55406c..7049d43a5e33 100644
--- a/packages/firebase_ai/firebase_ai/test/model_test.dart
+++ b/packages/firebase_ai/firebase_ai/test/model_test.dart
@@ -17,12 +17,12 @@ import 'package:firebase_ai/src/base_model.dart';
 import 'package:firebase_core/firebase_core.dart';
 import 'package:flutter_test/flutter_test.dart';
 
-import 'mock.dart';
+import 'mocks/mock_core.dart';
 import 'utils/matchers.dart';
 import 'utils/stub_client.dart';
 
 void main() {
-  setupFirebaseVertexAIMocks();
+  setupFirebaseAIMocks();
   // ignore: unused_local_variable
   late FirebaseApp app;
   setUpAll(() async {
diff --git a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
index 0ab750b13fef..c5dc1e7fbcaf 100644
--- a/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
+++ b/packages/firebase_vertexai/firebase_vertexai/example/lib/pages/imagen_page.dart
@@ -12,8 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import 'package:flutter/material.dart';
 import 'package:firebase_vertexai/firebase_vertexai.dart';
+
+import 'package:flutter/material.dart';
 //import 'package:firebase_storage/firebase_storage.dart';
 import '../widgets/message_widget.dart';
 
@@ -80,45 +81,39 @@ class _ImagenPageState extends State<ImagenPage> {
                 vertical: 25,
                 horizontal: 15,
               ),
-              child: Row(
+              child: Column(
                 children: [
-                  Expanded(
-                    child: TextField(
-                      autofocus: true,
-                      focusNode: _textFieldFocus,
-                      controller: _textController,
-                    ),
-                  ),
-                  const SizedBox.square(
-                    dimension: 15,
-                  ),
-                  if (!_loading)
-                    IconButton(
-                      onPressed: () async {
-                        await _testImagen(_textController.text);
-                      },
-                      icon: Icon(
-                        Icons.image_search,
-                        color: Theme.of(context).colorScheme.primary,
+                  // Generate Image Row
+                  Row(
+                    children: [
+                      Expanded(
+                        child: TextField(
+                          autofocus: true,
+                          focusNode: _textFieldFocus,
+                          decoration: const InputDecoration(
+                            hintText: 'Enter a prompt...',
+                          ),
+                          controller: _textController,
+                        ),
                       ),
-                      tooltip: 'Imagen raw data',
-                    )
-                  else
-                    const CircularProgressIndicator(),
-                  // NOTE: Keep this API private until future release.
-                  // if (!_loading)
-                  //   IconButton(
-                  //     onPressed: () async {
-                  //       await _testImagenGCS(_textController.text);
-                  //     },
-                  //     icon: Icon(
-                  //       Icons.imagesearch_roller,
-                  //       color: Theme.of(context).colorScheme.primary,
-                  //     ),
-                  //     tooltip: 'Imagen GCS',
-                  //   )
-                  // else
-                  //   const CircularProgressIndicator(),
+                      const SizedBox.square(dimension: 15),
+                      if (!_loading)
+                        IconButton(
+                          onPressed: () async {
+                            await _generateImageFromPrompt(
+                              _textController.text,
+                            );
+                          },
+                          icon: Icon(
+                            Icons.image_search,
+                            color: Theme.of(context).colorScheme.primary,
+                          ),
+                          tooltip: 'Generate Image',
+                        )
+                      else
+                        const CircularProgressIndicator(),
+                    ],
+                  ),
                 ],
               ),
             ),
@@ -128,7 +123,7 @@ class _ImagenPageState extends State<ImagenPage> {
     );
   }
 
-  Future<void> _testImagen(String prompt) async {
+  Future<void> _generateImageFromPrompt(String prompt) async {
     setState(() {
       _loading = true;
     });