Update app.py

app.py

@@ -47,18 +47,17 @@ ocr_model = ocr_predictor(
 
 
 if torch.cuda.is_available():
-    # Load the processor and model
+    # Load the processor (no need to move to GPU)
     processor = LlavaNextProcessor.from_pretrained("llava-hf/llava-v1.6-mistral-7b-hf")
+
+    # Load the model and move it to GPU
     vision_model = LlavaNextForConditionalGeneration.from_pretrained(
         "llava-hf/llava-v1.6-mistral-7b-hf",
         torch_dtype=torch.float16,
         low_cpu_mem_usage=True,
         load_in_4bit=True,
-    )
+    ).to("cuda")  # Move the model to GPU
 
-    # Move the processor and model to the GPU
-    processor = processor.to("cuda")
-    vision_model = vision_model.to("cuda")