Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers import AutoModelForVision2Seq, AutoProcessor
 from PIL import Image
 import numpy as np
 from fastapi import FastAPI, UploadFile, File
@@ -14,7 +14,7 @@ app = FastAPI()
 # Load SmolVLM-Instruct model and processor
 model_id = "HuggingFaceTB/SmolVLM-Instruct"
 processor = AutoProcessor.from_pretrained(model_id, token=os.environ.get("HF_TOKEN"))
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, token=os.environ.get("HF_TOKEN")).to("cuda")
+model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.bfloat16, token=os.environ.get("HF_TOKEN")).to("cuda")
 
 # Harmful objects list for detection
 harmful_objects = ["knife", "gun", "weapon", "blood", "syringe", "bomb", "blade"]
@@ -49,7 +49,8 @@ async def predict(files: List[UploadFile] = File(...)):
             # Get image embedding for similarity
             inputs_emb = processor(images=image, return_tensors="pt").to("cuda")
             with torch.no_grad():
-                emb = model.vision_tower(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
+                # Use vision tower for embeddings
+                emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
             image_embeddings.append(emb)
             
             results.append({
@@ -107,7 +108,7 @@ def gradio_predict(*images):
             # Get image embedding for similarity
             inputs_emb = processor(images=image, return_tensors="pt").to("cuda")
             with torch.no_grad():
-                emb = model.vision_tower(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
+                emb = model.vision_model(inputs_emb["pixel_values"]).last_hidden_state.mean(dim=1).cpu().numpy()
             image_embeddings.append(emb)
             
             results.append({