diff --git a/README.md b/README.md index 16c6b4a..fafc217 100644 --- a/README.md +++ b/README.md @@ -58,9 +58,21 @@ py -3.13 -m venv .venv 2. Install Python dependencies (this installs PyTorch as well): ```bash -pip install -r requirements.txt +python -m pip install -r requirements.txt ``` +Keep the virtual environment active when running `python main.py`. If you see +`ModuleNotFoundError: No module named 'latex2mathml'`, your Python may be using +another environment. On Linux/Mac, run with the project interpreter directly: + +```bash +.venv/bin/python main.py ~/scan-pages/Book1.pdf +``` + +If that interpreter also reports a missing dependency, install the full +requirements into `.venv`. For environments created with uv, use +`uv pip install --python .venv/bin/python -r requirements.txt`. + 3. GPU acceleration (optional): PyTorch is installed as a dependency in step 2. On Apple Silicon that wheel @@ -132,7 +144,8 @@ python main.py [input_path] [output_path] [options] Options: --max-pages INT Maximum number of pages to process - --start-page INT Page number to start from + --start-page INT Zero-based page index to start from + --chunk-pages INT Pages processed at once (default: 10; 0 disables chunking) --skip-epub Skip EPUB generation, only create markdown --skip-md Skip markdown generation, use existing markdown files ``` @@ -151,6 +164,26 @@ Convert to markdown only: python main.py thesis.pdf --skip-epub ``` +### GPU memory usage + +PDFs are processed in 10-page chunks by default, then joined into one Markdown +file and one EPUB. This bounds the page-image memory used by Marker; users do +not need to split or reassemble the PDF themselves. On an 8 GB GPU, the default +also sets all of Marker's layout, detection, OCR, recognition, table, and +equation batch sizes to 1. + +If a particularly image-heavy PDF still runs out of memory, use a smaller +chunk (for example `--chunk-pages 4`). `PDF2EPUB_GPU_BATCH_SIZE` can raise the +model batch sizes on GPUs with more memory. `DETECTOR_BATCH_SIZE` alone does not +control Marker's internal detection batch in the pinned Marker version. + +PyTorch's allocator variable is named `PYTORCH_CUDA_ALLOC_CONF` (including +`CUDA`), for example: + +```bash +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python main.py book.pdf +``` + ### Output Structure ``` diff --git a/main.py b/main.py index 14a0f74..71ba85e 100755 --- a/main.py +++ b/main.py @@ -41,7 +41,16 @@ def main(): '--start-page', type=int, default=None, - help='Page number to start from' + help='Zero-based page index to start from' + ) + parser.add_argument( + '--chunk-pages', + type=int, + default=pdf2md.DEFAULT_CHUNK_PAGES, + help=( + 'Pages processed at once (default: 10; lower this if GPU memory is ' + 'still exhausted, or use 0 to disable chunking)' + ) ) parser.add_argument( '--skip-epub', @@ -93,6 +102,7 @@ def main(): markdown_dir, args.max_pages, args.start_page, + args.chunk_pages, ) # Convert Markdown to EPUB unless skipped @@ -115,4 +125,4 @@ def main(): sys.exit(1) if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/modules/pdf2md.py b/modules/pdf2md.py index 7e15469..d838a00 100644 --- a/modules/pdf2md.py +++ b/modules/pdf2md.py @@ -1,172 +1,207 @@ from pathlib import Path -import sys +import gc +import io import json +import os +import sys + from PIL import Image -import io + + +DEFAULT_CHUNK_PAGES = 10 +DEFAULT_GPU_BATCH_SIZE = 1 + def get_default_output_dir(input_path: Path) -> Path: - """ - Generate default output directory path based on input PDF path. - Creates a directory with same name as PDF (without extension) next to the PDF. - """ + """Return the default output directory for a PDF.""" return input_path.parent / input_path.stem + def get_default_input_dir() -> Path: - """ - Get default input directory (./input) relative to current working directory. - Creates it if it doesn't exist. - """ - input_dir = Path.cwd() / 'input' + """Return (and create) the default input directory.""" + input_dir = Path.cwd() / "input" input_dir.mkdir(exist_ok=True) return input_dir def save_images(images: dict, image_dir: Path) -> None: - """ - Save images with proper error handling and format detection. - Preserves original image filenames from the input. - - Args: - images: Dictionary of images from marker-pdf conversion where keys are filenames - image_dir: Directory to save images to - """ + """Save images returned by Marker and close their PIL resources.""" if not images: - print("No images found in document") return - - image_dir.mkdir(exist_ok=True) - saved_count = 0 - + image_dir.mkdir(parents=True, exist_ok=True) for filename, image_data in images.items(): try: - # Skip if image data is None or empty - if not image_data: - continue - image_path = image_dir / filename - - # Handle different image data formats + image_path.parent.mkdir(parents=True, exist_ok=True) if isinstance(image_data, Image.Image): image_data.save(image_path) - saved_count += 1 - elif isinstance(image_data, bytes): - img = Image.open(io.BytesIO(image_data)) - img.save(image_path) - saved_count += 1 - - elif isinstance(image_data, str): - if Path(image_data).exists(): - img = Image.open(image_data) - img.save(image_path) - saved_count += 1 - else: - print(f"Image path does not exist: {image_data}") + with Image.open(io.BytesIO(image_data)) as image: + image.save(image_path) + elif isinstance(image_data, str) and Path(image_data).exists(): + with Image.open(image_data) as image: + image.save(image_path) else: - print(f"Unsupported image data type for {filename}: {type(image_data)}") - - except Exception as e: - print(f"Error saving image {filename}: {str(e)}") - continue - - if saved_count > 0: - print(f"Successfully saved {saved_count} images to: {image_dir}") - else: - print("No valid images were found to save") + print(f"Unsupported image data for {filename}: {type(image_data)}") + except Exception as exc: + print(f"Error saving image {filename}: {exc}") + finally: + if isinstance(image_data, Image.Image): + image_data.close() + + +def _page_count(input_path: str) -> int: + """Read a PDF's page count without rendering its pages.""" + from pypdfium2 import PdfDocument + + document = PdfDocument(input_path) + try: + return len(document) + finally: + document.close() + + +def _selected_pages(page_count: int, start_page: int | None, max_pages: int | None) -> list[int]: + start = 0 if start_page is None else start_page + if start < 0: + raise ValueError("start_page must be zero or greater") + if max_pages is not None and max_pages < 1: + raise ValueError("max_pages must be greater than zero") + if start >= page_count: + raise ValueError(f"start_page {start} is outside this {page_count}-page PDF") + stop = page_count if max_pages is None else min(page_count, start + max_pages) + return list(range(start, stop)) + + +def _chunks(items: list[int], size: int) -> list[list[int]]: + if size <= 0: + return [items] + return [items[index:index + size] for index in range(0, len(items), size)] + + +def _low_memory_config(page_range: list[int]) -> dict: + """Set Marker's own batch controls; Surya environment variables are insufficient.""" + batch_size = int(os.environ.get("PDF2EPUB_GPU_BATCH_SIZE", DEFAULT_GPU_BATCH_SIZE)) + if batch_size < 1: + raise ValueError("PDF2EPUB_GPU_BATCH_SIZE must be greater than zero") + return { + "page_range": page_range, + "layout_batch_size": batch_size, + "detection_batch_size": batch_size, + "ocr_error_batch_size": batch_size, + "recognition_batch_size": batch_size, + "table_rec_batch_size": batch_size, + "equation_batch_size": batch_size, + } + + +def _merge_metadata(all_metadata: list[dict]) -> dict: + """Combine Marker metadata from page chunks into one document-level object.""" + if not all_metadata: + return {} + merged = dict(all_metadata[0]) + merged["page_stats"] = [ + stat for metadata in all_metadata for stat in metadata.get("page_stats", []) + ] + merged["table_of_contents"] = [ + entry + for metadata in all_metadata + for entry in metadata.get("table_of_contents", []) + ] + merged["processing_chunks"] = len(all_metadata) + return merged + + +def _release_chunk(rendered) -> None: + images = getattr(rendered, "images", {}) + for image in images.values(): + if isinstance(image, Image.Image): + image.close() + images.clear() + gc.collect() + try: + import torch + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except ImportError: + pass + def convert_pdf( input_path: str, output_dir: Path, - max_pages: int = None, - start_page: int = None, + max_pages: int | None = None, + start_page: int | None = None, + chunk_pages: int = DEFAULT_CHUNK_PAGES, ) -> None: - """ - Convert a single PDF file to markdown format with enhanced image handling. - """ + """Convert a PDF in bounded page chunks and transparently join the result.""" try: - from marker.models import create_model_dict from marker.converters.pdf import PdfConverter + from marker.models import create_model_dict + pages = _selected_pages(_page_count(input_path), start_page, max_pages) + page_chunks = _chunks(pages, chunk_pages) + output_dir.mkdir(parents=True, exist_ok=True) + image_dir = output_dir / "images" + per_chunk = chunk_pages if chunk_pages > 0 else len(pages) + print( + f"Processing {len(pages)} page(s) in {len(page_chunks)} chunk(s) " + f"of up to {per_chunk} page(s)." + ) + + # Models are expensive but fixed-size. Reuse them while discarding each + # chunk's rendered pages and tensors before starting the next chunk. models = create_model_dict() - - # Build page_range config. marker-pdf 1.x requires an explicit list of - # page indices; there is no "start to end" shorthand. If start_page is - # given without max_pages we cannot construct the range without knowing - # the document's page count, so we warn and process all pages instead. - config = {} - if max_pages is not None: - s = start_page or 0 - config["page_range"] = list(range(s, s + max_pages)) - elif start_page is not None: + markdown_parts = [] + metadata_parts = [] + for chunk_number, page_range in enumerate(page_chunks, start=1): print( - "Warning: --start-page requires --max-pages in marker-pdf 1.x " - "(no open-ended page range is supported). Processing all pages." + f"Converting chunk {chunk_number}/{len(page_chunks)} " + f"(pages {page_range[0] + 1}-{page_range[-1] + 1})..." ) - - converter = PdfConverter(config=config, artifact_dict=models) - rendered = converter(input_path) - - full_text = rendered.markdown - images = rendered.images - metadata = rendered.metadata - - # All output will go to the output directory - output_dir.mkdir(parents=True, exist_ok=True) - - # Save markdown content - md_output = output_dir / f"{Path(input_path).stem}.md" - md_output.write_text(full_text, encoding='utf-8') + converter = PdfConverter( + config=_low_memory_config(page_range), artifact_dict=models + ) + rendered = converter(input_path) + markdown_parts.append(rendered.markdown.strip()) + metadata_parts.append(rendered.metadata) + save_images(rendered.images, image_dir) + del converter + _release_chunk(rendered) + del rendered + + stem = Path(input_path).stem + md_output = output_dir / f"{stem}.md" + md_output.write_text("\n\n".join(markdown_parts) + "\n", encoding="utf-8") print(f"Markdown saved to: {md_output}") - - # Save metadata as JSON - meta_output = output_dir / f"{Path(input_path).stem}_metadata.json" - with open(meta_output, 'w', encoding='utf-8') as f: - json.dump(metadata, f, indent=2) + meta_output = output_dir / f"{stem}_metadata.json" + meta_output.write_text( + json.dumps(_merge_metadata(metadata_parts), indent=2), encoding="utf-8" + ) print(f"Metadata saved to: {meta_output}") + if image_dir.exists(): + print(f"Images saved to: {image_dir}") - # Enhanced image handling - try: - if images: - image_dir = output_dir / "images" - save_images(images, image_dir) - - # Cleanup PIL Images - for img in images.values(): - if isinstance(img, Image.Image): - try: - img.close() - except Exception as e: - print(f"Warning: Failed to close image: {e}") - images.clear() - except Exception as e: - print(f"Warning: Error during image cleanup: {e}") - - except Exception as e: - print(f"Error converting {input_path}: {str(e)}", file=sys.stderr) + del models + gc.collect() + except Exception as exc: + print(f"Error converting {input_path}: {exc}", file=sys.stderr) raise - + def add_pdfs_to_queue(input_path: Path) -> list[Path]: - """ - Add PDF files to the processing queue. - If input_path is a directory, add all PDFs in it. - If input_path is a file, add just that file. - """ - queue = [] - + """Return PDFs represented by an input file or directory.""" if input_path.is_dir(): - pdfs = list(input_path.glob('*.pdf')) + pdfs = sorted(input_path.glob("*.pdf")) if not pdfs: print(f"No PDF files found in directory: {input_path}", file=sys.stderr) sys.exit(1) - queue.extend(pdfs) - else: - if not input_path.is_file(): - print(f"Error: Input file does not exist: {input_path}", file=sys.stderr) - sys.exit(1) - if input_path.suffix.lower() != '.pdf': - print(f"Error: Input file must be a PDF: {input_path}", file=sys.stderr) - sys.exit(1) - queue.append(input_path) - - return queue + return pdfs + if not input_path.is_file(): + print(f"Error: Input file does not exist: {input_path}", file=sys.stderr) + sys.exit(1) + if input_path.suffix.lower() != ".pdf": + print(f"Error: Input file must be a PDF: {input_path}", file=sys.stderr) + sys.exit(1) + return [input_path]