fix: cuda device

app.py

@@ -4,15 +4,21 @@ import uuid
 
 import fitz  # PyMuPDF
 import gradio as gr
-from modelscope import AutoModel, AutoTokenizer
+import torch
 from PIL import Image, ImageEnhance
+from transformers import AutoConfig, AutoModel, AutoTokenizer
 
 from got_ocr import got_ocr
 
 # 初始化模型和分词器
-tokenizer = AutoTokenizer.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True)
-model = AutoModel.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True, low_cpu_mem_usage=True, device_map="cuda", use_safetensors=True)
-model = model.eval().cuda()
+model_name = "stepfun-ai/GOT-OCR2_0"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_name, trust_remote_code=True, low_cpu_mem_usage=True, device_map="cuda", use_safetensors=True)
+model = model.eval().to(device)
+model.config.pad_token_id = tokenizer.eos_token_id
 
 UPLOAD_FOLDER = "./uploads"
 RESULTS_FOLDER = "./results"