From eab1638686f42585c03f0283e2e471349de10cf8 Mon Sep 17 00:00:00 2001
From: bursteratom <sunny19981005@outlook.com>
Date: Fri, 6 Dec 2024 14:37:32 -0500
Subject: [PATCH] lint

---
 src/axolotl/utils/models.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
index 3da5cc0dd..613bf9b21 100644
--- a/src/axolotl/utils/models.py
+++ b/src/axolotl/utils/models.py
@@ -30,6 +30,7 @@ from transformers import (  # noqa: F401
     AddedToken,
     AutoConfig,
     AutoModelForCausalLM,
+    AutoModelForImageTextToText,
     AutoModelForVision2Seq,
     AutoProcessor,
     AutoTokenizer,
@@ -553,6 +554,10 @@ class ModelLoader:
                 self.AutoModelLoader = (  # pylint: disable=invalid-name
                     MllamaForConditionalGeneration
                 )
+            elif self.model_config.model_type == "qwen2_vl":
+                self.AutoModelLoader = (  # pylint: disable=invalid-name
+                    AutoModelForImageTextToText
+                )
             else:
                 self.AutoModelLoader = (
                     AutoModelForVision2Seq  # pylint: disable=invalid-name