Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import streamlit as st
 from transformers import AutoModel, AutoTokenizer
 from PIL import Image
@@ -6,12 +5,11 @@ import os
 import base64
 import uuid
 import time
-import shutil
 from pathlib import Path
 
 # 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)
+model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True).to("cpu")
 model = model.eval()  # Use CPU
 
 # Define folders for uploads and results
@@ -75,21 +73,6 @@ def cleanup_old_files():
 # Streamlit App
 st.set_page_config(page_title="GOT-OCR-2.0 Demo", layout="wide")
 
-st.markdown("""
-<h2> <span style="color: #ff6600">General OCR Theory</span>: Towards OCR-2.0 via a Unified End-to-end Model</h2>
-<a href="https://huggingface.co/ucaslcl/GOT-OCR2_0">[😊 Hugging Face]</a> 
-<a href="https://arxiv.org/abs/2409.01704">[📜 Paper]</a> 
-<a href="https://github.com/Ucas-HaoranWei/GOT-OCR2.0/">[🌟 GitHub]</a> 
-""", unsafe_allow_html=True)
-
-st.markdown("""
-"🔥🔥🔥This is the official online demo of the GOT-OCR-2.0 model!!!"
-### Demo Guidelines
-- You need to upload your image below and choose one mode of GOT, then click "Submit" to run the GOT model. More characters will result in longer wait times.
-- **plain texts OCR & format texts OCR**: The two modes are for the image-level OCR.
-- **plain multi-crop OCR & format multi-crop OCR**: For images with more complex content, you can achieve higher-quality results with these modes.
-- **plain fine-grained OCR & format fine-grained OCR**: In these modes, you can specify fine-grained regions on the input image for more flexible OCR. Fine-grained regions can be coordinates of the box, red color, blue color, or green color.
-""")
 
 uploaded_image = st.file_uploader("Upload your image", type=["png", "jpg", "jpeg"])