Update app.py

app.py

@@ -4,14 +4,18 @@ from PIL import Image
 import fitz  # PyMuPDF
 import torch
 
-# 指定设备
+# 指定设备 (在免费Space上，这里会自动选择 'cpu')
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 
 # 加载模型和分词器
+# 首次加载会下载模型，可能需要很长时间
+print("Loading DeepSeek-OCR model...")
 model_path = 'deepseek-ai/DeepSeek-OCR'
 tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).to(device)
 model.eval()
+print("Model loaded successfully.")
 
 def pdf_to_images(pdf_path):
     """将PDF文件转换为PIL图像列表"""
@@ -19,7 +23,7 @@ def pdf_to_images(pdf_path):
     images = []
     for page_num in range(len(doc)):
         page = doc.load_page(page_num)
-        pix = page.get_pixmap()
+        pix = page.get_pixmap(dpi=200) # 适当降低dpi以减少内存消耗
         img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
         images.append(img)
     doc.close()
@@ -35,44 +39,38 @@ def ocr_process(pdf_file):
         images = pdf_to_images(pdf_path)
         
         full_text = ""
+        # 提示用户进程开始
+        yield "PDF处理完成，共 {} 页。开始逐页识别，请耐心等待...".format(len(images))
+
         for i, pil_img in enumerate(images):
-            # 准备模型输入
+            yield f"正在识别第 {i+1}/{len(images)} 页..."
+            
             messages = [
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "image", "image": pil_img},
-                        {"type": "text", "text": "Please perform OCR on this image."}
-                    ]
-                }
+                {"role": "user", "content": [{"type": "image", "image": pil_img}, {"type": "text", "text": "recognize characters in this image"}]}
             ]
             
-            text_input = tokenizer.apply_chat_template(
-                messages, add_generation_prompt=True, return_tensors="pt"
-            ).to(device)
+            text_input = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(device)
             
-            # 生成OCR结果
             outputs = model.generate(text_input, max_new_tokens=2048, do_sample=False)
             result_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
             
-            # 清理和提取识别出的文本
             # 简单的后处理，移除提示词部分
-            cleaned_text = result_text.split("Please perform OCR on this image.\n")[-1]
+            cleaned_text = result_text.split("recognize characters in this image")[-1].strip()
             
             full_text += f"--- Page {i+1} ---\n{cleaned_text}\n\n"
 
-        return full_text
+        yield full_text
 
     except Exception as e:
-        return f"处理时发生错误: {str(e)}"
+        yield f"处理时发生错误: {str(e)}"
 
 # 创建Gradio界面
 iface = gr.Interface(
     fn=ocr_process,
     inputs=gr.File(label="上传PDF文件", file_types=[".pdf"]),
-    outputs=gr.Textbox(label="识别结果", lines=20),
-    title="DeepSeek OCR PDF识别",
-    description="上传一个PDF文件，模型将识别其中的文本内容。免费CPU服务器处理较慢，请耐心等待。"
+    outputs=gr.Textbox(label="识别结果 (DeepSeek-OCR)", lines=20, show_copy_button=True),
+    title="DeepSeek OCR PDF识别 (CPU运行)",
+    description="上传PDF文件进行识别。警告：此模型在免费CPU服务器上运行会【极其缓慢】，处理多页或复杂PDF极有可能因超时而失败。"
 )
 
 # 启动应用