From a9687f6224eea298ce335dd06fc5f31466fef84a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 11 Mar 2025 22:36:51 -0700
Subject: [PATCH] Update vision.py

---
 unsloth/models/vision.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py
index 31d0acf62a..77774393c9 100644
--- a/unsloth/models/vision.py
+++ b/unsloth/models/vision.py
@@ -58,7 +58,10 @@ def unsloth_base_fast_generate(
     dtype = _get_dtype(self.config.torch_dtype)
 
     # Check if VLM
-    is_vlm = (x.endswith("ForConditionalGeneration") for x in self.config.architectures)
+    is_vlm = (
+        x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
+        for x in self.config.architectures
+    )
     is_vlm = is_vlm or hasattr(self.config, "vision_config")
 
     # Remove token_type_ids