initial commit

app.py

@@ -6,23 +6,28 @@ import base64
 import spaces
 
 tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
+config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
 model = model.eval().cuda()
+model.config.pad_token_id = tokenizer.eos_token_id
 
 @spaces.GPU
 def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, render=False):
+    # Create attention mask
+    attention_mask = torch.ones((1, model.config.max_position_embeddings), dtype=torch.long, device=model.device)
+    
     if task == "Plain Text OCR":
-        res = model.chat(tokenizer, image, ocr_type='ocr')
+        res = model.chat(tokenizer, image, ocr_type='ocr', attention_mask=attention_mask)
     elif task == "Format Text OCR":
-        res = model.chat(tokenizer, image, ocr_type='format')
+        res = model.chat(tokenizer, image, ocr_type='format', attention_mask=attention_mask)
     elif task == "Fine-grained OCR (Box)":
-        res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_box=ocr_box)
+        res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_box=ocr_box, attention_mask=attention_mask)
     elif task == "Fine-grained OCR (Color)":
-        res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_color=ocr_color)
+        res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_color=ocr_color, attention_mask=attention_mask)
     elif task == "Multi-crop OCR":
-        res = model.chat_crop(tokenizer, image_file=image)
+        res = model.chat_crop(tokenizer, image_file=image, attention_mask=attention_mask)
     elif task == "Render Formatted OCR":
-        res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file='./demo.html')
+        res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file='./demo.html', attention_mask=attention_mask)
         with open('./demo.html', 'r') as f:
             html_content = f.read()
         return res, html_content