Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import streamlit as st
 from transformers import AutoModel, AutoTokenizer
 from PIL import Image
@@ -7,6 +8,9 @@ import uuid
 import time
 from pathlib import Path
 
+# Force the use of CPU
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
 # Load tokenizer and model on 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).to("cpu")
@@ -73,7 +77,6 @@ def cleanup_old_files():
 # Streamlit App
 st.set_page_config(page_title="GOT-OCR-2.0 Demo", layout="wide")
 
-
 uploaded_image = st.file_uploader("Upload your image", type=["png", "jpg", "jpeg"])
 
 if uploaded_image: