Update app.py

app.py

@@ -10,15 +10,20 @@ from unsloth import FastVisionModel
 # Define the model and processor
 model_id = "0llheaven/Llama-3.2-11B-Vision-Radiology-mini"
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
 # device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModelForImageTextToText.from_pretrained(
     model_id,
     # load_in_4bit=True,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-)
+    torch_dtype=torch.float32 if device.type == "cpu" else torch.bfloat16,
+    device_map="auto" if device.type == "cuda" else None,
+).to(device)
 
-model.gradient_checkpointing_enable()
+if device.type == "cuda":
+    model.gradient_checkpointing_enable()
+# model.gradient_checkpointing_enable()
 
 processor = AutoProcessor.from_pretrained(model_id)
 
@@ -53,7 +58,7 @@ def generate_description(image: Image.Image, instruction: str):
         input_text,
         add_special_tokens=False,
         return_tensors="pt"
-    ).to(model.device)
+    ).to(device)
 
     # Generate the output from the model
     # output = model.generate(**inputs, max_new_tokens=256)