Spaces:
Runtime error
Runtime error
| import spaces # ZeroGPU: must be imported before torch | |
| import time | |
| import gradio as gr | |
| import torch | |
| from paddleocr import PaddleOCR | |
| from transformers import ( | |
| AutoModelForImageTextToText, | |
| AutoProcessor, | |
| Florence2ForConditionalGeneration, | |
| ) | |
| GOT_ID = "stepfun-ai/GOT-OCR-2.0-hf" # 0.6B, end-to-end OCR-2.0 VLM | |
| FLORENCE_ID = "microsoft/Florence-2-large" # 0.77B, unified vision model, <OCR> task | |
| # --------------------------------------------------------------------------- | |
| # Model loading (module level, on cuda — ZeroGPU optimizes startup placement) | |
| # --------------------------------------------------------------------------- | |
| print("Loading GOT-OCR-2.0-hf ...", flush=True) | |
| got_model = AutoModelForImageTextToText.from_pretrained( | |
| GOT_ID, dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map="cuda" | |
| ) | |
| got_processor = AutoProcessor.from_pretrained(GOT_ID, use_fast=True) | |
| print("Loading Florence-2-large ...", flush=True) | |
| try: | |
| florence_model = Florence2ForConditionalGeneration.from_pretrained( | |
| FLORENCE_ID, dtype=torch.float16 | |
| ) | |
| except Exception as e: # native class unavailable / load issue -> community pattern | |
| print(f"Native Florence2 load failed ({e}); falling back to trust_remote_code", flush=True) | |
| from transformers import AutoModelForCausalLM | |
| florence_model = AutoModelForCausalLM.from_pretrained( | |
| FLORENCE_ID, dtype=torch.float16, trust_remote_code=True | |
| ) | |
| florence_model = florence_model.to("cuda") | |
| florence_processor = AutoProcessor.from_pretrained( | |
| FLORENCE_ID, trust_remote_code=True | |
| ) | |
| print("Initializing PaddleOCR PP-OCRv5 (runs on CPU) ...", flush=True) | |
| paddle_ocr = PaddleOCR( | |
| text_detection_model_name="PP-OCRv5_server_det", | |
| text_recognition_model_name="PP-OCRv5_server_rec", | |
| use_doc_orientation_classify=False, | |
| use_doc_unwarping=False, | |
| use_textline_orientation=False, | |
| device="cpu", # paddle leg stays on CPU; ZeroGPU quota is for the VLMs | |
| ) | |
| print("All models ready.", flush=True) | |
| # --------------------------------------------------------------------------- | |
| # Inference | |
| # --------------------------------------------------------------------------- | |
| def run_vlm_ocr(image: Image.Image): | |
| """Run GOT-OCR-2.0 and Florence-2 on the same image; returns (got_text, florence_text).""" | |
| # --- GOT-OCR-2.0: plain page OCR (processor injects the "OCR: " prompt) --- | |
| got_inputs = got_processor(image, return_tensors="pt").to("cuda") | |
| with torch.inference_mode(): | |
| got_ids = got_model.generate( | |
| **got_inputs, | |
| do_sample=False, | |
| tokenizer=got_processor.tokenizer, # required for stop_strings | |
| stop_strings="<|im_end|>", # required, else generation runs to the cap | |
| max_new_tokens=4096, | |
| ) | |
| got_text = got_processor.decode( | |
| got_ids[0, got_inputs["input_ids"].shape[1]:], skip_special_tokens=True | |
| ) | |
| # --- Florence-2: <OCR> task --- | |
| task_prompt = "<OCR>" | |
| flo_inputs = florence_processor( | |
| text=task_prompt, images=image, return_tensors="pt" | |
| ).to("cuda", torch.float16) | |
| with torch.inference_mode(): | |
| flo_ids = florence_model.generate( | |
| **flo_inputs, max_new_tokens=1024, num_beams=3, do_sample=False | |
| ) | |
| flo_decoded = florence_processor.batch_decode(flo_ids, skip_special_tokens=False)[0] | |
| flo_parsed = florence_processor.post_process_generation( | |
| flo_decoded, task=task_prompt, image_size=(image.width, image.height) | |
| ) | |
| florence_text = flo_parsed[task_prompt] or "(no text detected)" | |
| return got_text.strip() or "(no text detected)", florence_text | |
| def run_paddle_ocr(image: Image.Image): | |
| """Run PaddleOCR PP-OCRv5 det+rec on CPU.""" | |
| result = paddle_ocr.predict(image) | |
| lines = [] | |
| for res in result: | |
| texts = res.get("rec_texts", []) if hasattr(res, "get") else res["rec_texts"] | |
| lines.extend(texts) | |
| return "\n".join(lines) if lines else "(no text detected)" | |
| def compare(image: Image.Image): | |
| if image is None: | |
| raise gr.Error("Please upload or paste an image first.") | |
| t0 = time.perf_counter() | |
| paddle_text = run_paddle_ocr(image) | |
| paddle_time = time.perf_counter() - t0 | |
| t1 = time.perf_counter() | |
| got_text, florence_text = run_vlm_ocr(image) | |
| vlm_time = time.perf_counter() - t1 | |
| timings = ( | |
| f"PaddleOCR PP-OCRv5: {paddle_time:.2f}s (CPU) | " | |
| f"GOT-OCR-2.0 + Florence-2: {vlm_time:.2f}s (GPU)" | |
| ) | |
| return paddle_text, got_text, florence_text, timings | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| DESCRIPTION = """ | |
| # OCR Model Comparison — 3 models under 3B | |
| Upload an image and run it through all three models at once: | |
| | Model | Params | Type | | |
| | --- | --- | --- | | |
| | [PaddleOCR PP-OCRv5](https://huggingface.co/PaddlePaddle/PP-OCRv5_server_det) | < 0.1B | classic det+rec pipeline (CPU) | | |
| | [GOT-OCR-2.0-hf](https://huggingface.co/stepfun-ai/GOT-OCR-2.0-hf) | 0.6B | end-to-end OCR VLM | | |
| | [Florence-2-large](https://huggingface.co/microsoft/Florence-2-large) | 0.77B | unified vision model, `<OCR>` task | | |
| Text is shown in reading order per model; lines may differ between models. | |
| """ | |
| with gr.Blocks(title="OCR Model Comparison") as demo: | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| image_in = gr.Image(type="pil", label="Input image", height=420) | |
| run_btn = gr.Button("Run all three models", variant="primary") | |
| with gr.Row(): | |
| paddle_out = gr.Textbox(label="PaddleOCR PP-OCRv5 (< 0.1B)", lines=14, show_copy_button=True) | |
| got_out = gr.Textbox(label="GOT-OCR-2.0-hf (0.6B)", lines=14, show_copy_button=True) | |
| florence_out = gr.Textbox(label="Florence-2-large (0.77B)", lines=14, show_copy_button=True) | |
| timings_out = gr.Textbox(label="Runtime", interactive=False) | |
| run_btn.click( | |
| fn=compare, | |
| inputs=image_in, | |
| outputs=[paddle_out, got_out, florence_out, timings_out], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |