Spaces:
Running
Running
File size: 5,752 Bytes
5827499 af75cff 5827499 af75cff d45f3e7 5827499 8be5494 af75cff 5827499 89a1632 5827499 89a1632 af75cff afbaa03 5827499 afbaa03 5827499 8be5494 5827499 fff0f58 5827499 fff0f58 5827499 6ba101c 5827499 6ba101c 5827499 8d1fa76 2ac226e 5201e8a 6ba101c 5201e8a 2ac226e 6ba101c 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 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
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
print(f"Processing page {page_num}...") # Debugging line
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)
print(f"Anchor text for page {page_num}: {anchor_text}") # Debugging line
prompt = build_finetuning_prompt(anchor_text)
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:].detach().cpu()
decoded = "[No output generated]"
if new_tokens.shape[1] > 0:
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)}]"
else:
decoded = "[Model returned no new tokens]"
except Exception as processing_error:
decoded = f"[Processing error on page {page_num}: {str(processing_error)}]"
else:
try:
# Check if the tokens are empty
if not new_tokens:
decoded = f"[No tokens generated for page {page_num}]"
else:
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)}]"
print(f"Decoded content for page {page_num}: {decoded}") # Debugging line
# 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
)
|