Spaces:

leonarb
/

olmocr-demo

Running

App Files Files Community

leonarb commited on 22 days ago

Commit

5827499

verified ·

1 Parent(s): b3d319d

Update app.py

Browse files

Files changed (1) hide show

app.py +93 -111

app.py CHANGED Viewed

@@ -1,137 +1,119 @@
-import os
 import base64
-import tempfile
 from io import BytesIO
-import torch
-import gradio as gr
 from PIL import Image
-from PyPDF2 import PdfReader
-from ebooklib import epub
-from pdf2image import convert_from_path
 from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
-# Set cache and log paths
-cache_dir = "/tmp/huggingface_cache"
-os.environ["HF_HOME"] = cache_dir
-os.environ["TORCH_HOME"] = cache_dir
-os.makedirs(cache_dir, exist_ok=True)
-# Patch logging to avoid permission errors
-import logging
-from logging import FileHandler
-class SafeFileHandler(FileHandler):
-    def __init__(self, filename, mode='a', encoding=None, delay=False, errors=None):
-        # Redirect all logs to tmp
-        safe_path = os.environ.get("OLMOCR_LOG_PATH", "/tmp/olmocr-pipeline-debug.log")
-        super().__init__(safe_path, mode, encoding, delay, errors)
-logging.FileHandler = SafeFileHandler
-# Now import olmocr
-from olmocr.run_ocr import ocr_pdf_to_text
 from olmocr.prompts import build_finetuning_prompt
 from olmocr.prompts.anchor import get_anchor_text
-from olmocr.data.renderpdf import render_pdf_to_base64png
 # Load model and processor
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = Qwen2VLForConditionalGeneration.from_pretrained(
-    "allenai/olmOCR-7B-0225-preview",
-    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
-).eval().to(device)
 processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
-def ocr_page(pdf_path, page_num):
-    image_b64 = render_pdf_to_base64png(pdf_path, page_num + 1, target_longest_image_dim=1024)
-    anchor_text = get_anchor_text(pdf_path, page_num + 1, pdf_engine="pdfreport", target_length=4000)
-    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_b64}"}}
-        ],
-    }]
-    prompt_text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    main_image = Image.open(BytesIO(base64.b64decode(image_b64)))
-    inputs = processor(text=[prompt_text], images=[main_image], return_tensors="pt", padding=True)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            temperature=0.8,
-            max_new_tokens=1024,
-            do_sample=True,
-        )
-    prompt_len = inputs["input_ids"].shape[1]
-    new_tokens = outputs[:, prompt_len:]
-    decoded = processor.tokenizer.batch_decode(new_tokens, skip_special_tokens=True)
-    return decoded[0] if decoded else ""
-def create_epub_from_text(text, output_path, title, author, language, cover_image):
     book = epub.EpubBook()
     book.set_title(title)
-    book.set_language(language)
     book.add_author(author)
-    with open(cover_image, "rb") as cover_file:
-        cover_data = cover_file.read()
-    cover_item = epub.EpubItem(uid="cover", file_name="cover.jpg", media_type="image/jpeg", content=cover_data)
-    book.add_item(cover_item)
-    chapter = epub.EpubHtml(title="Content", file_name="content.xhtml", lang=language)
-    chapter.set_content(f"<html><body><h1>{title}</h1><p>{text}</p></body></html>")
-    book.add_item(chapter)
-    book.toc = (epub.Link("content.xhtml", "Content", "content"),)
     book.add_item(epub.EpubNav())
     epub.write_epub(output_path, book)
-def convert_pdf_to_epub(pdf_file, title, author, language):
-    tmp_pdf_path = pdf_file.name
-    reader = PdfReader(tmp_pdf_path)
-    cover_path = "/tmp/cover.jpg"
-    images = convert_from_path(tmp_pdf_path, first_page=1, last_page=1)
-    images[0].save(cover_path, "JPEG")
-    # Use official AllenAI OCR function
-    ocr_text = ocr_pdf_to_text(
-        pdf_path=tmp_pdf_path,
-        model=model,
-        processor=processor
-    )
-    epub_path = "/tmp/output.epub"
-    create_epub_from_text(
-        text=ocr_text,
-        output_path=epub_path,
-        title=title,
-        author=author,
-        language=language,
-        cover_image=cover_path
-    )
-    return epub_path, cover_path
-def interface_fn(pdf, title, author, language):
-    epub_path, _ = convert_pdf_to_epub(pdf, title, author, language)
-    return epub_path
-demo = gr.Interface(
-    fn=interface_fn,
     inputs=[
         gr.File(label="Upload PDF", file_types=[".pdf"]),
-        gr.Textbox(label="EPUB Title", placeholder="e.g. Understanding AI"),
-        gr.Textbox(label="Author", placeholder="e.g. Allen AI"),
-        gr.Textbox(label="Language", placeholder="e.g. en", value="en"),
     ],
     outputs=gr.File(label="Download EPUB"),
-    title="PDF to EPUB Converter (olmOCR)",
-    description="Upload a PDF to convert it into a structured EPUB. The first page is used as the cover. OCR is performed with the olmOCR model.",
-    allow_flagging="never",
 )
 if __name__ == "__main__":
-    demo.launch(share=True)

+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:]
+            decoded = processor.tokenizer.batch_decode(new_tokens, skip_special_tokens=True)[0]
+        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."
 )
 if __name__ == "__main__":
+    iface.launch()