Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # must come before torch / transformers | |
| import time | |
| import torch | |
| from PIL import Image | |
| import gradio as gr | |
| from transformers import AutoModelForImageTextToText, AutoProcessor | |
| MODEL_ID = "moganai/MoganOCR-TR" | |
| PROMPT = "Sayfayı markdown olarak yaz." | |
| # The model card allows up to 8,000 new tokens; we default lower so a dense page | |
| # still fits comfortably inside the ZeroGPU time budget (greedy decoding). | |
| # Measured on ZeroGPU: ~37 tokens/s, densest bench page = 840 tokens in 23 s. | |
| MAX_NEW_TOKENS = 3072 | |
| MAX_ALLOWED_TOKENS = 6000 | |
| model = AutoModelForImageTextToText.from_pretrained( | |
| MODEL_ID, | |
| dtype=torch.bfloat16, | |
| attn_implementation="sdpa", | |
| ).to("cuda") | |
| processor = AutoProcessor.from_pretrained(MODEL_ID) | |
| model.eval() | |
| def _load_image(image): | |
| """Accept a Gradio image (PIL / numpy / path) and return an RGB PIL image.""" | |
| if image is None: | |
| raise ValueError("Please upload a page image first.") | |
| if isinstance(image, str): | |
| return Image.open(image).convert("RGB") | |
| if isinstance(image, Image.Image): | |
| return image.convert("RGB") | |
| return Image.fromarray(image).convert("RGB") | |
| def ocr_page( | |
| image, | |
| max_new_tokens: int = MAX_NEW_TOKENS, | |
| progress=gr.Progress(track_tqdm=True), | |
| ): | |
| """Transcribe one document page image to Markdown. | |
| Turkish OCR: text, headings and lists come out as Markdown, tables as | |
| HTML and formulas as LaTeX, in reading order. | |
| Args: | |
| image: the page image (scan, photo or screenshot). | |
| max_new_tokens: generation budget for the transcription. | |
| """ | |
| t0 = time.perf_counter() | |
| img = _load_image(image) | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": img}, | |
| {"type": "text", "text": PROMPT}, | |
| ], | |
| } | |
| ] | |
| inputs = processor.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| tokenize=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| ).to(model.device) | |
| with torch.inference_mode(): | |
| out = model.generate( | |
| **inputs, | |
| max_new_tokens=int(max_new_tokens), | |
| do_sample=False, | |
| ) | |
| generated = out[0, inputs["input_ids"].shape[1] :] | |
| text = processor.decode(generated, skip_special_tokens=True) | |
| print(f"[ocr_page] {max_new_tokens=} tokens_out={generated.shape[0]} " | |
| f"elapsed={time.perf_counter() - t0:.1f}s", flush=True) | |
| return text, text | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| """ | |
| # MoganOCR-TR 📄→️📝 | |
| **2B-parameter OCR model for Turkish documents.** Upload one page — a scan, photo or | |
| screenshot — and it writes the whole page in reading order as Markdown: | |
| text, headings and lists as Markdown, tables as HTML, formulas as LaTeX, | |
| running headers/footers as `<header>`/`<footer>` tags. | |
| [Model card](https://huggingface.co/moganai/MoganOCR-TR) · | |
| [TurkishOCR-Bench](https://huggingface.co/datasets/moganai/TurkishOCR-Bench) | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_in = gr.Image(label="Page image", type="pil", height=420) | |
| run_btn = gr.Button("Transcribe page", variant="primary") | |
| with gr.Accordion("Advanced settings", open=False): | |
| max_new_tokens = gr.Slider( | |
| minimum=512, | |
| maximum=MAX_ALLOWED_TOKENS, | |
| value=MAX_NEW_TOKENS, | |
| step=256, | |
| label="Max new tokens", | |
| info="Generation budget for the transcription. " | |
| "Denser pages need more tokens; each page gets " | |
| "about 60 s of GPU time.", | |
| ) | |
| with gr.Column(): | |
| markdown_out = gr.Markdown(label="Transcription (rendered)") | |
| with gr.Accordion("Raw Markdown", open=True): | |
| raw_out = gr.Textbox( | |
| show_label=False, | |
| lines=18, | |
| max_lines=40, | |
| interactive=False, | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/scanned_book_page.png"], | |
| ["examples/scanned_newspaper_page.png"], | |
| ["examples/invoice_form.jpg"], | |
| ["examples/slide.png"], | |
| ], | |
| inputs=[image_in], | |
| outputs=[markdown_out, raw_out], | |
| fn=ocr_page, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Examples (pages from TurkishOCR-Bench)", | |
| ) | |
| run_btn.click( | |
| fn=ocr_page, | |
| inputs=[image_in, max_new_tokens], | |
| outputs=[markdown_out, raw_out], | |
| api_name="transcribe", | |
| ) | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |