From b696a7f084c0df48dd94c6cacf15cbfdbeb5b2f1 Mon Sep 17 00:00:00 2001
From: chongkuiqi <18810608559@163.com>
Date: Mon, 3 Mar 2025 14:51:54 +0800
Subject: [PATCH] only use clip-vision-model

---
 model/model_vlm.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/model/model_vlm.py b/model/model_vlm.py
index f9b10d7..2b8fe01 100644
--- a/model/model_vlm.py
+++ b/model/model_vlm.py
@@ -3,7 +3,7 @@ from .model import *
 from typing import Optional, Tuple, List
 from torch import nn
 import warnings
-from transformers import CLIPProcessor, CLIPModel
+from transformers import CLIPProcessor, CLIPVisionModel
 import torch
 
 warnings.filterwarnings('ignore')
@@ -36,7 +36,7 @@ class MiniMindVLM(MiniMindLM):
 
     @staticmethod
     def get_vision_model(model_path="./model/vision_model/clip-vit-base-patch16"):
-        model = CLIPModel.from_pretrained(model_path)
+        model = CLIPVisionModel.from_pretrained(model_path)
         processor = CLIPProcessor.from_pretrained(model_path)
         # 冻结 vision_encoder 的所有参数
         for param in model.parameters():
@@ -52,7 +52,7 @@ class MiniMindVLM(MiniMindLM):
     @staticmethod
     def get_image_embeddings(image_tensors, vision_model):
         with torch.no_grad():
-            outputs = vision_model.vision_model(pixel_values=image_tensors)
+            outputs = vision_model(pixel_values=image_tensors)
         img_embedding = outputs.last_hidden_state[:, 1:, :].squeeze()
         return img_embedding