Update app.py

app.py

@@ -9,11 +9,9 @@ import gradio as gr
 import argparse
 import tempfile
 
-import os
 from PIL import Image
-import json
-from tqdm import tqdm
 import easyocr
+import torch
 
 assert (
     __version__ == "4.32.0"
@@ -31,10 +29,14 @@ def get_easy_text(img_file):
     return out
 
 model_name = "DigitalAgent/Captioner"
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
 model = (
     AutoModelForCausalLM.from_pretrained(
-        model_name, device_map="cuda", trust_remote_code=True
-    )
+        model_name, trust_remote_code=True
+    ).to(device)
     .eval()
     .half()
 )