Update app.py

app.py

@@ -1,16 +1,15 @@
 import random
-import torch
 import spaces
 import gradio as gr
-from transformers import LlavaNextProcessor, LlavaNextForConditionalGeneration
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from huggingface_hub import hf_hub_download
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 model_path = "ibm-granite/granite-vision-3.1-2b-preview"
-processor = LlavaNextProcessor.from_pretrained(
-    model_path, 
-    use_fast=True,
-    image_processor_type = 'llava_next'
-)
-model = LlavaNextForConditionalGeneration.from_pretrained(model_path, torch_dtype="auto", device_map="auto")
+processor = AutoProcessor.from_pretrained(model_path)
+model = AutoModelForVision2Seq.from_pretrained(model_path).to(device)
 
 @spaces.GPU()
 def get_text_from_content(content):