Update app.py

app.py

@@ -11,9 +11,12 @@ import tempfile
 import time
 import shutil
 from pathlib import Path
+import torch
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
-model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True)
+model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map=DEVICE, use_safetensors=True)
 model = model.eval().cuda()
 
 UPLOAD_FOLDER = "./uploads"