import os import time import torch import spaces import warnings import tempfile import sys from io import StringIO from contextlib import contextmanager from threading import Thread from PIL import Image from transformers import ( AutoProcessor, AutoModelForCausalLM, AutoModel, AutoTokenizer, Qwen2_5_VLForConditionalGeneration, TextIteratorStreamer ) from huggingface_hub import snapshot_download from qwen_vl_utils import process_vision_info # Suppress the warning about uninitialized weights warnings.filterwarnings('ignore', message='Some weights.*were not initialized') # Try importing Qwen3VL if available try: from transformers import Qwen3VLForConditionalGeneration except ImportError: Qwen3VLForConditionalGeneration = None MAX_MAX_NEW_TOKENS = 4096 DEFAULT_MAX_NEW_TOKENS = 2048 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096")) CACHE_DIR = os.getenv("HF_CACHE_DIR", "./models") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") print(f"Initial Device: {device}") print(f"CUDA Available: {torch.cuda.is_available()}") # Load Chandra-OCR try: MODEL_ID_V = "datalab-to/chandra" processor_v = AutoProcessor.from_pretrained(MODEL_ID_V, trust_remote_code=True) if Qwen3VLForConditionalGeneration: model_v = Qwen3VLForConditionalGeneration.from_pretrained( MODEL_ID_V, trust_remote_code=True, torch_dtype=torch.float16, device_map="auto" ).eval() print("✓ Chandra-OCR loaded") else: model_v = None print("✗ Chandra-OCR: Qwen3VL not available") except Exception as e: model_v = None processor_v = None print(f"✗ Chandra-OCR: Failed to load - {str(e)}") # Load Nanonets-OCR2-3B try: MODEL_ID_X = "nanonets/Nanonets-OCR2-3B" processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True) model_x = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_X, trust_remote_code=True, torch_dtype=torch.float16, device_map="auto" ).eval() print("✓ Nanonets-OCR2-3B loaded") except Exception as e: model_x = None processor_x = None print(f"✗ Nanonets-OCR2-3B: Failed to load - {str(e)}") # Load olmOCR-2-7B-1025 try: MODEL_ID_M = "allenai/olmOCR-2-7B-1025" processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True) model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained( MODEL_ID_M, trust_remote_code=True, torch_dtype=torch.float16, device_map="auto" ).eval() print("✓ olmOCR-2-7B-1025 loaded") except Exception as e: model_m = None processor_m = None print(f"✗ olmOCR-2-7B-1025: Failed to load - {str(e)}") @spaces.GPU def generate_image(model_name: str, text: str, image: Image.Image, max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float): """ Generates responses using the selected model for image input. Yields raw text and Markdown-formatted text. This function is decorated with @spaces.GPU to ensure it runs on GPU when available in Hugging Face Spaces. Args: model_name: Name of the OCR model to use text: Prompt text for the model image: PIL Image object to process max_new_tokens: Maximum number of tokens to generate temperature: Sampling temperature top_p: Nucleus sampling parameter top_k: Top-k sampling parameter repetition_penalty: Penalty for repeating tokens Yields: tuple: (raw_text, markdown_text) """ # Device will be cuda when @spaces.GPU decorator activates device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Select model and processor based on model_name if model_name == "olmOCR-2-7B-1025": if model_m is None: yield "olmOCR-2-7B-1025 is not available.", "olmOCR-2-7B-1025 is not available." return processor = processor_m model = model_m elif model_name == "Nanonets-OCR2-3B": if model_x is None: yield "Nanonets-OCR2-3B is not available.", "Nanonets-OCR2-3B is not available." return processor = processor_x model = model_x elif model_name == "Chandra-OCR": if model_v is None: yield "Chandra-OCR is not available.", "Chandra-OCR is not available." return processor = processor_v model = model_v else: yield "Invalid model selected.", "Invalid model selected." return if image is None: yield "Please upload an image.", "Please upload an image." return try: # Prepare messages in chat format messages = [{ "role": "user", "content": [ {"type": "image"}, {"type": "text", "text": text}, ] }] # Apply chat template with fallback try: prompt_full = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) except Exception as template_error: # Fallback: create a simple prompt without chat template print(f"Chat template error: {template_error}. Using fallback prompt.") prompt_full = f"{text}" # Process inputs inputs = processor( text=[prompt_full], images=[image], return_tensors="pt", padding=True ).to(device) # Setup streaming generation streamer = TextIteratorStreamer( processor.tokenizer if hasattr(processor, 'tokenizer') else processor, skip_prompt=True, skip_special_tokens=True ) generation_kwargs = { **inputs, "streamer": streamer, "max_new_tokens": max_new_tokens, "do_sample": True, "temperature": temperature, "top_p": top_p, "top_k": top_k, "repetition_penalty": repetition_penalty, } # Start generation in separate thread thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() # Stream the results buffer = "" for new_text in streamer: buffer += new_text buffer = buffer.replace("<|im_end|>", "") time.sleep(0.01) yield buffer, buffer # Ensure thread completes thread.join() except Exception as e: error_msg = f"Error during generation: {str(e)}" print(f"Full error: {e}") import traceback traceback.print_exc() yield error_msg, error_msg # Example usage for Gradio interface if __name__ == "__main__": import gradio as gr # Determine available models available_models = [] if model_m is not None: available_models.append("olmOCR-2-7B-1025") print(" Added: olmOCR-2-7B-1025") if model_x is not None: available_models.append("Nanonets-OCR2-3B") print(" Added: Nanonets-OCR2-3B") if model_v is not None: available_models.append("Chandra-OCR") print(" Added: Chandra-OCR") if not available_models: print("ERROR: No models were loaded successfully!") exit(1) print(f"\n✓ Available models for dropdown: {', '.join(available_models)}") with gr.Blocks(title="Multi-Model OCR") as demo: gr.Markdown("# 🔍 Multi-Model OCR Application") gr.Markdown("Upload an image and select a model to extract text. Models run on GPU via Hugging Face Spaces.") with gr.Row(): with gr.Column(): model_selector = gr.Dropdown( choices=available_models, value=available_models[0] if available_models else None, label="Select OCR Model" ) image_input = gr.Image(type="pil", label="Upload Image") text_input = gr.Textbox( value="Extract all text from this image.", label="Prompt", lines=2 ) with gr.Accordion("Advanced Settings", open=False): max_tokens = gr.Slider( minimum=1, maximum=MAX_MAX_NEW_TOKENS, value=DEFAULT_MAX_NEW_TOKENS, step=1, label="Max New Tokens" ) temperature = gr.Slider( minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature" ) top_p = gr.Slider( minimum=0.0, maximum=1.0, value=0.9, step=0.05, label="Top P" ) top_k = gr.Slider( minimum=1, maximum=100, value=50, step=1, label="Top K" ) repetition_penalty = gr.Slider( minimum=1.0, maximum=2.0, value=1.1, step=0.1, label="Repetition Penalty" ) submit_btn = gr.Button("Extract Text", variant="primary") with gr.Column(): output_text = gr.Textbox(label="Extracted Text", lines=20) output_markdown = gr.Markdown(label="Formatted Output") gr.Markdown(""" ### Available Models: - **olmOCR-2-7B-1025**: Allen AI's OCR model - **Nanonets-OCR2-3B**: Nanonets OCR model - **Chandra-OCR**: Datalab OCR model """) submit_btn.click( fn=generate_image, inputs=[ model_selector, text_input, image_input, max_tokens, temperature, top_p, top_k, repetition_penalty ], outputs=[output_text, output_markdown] ) # Launch with share=True for Hugging Face Spaces demo.launch(share=True)