Refactor model initialization to always load on GPU

app.py

@@ -46,14 +46,15 @@ ocr_model = ocr_predictor(
 )
 
 
-if torch.cuda.is_available():
-    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",
-        torch_dtype=torch.float16,
-        low_cpu_mem_usage=True,
-        load_in_4bit=True,
-    )
+
+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",
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
+    load_in_4bit=True,
+)
+vision_model.to("cuda:0")
 
 
 @spaces.GPU