Update app.py

app.py

@@ -1,8 +1,11 @@
 import os
 import gradio as gr
+import torch
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForCausalLM 
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 #workaround for unnecessary flash_attn requirement
 from unittest.mock import patch
 from transformers.dynamic_module_utils import get_imports
@@ -16,7 +19,7 @@ def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
     return imports
 
 with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): #workaround for unnecessary flash_attn requirement
-    model = AutoModelForCausalLM.from_pretrained("Oysiyl/Florence-2-FT-OCR-Cauldron-IAM", attn_implementation="sdpa", trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained("Oysiyl/Florence-2-FT-OCR-Cauldron-IAM", attn_implementation="sdpa", trust_remote_code=True).to(device)
 
 processor = AutoProcessor.from_pretrained("Oysiyl/Florence-2-FT-OCR-Cauldron-IAM", trust_remote_code=True)