diff --git a/autogen/token_count_utils.py b/autogen/token_count_utils.py
index 619a9116e9..8ba5534d9e 100644
--- a/autogen/token_count_utils.py
+++ b/autogen/token_count_utils.py
@@ -11,7 +11,7 @@
 
 import tiktoken
 
-# from .agentchat.contrib.img_utils import num_tokens_from_gpt_image
+from .agentchat.contrib.img_utils import num_tokens_from_gpt_image
 from .import_utils import optional_import_block
 
 # if PIL is not imported, we will redefine num_tokens_from_gpt_image to return 0 tokens for images
@@ -19,8 +19,6 @@
 with optional_import_block() as result:
     import PIL  # noqa: F401
 
-    from .agentchat.contrib.img_utils import num_tokens_from_gpt_image
-
 pil_imported = result.is_successful
 if not pil_imported:
 
diff --git a/test/test_token_count.py b/test/test_token_count.py
index 8118f999e6..8461be8fb0 100755
--- a/test/test_token_count.py
+++ b/test/test_token_count.py
@@ -11,7 +11,7 @@
 from autogen.import_utils import optional_import_block
 
 with optional_import_block() as result:
-    from autogen.agentchat.contrib.img_utils import num_tokens_from_gpt_image  # noqa: F401
+    from PIL import Image  # noqa: F401
 
 img_util_imported = result.is_successful