Spaces:
Sleeping
Sleeping
File size: 2,402 Bytes
64cf084 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | import io
import litserve as ls
from transformers import NougatProcessor, VisionEncoderDecoderModel
import torch
from PIL import Image
from fastapi import UploadFile
from pdf2image import convert_from_bytes
class NougatLitAPI(ls.LitAPI):
def setup(self, device):
model_name = "facebook/nougat-base"
self.processor = NougatProcessor.from_pretrained(model_name)
self.model = VisionEncoderDecoderModel.from_pretrained(model_name)
self.model.to(device)
self.model.eval()
def decode_request(self, request: UploadFile):
file_bytes = request.file.read()
# convert PDF bytes to images
try:
images = convert_from_bytes(file_bytes, dpi=300)
images = [img.convert("RGB") for img in images]
except Exception as e:
print(f"couldn't convert pdf to images: {e}")
images = []
if not images:
# fallback: blank white image if PDF fails
img = Image.new('RGB', (500, 500), color="white")
images.append(img)
print("warning: no images found in pdf, using blank fallback.")
return images
def predict(self, images):
markdown_output = []
for i, img in enumerate(images):
# preprocess single page
pixel_values = self.processor(images=img, return_tensors="pt").pixel_values
if torch.cuda.is_available():
pixel_values = pixel_values.to(self.model.device)
# generate markdown for page
with torch.no_grad():
outputs = self.model.generate(
pixel_values,
min_length=1,
max_new_tokens=2048, # enough for one page
bad_words_ids=[[self.processor.tokenizer.unk_token_id]],
)
decoded_text = self.processor.batch_decode(outputs, skip_special_tokens=True)[0]
decoded_text = self.processor.post_process_generation(decoded_text, fix_markdown=False)
markdown_output.append(f"\n\n--- Page {i + 1} ---\n\n{decoded_text}")
return "".join(markdown_output)
def encode_response(self, output_text):
return {"markdown_output": output_text}
if __name__ == "__main__":
api = NougatLitAPI()
server = ls.LitServer(api, accelerator="cuda", devices=1)
server.run(port=8000)
|