Update vision_model.py

vision_model.py

@@ -1,17 +1,19 @@
 from transformers import AutoProcessor, AutoModelForVision2Seq
 from PIL import Image
 import torch
-
 import os
 os.environ["HF_HOME"] = "/app/.cache"
 
-
 model_id = "llava-hf/llava-1.5-7b-hf"
 processor = AutoProcessor.from_pretrained(model_id)
-model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16).cuda()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = AutoModelForVision2Seq.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+).to(device)
 
 def describe_image(image_path, prompt="Describe this sculpt and any missing regions."):
     image = Image.open(image_path).convert("RGB")
-    inputs = processor(prompt, image, return_tensors="pt").to("cuda", torch.float16)
+    inputs = processor(prompt, image, return_tensors="pt").to(device, model.dtype)
     output = model.generate(**inputs, max_new_tokens=150)
-    return processor.decode(output[0], skip_special_tokens=True)
+    return processor.decode(output[0], skip_special_tokens=True)