fix: cuda device

app-ocr.py

@@ -6,9 +6,6 @@ import time
 import uuid
 from pathlib import Path
 
-# import numpy as np
-# import tempfile
-# from PIL import Image
 import gradio as gr
 from modelscope import AutoModel, AutoTokenizer
 

app.py

@@ -11,7 +11,8 @@ from transformers import AutoConfig, AutoModel, AutoTokenizer
 from got_ocr import got_ocr
 
 # 初始化模型和分词器
-model_name = "stepfun-ai/GOT-OCR2_0"
+# model_name = "stepfun-ai/GOT-OCR2_0"
+model_name = "ucaslcl/GOT-OCR2_0"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)

got_ocr.py

@@ -39,12 +39,18 @@ def got_ocr(model, tokenizer, image_path, got_mode="format texts OCR", fine_grai
 
 # 使用示例
 if __name__ == "__main__":
-    from modelscope import AutoModel, AutoTokenizer
+    import torch
+    from transformers import AutoConfig, AutoModel, AutoTokenizer
 
     # 初始化模型和分词器
-    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
 
     image_path = "path/to/your/image.png"
     result, html = got_ocr(model, tokenizer, image_path, got_mode="format texts OCR")