moebiusT7/book-ocr-studio

1"""Dedicated process: release Marker CUDA allocations at exit."""2import os, sys, time3from pathlib import Path4os.environ.setdefault('TORCH_DEVICE','cuda')5os.environ.setdefault('RECOGNITION_BATCH_SIZE','4')6os.environ.setdefault('DETECTOR_BATCH_SIZE','2')7from core import read, save8 9def run(job):10    from marker.converters.pdf import PdfConverter11    from marker.models import create_model_dict12    from marker.output import save_output13    cfg=read(job/'job.json')14    todo=[job/f'page-{idx+1:05d}' for idx in cfg['selected'] if not (job/f'page-{idx+1:05d}'/'marker.json').exists()]15    if not todo: return16    load_start=time.time()17    converter=PdfConverter(artifact_dict=create_model_dict(),config={'force_ocr':cfg['force_ocr'],'use_llm':False})18    save(job/'marker-load.json',dict(start=load_start,end=time.time(),gpu=os.environ.get('CUDA_VISIBLE_DEVICES')))19    for folder in todo:20        # Bound the OCR lead to four pages while the second GPU consumes them.21        parallel=(job/'execution.json').exists() and read(job/'execution.json')['mode']=='dual'22        while parallel and cfg['gemma']:23            pending=sum((p/'marker.json').exists() and not ((p/'review.json').exists() or (p/'review-error.json').exists()) for p in job.glob('page-*'))24            if pending<4 or (job/'cancel').exists(): break25            time.sleep(1)26        if (job/'cancel').exists(): break27        print('Marker:',folder.name,flush=True)28        source=folder/'source.pdf' if (folder/'source.pdf').exists() else folder/'source.png'29        started=time.time()30        result=converter(str(source))31        save_output(result,str(folder),'original')32        if not (folder/'original.md').exists(): raise RuntimeError('Marker output not found')33        save(folder/'marker.json',dict(engine='marker-pdf',version='1.10.2',source=source.name,start=started,end=time.time(),gpu=os.environ.get('CUDA_VISIBLE_DEVICES')))34if __name__=='__main__': run(Path(sys.argv[1]).resolve())35