Update app.py

app.py

@@ -11,12 +11,15 @@ 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"
+UPLOAD_FOLDER = "./uploads
 RESULTS_FOLDER = "./results"
 
 for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]: