Update app.py

app.py

@@ -1,11 +1,26 @@
 import gradio as gr
+import torch
 from transformers import AutoModel, AutoTokenizer
 from PIL import Image
 
+# Check GPU availability
+if torch.cuda.is_available():
+    print("CUDA is available! GPU is present.")
+    print(f"Number of GPUs: {torch.cuda.device_count()}")
+    print(f"GPU Name: {torch.cuda.get_device_name(0)}")
+else:
+    print("CUDA is not available. Running on CPU.")
+
 # Load the tokenizer and model
 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, pad_token_id=tokenizer.eos_token_id)
-model = model.eval().cuda()
+
+# Initialize the model
+if torch.cuda.is_available():
+    model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
+    model = model.eval().cuda()
+else:
+    model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, pad_token_id=tokenizer.eos_token_id)
+    model = model.eval()  # Keep model on CPU
 
 # Define the OCR function
 def perform_ocr(image):