Update app.py

app.py

@@ -5,6 +5,9 @@ import spaces
 import torch
 import os
 
+from transformers import AutoProcessor, AutoModelForCausalLM 
+
+
 
 access_token = os.getenv('HF_token')
 model_id = "selamw/BirdWatcher2"
@@ -41,11 +44,19 @@ def convert_to_markdown(input_text):
     
 @spaces.GPU
 def infer_fin_pali(image, question):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
-    model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, quantization_config=bnb_config, token=access_token)
-    processor = PaliGemmaProcessor.from_pretrained(model_id, token=access_token)
+    # model = PaliGemmaForConditionalGeneration.from_pretrained(model_id, quantization_config=bnb_config, token=access_token)
+    # processor = PaliGemmaProcessor.from_pretrained(model_id, token=access_token)
+
 
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch_dtype, trust_remote_code=True, token=access_token).to(device)
+    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, token=access_token)
+###
+    
     inputs = processor(images=image, text=question, return_tensors="pt").to(device)
 
     predictions = model.generate(**inputs, max_new_tokens=512)