IFMedTechdemo/Multi-Model-OCR

1import os2import time3import torch4import spaces5import warnings6import tempfile7import sys8from io import StringIO9from contextlib import contextmanager10from threading import Thread11from PIL import Image12from transformers import (13    AutoProcessor,14    AutoModelForCausalLM,15    AutoModel,16    AutoTokenizer,17    Qwen2_5_VLForConditionalGeneration,18    TextIteratorStreamer19)20from huggingface_hub import snapshot_download21from qwen_vl_utils import process_vision_info22 23 24 25 26# Suppress the warning about uninitialized weights27warnings.filterwarnings('ignore', message='Some weights.*were not initialized')28 29 30 31 32# Try importing Qwen3VL if available33try:34    from transformers import Qwen3VLForConditionalGeneration35except ImportError:36    Qwen3VLForConditionalGeneration = None37 38 39 40 41MAX_MAX_NEW_TOKENS = 409642DEFAULT_MAX_NEW_TOKENS = 204843MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))44CACHE_DIR = os.getenv("HF_CACHE_DIR", "./models")45 46 47 48 49device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")50 51 52 53 54print(f"Initial Device: {device}")55print(f"CUDA Available: {torch.cuda.is_available()}")56 57 58 59 60# Load Chandra-OCR61try:62    MODEL_ID_V = "datalab-to/chandra"63    processor_v = AutoProcessor.from_pretrained(MODEL_ID_V, trust_remote_code=True)64    if Qwen3VLForConditionalGeneration:65        model_v = Qwen3VLForConditionalGeneration.from_pretrained(66            MODEL_ID_V,67            trust_remote_code=True,68            torch_dtype=torch.float16,69            device_map="auto"70        ).eval()71        print("✓ Chandra-OCR loaded")72    else:73        model_v = None74        print("✗ Chandra-OCR: Qwen3VL not available")75except Exception as e:76    model_v = None77    processor_v = None78    print(f"✗ Chandra-OCR: Failed to load - {str(e)}")79 80 81 82 83# Load Nanonets-OCR2-3B84try:85    MODEL_ID_X = "nanonets/Nanonets-OCR2-3B"86    processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True)87    model_x = Qwen2_5_VLForConditionalGeneration.from_pretrained(88        MODEL_ID_X,89        trust_remote_code=True,90        torch_dtype=torch.float16,91        device_map="auto"92    ).eval()93    print("✓ Nanonets-OCR2-3B loaded")94except Exception as e:95    model_x = None96    processor_x = None97    print(f"✗ Nanonets-OCR2-3B: Failed to load - {str(e)}")98 99# Load olmOCR-2-7B-1025100try:101    MODEL_ID_M = "allenai/olmOCR-2-7B-1025"102    processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)103    model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(104        MODEL_ID_M,105        trust_remote_code=True,106        torch_dtype=torch.float16,107        device_map="auto"108    ).eval()109    print("✓ olmOCR-2-7B-1025 loaded")110except Exception as e:111    model_m = None112    processor_m = None113    print(f"✗ olmOCR-2-7B-1025: Failed to load - {str(e)}")114 115 116 117 118@spaces.GPU119def generate_image(model_name: str, text: str, image: Image.Image,120                   max_new_tokens: int, temperature: float, top_p: float,121                   top_k: int, repetition_penalty: float):122    """123    Generates responses using the selected model for image input.124    Yields raw text and Markdown-formatted text.125    This function is decorated with @spaces.GPU to ensure it runs on GPU126    when available in Hugging Face Spaces.127    Args:128        model_name: Name of the OCR model to use129        text: Prompt text for the model130        image: PIL Image object to process131        max_new_tokens: Maximum number of tokens to generate132        temperature: Sampling temperature133        top_p: Nucleus sampling parameter134        top_k: Top-k sampling parameter135        repetition_penalty: Penalty for repeating tokens136    Yields:137        tuple: (raw_text, markdown_text)138    """139    # Device will be cuda when @spaces.GPU decorator activates140    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")141 142 143    # Select model and processor based on model_name144    if model_name == "olmOCR-2-7B-1025":145        if model_m is None:146            yield "olmOCR-2-7B-1025 is not available.", "olmOCR-2-7B-1025 is not available."147            return148        processor = processor_m149        model = model_m150    elif model_name == "Nanonets-OCR2-3B":151        if model_x is None:152            yield "Nanonets-OCR2-3B is not available.", "Nanonets-OCR2-3B is not available."153            return154        processor = processor_x155        model = model_x156    elif model_name == "Chandra-OCR":157        if model_v is None:158            yield "Chandra-OCR is not available.", "Chandra-OCR is not available."159            return160        processor = processor_v161        model = model_v162    else:163        yield "Invalid model selected.", "Invalid model selected."164        return165 166 167 168 169    if image is None:170        yield "Please upload an image.", "Please upload an image."171        return172 173 174    try:175        # Prepare messages in chat format176        messages = [{177            "role": "user",178            "content": [179                {"type": "image"},180                {"type": "text", "text": text},181            ]182        }]183 184 185        # Apply chat template with fallback186        try:187            prompt_full = processor.apply_chat_template(188                messages, 189                tokenize=False, 190                add_generation_prompt=True191            )192        except Exception as template_error:193            # Fallback: create a simple prompt without chat template194            print(f"Chat template error: {template_error}. Using fallback prompt.")195            prompt_full = f"{text}"196 197 198 199 200        # Process inputs201        inputs = processor(202            text=[prompt_full],203            images=[image],204            return_tensors="pt",205            padding=True206        ).to(device)207 208 209 210 211        # Setup streaming generation212        streamer = TextIteratorStreamer(213            processor.tokenizer if hasattr(processor, 'tokenizer') else processor, 214            skip_prompt=True, 215            skip_special_tokens=True216        )217 218 219        generation_kwargs = {220            **inputs,221            "streamer": streamer,222            "max_new_tokens": max_new_tokens,223            "do_sample": True,224            "temperature": temperature,225            "top_p": top_p,226            "top_k": top_k,227            "repetition_penalty": repetition_penalty,228        }229 230 231        # Start generation in separate thread232        thread = Thread(target=model.generate, kwargs=generation_kwargs)233        thread.start()234 235 236        # Stream the results237        buffer = ""238        for new_text in streamer:239            buffer += new_text240            buffer = buffer.replace("<|im_end|>", "")241            time.sleep(0.01)242            yield buffer, buffer243 244 245        # Ensure thread completes246        thread.join()247 248 249    except Exception as e:250        error_msg = f"Error during generation: {str(e)}"251        print(f"Full error: {e}")252        import traceback253        traceback.print_exc()254        yield error_msg, error_msg255 256 257 258 259# Example usage for Gradio interface260if __name__ == "__main__":261    import gradio as gr262 263 264    # Determine available models265    available_models = []266    if model_m is not None:267        available_models.append("olmOCR-2-7B-1025")268        print("  Added: olmOCR-2-7B-1025")269    if model_x is not None:270        available_models.append("Nanonets-OCR2-3B")271        print("  Added: Nanonets-OCR2-3B")272    if model_v is not None:273        available_models.append("Chandra-OCR")274        print("  Added: Chandra-OCR")275    if not available_models:276        print("ERROR: No models were loaded successfully!")277        exit(1)278 279 280    print(f"\n✓ Available models for dropdown: {', '.join(available_models)}")281 282 283    with gr.Blocks(title="Multi-Model OCR") as demo:284        gr.Markdown("# 🔍 Multi-Model OCR Application")285        gr.Markdown("Upload an image and select a model to extract text. Models run on GPU via Hugging Face Spaces.")286 287 288        with gr.Row():289            with gr.Column():290                model_selector = gr.Dropdown(291                    choices=available_models,292                    value=available_models[0] if available_models else None,293                    label="Select OCR Model"294                )295                image_input = gr.Image(type="pil", label="Upload Image")296                text_input = gr.Textbox(297                    value="Extract all text from this image.",298                    label="Prompt",299                    lines=2300                )301 302 303                with gr.Accordion("Advanced Settings", open=False):304                    max_tokens = gr.Slider(305                        minimum=1,306                        maximum=MAX_MAX_NEW_TOKENS,307                        value=DEFAULT_MAX_NEW_TOKENS,308                        step=1,309                        label="Max New Tokens"310                    )311                    temperature = gr.Slider(312                        minimum=0.1,313                        maximum=2.0,314                        value=0.7,315                        step=0.1,316                        label="Temperature"317                    )318                    top_p = gr.Slider(319                        minimum=0.0,320                        maximum=1.0,321                        value=0.9,322                        step=0.05,323                        label="Top P"324                    )325                    top_k = gr.Slider(326                        minimum=1,327                        maximum=100,328                        value=50,329                        step=1,330                        label="Top K"331                    )332                    repetition_penalty = gr.Slider(333                        minimum=1.0,334                        maximum=2.0,335                        value=1.1,336                        step=0.1,337                        label="Repetition Penalty"338                    )339 340 341                submit_btn = gr.Button("Extract Text", variant="primary")342 343 344            with gr.Column():345                output_text = gr.Textbox(label="Extracted Text", lines=20)346                output_markdown = gr.Markdown(label="Formatted Output")347 348 349        gr.Markdown("""350        ### Available Models:351        - **olmOCR-2-7B-1025**: Allen AI's OCR model352        - **Nanonets-OCR2-3B**: Nanonets OCR model353        - **Chandra-OCR**: Datalab OCR model354        """)355 356 357        submit_btn.click(358            fn=generate_image,359            inputs=[360                model_selector,361                text_input,362                image_input,363                max_tokens,364                temperature,365                top_p,366                top_k,367                repetition_penalty368            ],369            outputs=[output_text, output_markdown]370        )371 372 373    # Launch with share=True for Hugging Face Spaces374    demo.launch(share=True)375