Spaces:
Running
Running
File size: 4,734 Bytes
5827499 af75cff 5827499 af75cff d45f3e7 5827499 8be5494 af75cff 5827499 89a1632 5827499 89a1632 af75cff afbaa03 5827499 afbaa03 5827499 8be5494 5827499 fff0f58 5827499 fff0f58 5827499 2ac226e 5827499 fff0f58 5827499 fff0f58 5827499 fff0f58 5827499 8be5494 afbaa03 5827499 8be5494 f01e8a4 5827499 99e3331 d45f3e7 9d46b18 |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 |
import gradio as gr
import torch
import base64
import fitz # PyMuPDF
from io import BytesIO
from PIL import Image
from pathlib import Path
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
from olmocr.data.renderpdf import render_pdf_to_base64png
from olmocr.prompts import build_finetuning_prompt
from olmocr.prompts.anchor import get_anchor_text
from ebooklib import epub
# Load model and processor
model = Qwen2VLForConditionalGeneration.from_pretrained(
"allenai/olmOCR-7B-0225-preview", torch_dtype=torch.bfloat16
).eval()
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
def process_pdf_to_epub(pdf_file, title, author):
pdf_path = pdf_file.name
doc = fitz.open(pdf_path)
num_pages = len(doc)
# Create EPUB book
book = epub.EpubBook()
book.set_identifier("id123456")
book.set_title(title)
book.add_author(author)
chapters = []
for i in range(num_pages):
page_num = i + 1
try:
# Render page to base64 image
image_base64 = render_pdf_to_base64png(pdf_path, page_num, target_longest_image_dim=1024)
anchor_text = get_anchor_text(pdf_path, page_num, pdf_engine="pdfreport", target_length=4000)
prompt = build_finetuning_prompt(anchor_text)
# Format prompt
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
],
}
]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image = Image.open(BytesIO(base64.b64decode(image_base64)))
inputs = processor(
text=[text],
images=[image],
padding=True,
return_tensors="pt",
)
inputs = {k: v.to(device) for k, v in inputs.items()}
output = model.generate(
**inputs,
temperature=0.8,
max_new_tokens=512,
num_return_sequences=1,
do_sample=True,
)
prompt_length = inputs["input_ids"].shape[1]
new_tokens = output[:, prompt_length:]
try:
decoded_list = processor.tokenizer.batch_decode(new_tokens, skip_special_tokens=True)
decoded = decoded_list[0].strip() if decoded_list else "[No output generated]"
except Exception as decode_error:
decoded = f"[Decoding error on page {page_num}: {str(decode_error)}]"
except Exception as e:
decoded = f"[Error processing page {page_num}: {str(e)}]"
# Create chapter
chapter = epub.EpubHtml(title=f"Page {page_num}", file_name=f"page_{page_num}.xhtml", lang="en")
chapter.content = f"<h1>Page {page_num}</h1><p>{decoded}</p>"
book.add_item(chapter)
chapters.append(chapter)
# Save cover image from page 1
if page_num == 1:
cover_image = Image.open(BytesIO(base64.b64decode(image_base64)))
cover_io = BytesIO()
cover_image.save(cover_io, format='PNG')
book.set_cover("cover.png", cover_io.getvalue())
# Assemble EPUB
book.toc = tuple(chapters)
book.add_item(epub.EpubNcx())
book.add_item(epub.EpubNav())
book.spine = ['nav'] + chapters
output_path = "/tmp/output.epub"
epub.write_epub(output_path, book)
return output_path
# Gradio Interface
iface = gr.Interface(
fn=process_pdf_to_epub,
inputs=[
gr.File(label="Upload PDF", file_types=[".pdf"]),
gr.Textbox(label="EPUB Title"),
gr.Textbox(label="Author(s)")
],
outputs=gr.File(label="Download EPUB"),
title="PDF to EPUB Converter (with olmOCR)",
description="Uploads a PDF, extracts text from each page with vision + prompt, and builds an EPUB using the outputs. Sets the first page as cover.",
allow_flagging="never" # Add this line to avoid the flagged directory issue
)
if __name__ == "__main__":
iface.launch(
server_name="0.0.0.0", # Required to make app publicly accessible
server_port=7860, # Can be changed if needed
share=True, # Optional: creates a public Gradio link if supported
debug=True, # Optional: helpful if you're troubleshooting
allowed_paths=["/tmp"] # Optional: makes it explicit that Gradio can write here
)
|