Update vision_analyzer.py

vision_analyzer.py

@@ -21,6 +21,17 @@ import numpy as np
 from PIL import Image
 import torch
 from transformers import AutoProcessor, AutoModelForCausalLM
+# from transformers import AutoProcessor, AutoModelForCausalLM
+
+# Attempt to install flash-attn
+try:
+    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, check=True, shell=True)
+except subprocess.CalledProcessError as e:
+    print(f"Error installing flash-attn: {e}")
+    print("Continuing without flash-attn.")
+
+# Determine the device to use
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Initialize FastAPI
 app = FastAPI()
@@ -79,27 +90,14 @@ import sys
 device = "cpu" # Explicitly ensure CPU usage
 
 try:
-    # Load processor with padding configuration
-    vision_processor = AutoProcessor.from_pretrained(
-        "microsoft/git-base-coco",
-        padding="max_length",
-        truncation=True
-    )
-    
-    # Load model with CPU optimizations
-    vision_model = AutoModelForCausalLM.from_pretrained(
-        "microsoft/git-base-coco",
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True,
-        device_map="cpu"
-    ).eval()
-    
-    print("✅ Successfully loaded GIT model and processor")
-    
+    vision_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True, 
+            attn_implementation="eager" ).to(device).eval()
+    vision_processor = AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True)
 except Exception as e:
-    print(f"❌ Error loading model: {str(e)}")
-    vision_model = None
-    vision_processor = None
+    print(f"Error loading base model: {e}")
+    vision_language_model_base = None
+    vision_language_processor_base = None
+
 
 # Preprompt templates
 PREPROMPT_TEMPLATES = {