update'

main.py

@@ -4,6 +4,8 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
 )
+from transformers import Qwen2_5_VLForConditionalGeneration
+
 from qwen_vl_utils import process_vision_info
 import torch
 import logging
@@ -13,22 +15,41 @@ logging.basicConfig(level=logging.INFO)
 app = FastAPI()
 
 # Qwen2.5-VL Model Setup
-qwen_checkpoint = "Qwen/Qwen2.5-VL-7B-Instruct"
-min_pixels = 256 * 28 * 28
-max_pixels = 1280 * 28 * 28
+# qwen_checkpoint = "Qwen/Qwen2.5-VL-7B-Instruct"
+# min_pixels = 256 * 28 * 28
+# max_pixels = 1280 * 28 * 28
+
+# processor = AutoProcessor.from_pretrained(
+#     qwen_checkpoint,
+#     min_pixels=min_pixels,
+#     max_pixels=max_pixels,
+# )
+
+# qwen_model = AutoModelForCausalLM.from_pretrained(
+#     qwen_checkpoint,
+#     torch_dtype=torch.bfloat16,
+#     device_map="auto",
+# )
+
 
+
+checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
+min_pixels = 256*28*28
+max_pixels = 1280*28*28
 processor = AutoProcessor.from_pretrained(
-    qwen_checkpoint,
+    checkpoint,
     min_pixels=min_pixels,
-    max_pixels=max_pixels,
+    max_pixels=max_pixels
 )
-
-qwen_model = AutoModelForCausalLM.from_pretrained(
-    qwen_checkpoint,
+qwen_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    checkpoint,
     torch_dtype=torch.bfloat16,
     device_map="auto",
+    # attn_implementation="flash_attention_2",
 )
 
+
+
 # LLaMA Model Setup
 llama_model_name = "Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2"
 llama_tokenizer = AutoTokenizer.from_pretrained(llama_model_name)