Update app.py

app.py

@@ -47,6 +47,7 @@ ocr_model = ocr_predictor(
 
 
 if torch.cuda.is_available():
+    # Load the processor and model
     processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
     vision_model = LlavaNextForConditionalGeneration.from_pretrained(
         "llava-hf/llava-v1.6-mistral-7b-hf",
@@ -55,6 +56,11 @@ if torch.cuda.is_available():
         load_in_4bit=True,
     )
 
+    # Move the processor and model to the GPU
+    processor = processor.to("cuda")
+    vision_model = vision_model.to("cuda")
+
+
 
 @spaces.GPU
 def get_image_description(image):