diff --git a/google/generativeai/vision_models/_vision_models.py b/google/generativeai/vision_models/_vision_models.py
index 3a50b3872..78f495d65 100644
--- a/google/generativeai/vision_models/_vision_models.py
+++ b/google/generativeai/vision_models/_vision_models.py
@@ -34,16 +34,22 @@
 from proto.marshal.collections import repeated
 
 
-# pylint: disable=g-import-not-at-top
-try:
+# pylint: disable=g-import-not-at-top\
+if typing.TYPE_CHECKING:
     from IPython import display as IPython_display
-except ImportError:
-    IPython_display = None
-
-try:
-    from PIL import Image as PIL_Image
-except ImportError:
-    PIL_Image = None
+else:
+    try:
+        from IPython import display as IPython_display
+    except ImportError:
+        IPython_display = None
+
+if typing.TYPE_CHECKING:
+    import PIL.Image as PIL_Image
+else:
+    try:
+        from PIL import Image as PIL_Image
+    except ImportError:
+        PIL_Image = None
 
 
 # This is to get around https://github.com/googleapis/proto-plus-python/issues/488