Update app.py

app.py

@@ -1,197 +1,658 @@
+
+
+
+
 import os
+
+import time
+
+from threading import Thread
+
+
+
 import gradio as gr
+
+import spaces
+
 from PIL import Image
+
 import torch
-from transformers import AutoProcessor, AutoModelForImageTextToText, Qwen2_5_VLForConditionalGeneration
-from reportlab.platypus import SimpleDocTemplate, Paragraph
-from reportlab.lib.styles import getSampleStyleSheet
-from docx import Document
-from gtts import gTTS
 
-# ---------------------------
-# Device & constants
-# ---------------------------
-DEVICE = "cuda"  # Force GPU usage
-MAX_NEW_TOKENS_DEFAULT = 512
+from transformers import (
+
+    AutoProcessor,
+
+    AutoModelForImageTextToText,
+
+    Qwen2_5_VLForConditionalGeneration,
+
+)
+
+
 
 # ---------------------------
-# Models config
+
+# Models
+
 # ---------------------------
+
 MODEL_PATHS = {
-    "Complex handwritings": ("prithivMLmods/Qwen2.5-VL-7B-Abliterated-Caption-it", Qwen2_5_VLForConditionalGeneration),
-    "Simple/scanned handwriting": ("nanonets/Nanonets-OCR-s", Qwen2_5_VLForConditionalGeneration),
-    "Structured handwriting": ("Emeritus-21/Finetuned-full-HTR-model", AutoModelForImageTextToText),
+
+    "Model 1 (Complex handwrittings )": (
+
+        "prithivMLmods/Qwen2.5-VL-7B-Abliterated-Caption-it",
+
+        Qwen2_5_VLForConditionalGeneration,
+
+    ),
+
+    "Model 2 (simple and scanned handwritting )": (
+
+        "nanonets/Nanonets-OCR-s",
+
+        Qwen2_5_VLForConditionalGeneration,
+
+    ),
+
+    "Model 3 (structured handwritting)": (
+
+        "Emeritus-21/Finetuned-full-HTR-model",
+
+        AutoModelForImageTextToText,
+
+    ),
+
 }
 
+
+
+MAX_NEW_TOKENS_DEFAULT = 512
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+
 # ---------------------------
-# Lazy load models
+
+# Preload models at startup
+
 # ---------------------------
+
 _loaded_processors = {}
+
 _loaded_models = {}
 
-def load_model(name):
-    if name in _loaded_models:
-        return _loaded_processors[name], _loaded_models[name]
-    repo_id, cls = MODEL_PATHS[name]
-    processor = AutoProcessor.from_pretrained(repo_id, trust_remote_code=True)
-    model = cls.from_pretrained(
-        repo_id,
-        trust_remote_code=True,
-        torch_dtype=torch.float16,
-        low_cpu_mem_usage=True
-    ).to(DEVICE).eval()
-    _loaded_processors[name] = processor
-    _loaded_models[name] = model
-    return processor, model
+
+
+print("🚀 Preloading models into GPU/CPU memory...")
+
+
+
+for name, (repo_id, cls) in MODEL_PATHS.items():
+
+    try:
+
+        print(f"Loading {name} ...")
+
+        processor = AutoProcessor.from_pretrained(repo_id, trust_remote_code=True)
+
+        model = cls.from_pretrained(
+
+            repo_id,
+
+            trust_remote_code=True,
+
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+
+            low_cpu_mem_usage=True,
+
+        ).to(device).eval()
+
+        _loaded_processors[name] = processor
+
+        _loaded_models[name] = model
+
+        print(f"✅ {name} ready.")
+
+    except Exception as e:
+
+        print(f"⚠️ Failed to load {name}: {e}")
+
+
 
 # ---------------------------
-# OCR function (GPU ready)
+
+# Warmup (GPU)
+
 # ---------------------------
-@gr.utils.space_decorator  # Spaces decorator to detect GPU
-def ocr_image_gpu(image: Image.Image, model_choice: str, query: str = None,
-                  max_new_tokens: int = MAX_NEW_TOKENS_DEFAULT, temperature: float = 0.1,
-                  top_p: float = 1.0, top_k: int = 0, repetition_penalty: float = 1.0):
-    if image is None:
-        return "Please upload an image."
-    if model_choice not in MODEL_PATHS:
-        return f"Invalid model: {model_choice}"
-
-    processor, model = load_model(model_choice)
-
-    prompt = query.strip() if query and query.strip() else (
-        "You are a professional Handwritten OCR system.\n"
-        "TASK: Read the handwritten image and transcribe the text EXACTLY as written.\n"
-        "- Preserve original structure and line breaks.\n"
-        "- Keep spacing, bullet points, numbering, and indentation.\n"
-        "- Render tables as Markdown tables if present.\n"
-        "- Do NOT autocorrect spelling or grammar.\n"
-        "- Do NOT merge lines.\n"
-        "Return RAW transcription only."
-    )
-
-    batch = processor(text=[prompt], images=[image], return_tensors="pt").to(DEVICE)
-
-    with torch.inference_mode():
-        output_ids = model.generate(
-            **batch,
-            max_new_tokens=max_new_tokens,
-            do_sample=False,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            repetition_penalty=repetition_penalty,
-        )
-
-    # decode safely
-    text = ""
-    if hasattr(processor, "batch_decode"):
-        text = processor.batch_decode(output_ids, skip_special_tokens=True)[0]
-    elif hasattr(model, "tokenizer") and model.tokenizer is not None:
-        text = model.tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
-    text = text.replace("<|im_end|>", "").strip()
-    return text
+
+@spaces.GPU
+
+def warmup(progress=gr.Progress(track_tqdm=True)):
+
+    try:
+
+        default_model_choice = next(iter(MODEL_PATHS.keys()))
+
+        processor = _loaded_processors[default_model_choice]
+
+        model = _loaded_models[default_model_choice]
+
+        tokenizer = getattr(processor, "tokenizer", None)
+
+
+
+        messages = [{"role": "user", "content": [{"type": "text", "text": "Warmup."}]}]
+
+        if tokenizer and hasattr(tokenizer, "apply_chat_template"):
+
+            chat_prompt = tokenizer.apply_chat_template(
+
+                messages, tokenize=False, add_generation_prompt=True
+
+            )
+
+        else:
+
+            chat_prompt = "Warmup."
+
+
+
+        inputs = processor(
+
+            text=[chat_prompt],
+
+            images=None,
+
+            return_tensors="pt"
+
+        ).to(device)
+
+
+
+        with torch.inference_mode():
+
+            _ = model.generate(**inputs, max_new_tokens=1)
+
+
+
+        return f"GPU warm and {default_model_choice} ready."
+
+    except Exception as e:
+
+        return f"Warmup skipped: {e}"
+
+
 
 # ---------------------------
-# Export helpers
+
+# Helpers
+
 # ---------------------------
+
+def _build_inputs(processor, tokenizer, image: Image.Image, prompt: str):
+
+    """Build processor inputs for text+image with/without chat template."""
+
+    messages = [
+
+        {
+
+            "role": "user",
+
+            "content": [
+
+                {"type": "image", "image": image},
+
+                {"type": "text", "text": prompt},
+
+            ],
+
+        }
+
+    ]
+
+    if tokenizer and hasattr(tokenizer, "apply_chat_template"):
+
+        chat_prompt = tokenizer.apply_chat_template(
+
+            messages, tokenize=False, add_generation_prompt=True
+
+        )
+
+        return processor(text=[chat_prompt], images=[image], return_tensors="pt")
+
+    # Fallback: plain prompt + image
+
+    return processor(text=[prompt], images=[image], return_tensors="pt")
+
+
+
+def _decode_text(model, processor, tokenizer, output_ids):
+
+    """Robust decode for different processor/tokenizer setups."""
+
+    text = ""
+
+    try:
+
+        if hasattr(processor, "batch_decode"):
+
+            text = processor.batch_decode(output_ids, skip_special_tokens=True)[0]
+
+            return text
+
+    except Exception:
+
+        pass
+
+    try:
+
+        if tokenizer is not None:
+
+            text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
+
+            return text
+
+    except Exception:
+
+        pass
+
+    try:
+
+        model_tok = getattr(model, "tokenizer", None)
+
+        if model_tok is not None:
+
+            text = model_tok.batch_decode(output_ids, skip_special_tokens=True)[0]
+
+            return text
+
+    except Exception:
+
+        pass
+
+    # Last-resort string
+
+    return str(output_ids)
+
+
+
+def _default_prompt(query: str | None) -> str:
+
+    if query and query.strip():
+
+        return query.strip()
+
+    return (
+
+        "You are a professional Handwritten OCR system.\n"
+
+        "TASK: Read the handwritten image and transcribe the text EXACTLY as written.\n"
+
+        "- Preserve original structure and line breaks.\n"
+
+        "- Keep spacing, bullet points, numbering, and indentation.\n"
+
+        "- Render tables as Markdown tables if present.\n"
+
+        "- Do NOT autocorrect spelling or grammar.\n"
+
+        "- Do NOT merge lines.\n"
+
+        "Return RAW transcription only."
+
+    )
+
+
+
+# ---------------------------
+
+# OCR Function (NO STREAMING / NO yield)  ✅ FIX
+
+# ---------------------------
+
+@spaces.GPU
+
+def ocr_image(
+
+    image: Image.Image,
+
+    model_choice: str,
+
+    query: str = None,
+
+    max_new_tokens: int = MAX_NEW_TOKENS_DEFAULT,
+
+    temperature: float = 0.1,
+
+    top_p: float = 1.0,
+
+    top_k: int = 0,
+
+    repetition_penalty: float = 1.0,
+
+    progress=gr.Progress(track_tqdm=True),
+
+):
+
+    if image is None:
+
+        return "Please upload or capture an image."
+
+
+
+    if model_choice not in _loaded_models:
+
+        return f"Invalid model: {model_choice}"
+
+
+
+    processor = _loaded_processors[model_choice]
+
+    model = _loaded_models[model_choice]
+
+    tokenizer = getattr(processor, "tokenizer", None)
+
+
+
+    prompt = _default_prompt(query)
+
+
+
+    # Build inputs
+
+    batch = _build_inputs(processor, tokenizer, image, prompt).to(device)
+
+
+
+    # Generate (no streaming)
+
+    with torch.inference_mode():
+
+        output_ids = model.generate(
+
+            **batch,
+
+            max_new_tokens=max_new_tokens,
+
+            do_sample=False,
+
+            temperature=temperature,
+
+            top_p=top_p,
+
+            top_k=top_k,
+
+            repetition_penalty=repetition_penalty,
+
+        )
+
+
+
+    # Decode
+
+    decoded = _decode_text(model, processor, tokenizer, output_ids)
+
+    cleaned = decoded.replace("<|im_end|>", "").strip()
+
+    return cleaned
+
+
+
+# ---------------------------
+
+# Export Helpers
+
+# ---------------------------
+
+from reportlab.platypus import SimpleDocTemplate, Paragraph
+
+from reportlab.lib.styles import getSampleStyleSheet
+
+from docx import Document
+
+
+
 def _safe_text(text: str) -> str:
-    return (text or "").strip()
+
+    return (text or "").strip()
+
+
 
 def save_as_pdf(text):
-    text = _safe_text(text)
-    if not text:
-        return None
-    filepath = "output.pdf"
-    doc = SimpleDocTemplate(filepath)
-    styles = getSampleStyleSheet()
-    flowables = [Paragraph(t, styles["Normal"]) for t in text.splitlines() if t != ""]
-    if not flowables:
-        flowables = [Paragraph(" ", styles["Normal"])]
-    doc.build(flowables)
-    return filepath
+
+    text = _safe_text(text)
+
+    if not text:
+
+        return None
+
+    filepath = "output.pdf"
+
+    doc = SimpleDocTemplate(filepath)
+
+    styles = getSampleStyleSheet()
+
+    flowables = [Paragraph(t, styles["Normal"]) for t in text.splitlines() if t != ""]
+
+    if not flowables:
+
+        flowables = [Paragraph(" ", styles["Normal"])]
+
+    doc.build(flowables)
+
+    return filepath
+
+
 
 def save_as_word(text):
-    text = _safe_text(text)
-    if not text:
-        return None
-    filepath = "output.docx"
-    doc = Document()
-    for line in text.splitlines():
-        doc.add_paragraph(line)
-    doc.save(filepath)
-    return filepath
+
+    text = _safe_text(text)
+
+    if not text:
+
+        return None
+
+    filepath = "output.docx"
+
+    doc = Document()
+
+    for line in text.splitlines():
+
+        doc.add_paragraph(line)
+
+    doc.save(filepath)
+
+    return filepath
+
+
+
+# gTTS uses Google TTS (requires outbound internet). Wrap in try/except so Space doesn't crash.
 
 def save_as_audio(text):
-    text = _safe_text(text)
-    if not text:
-        return None
-    filepath = "output.mp3"
-    tts = gTTS(text)
-    tts.save(filepath)
-    return filepath
+
+    text = _safe_text(text)
+
+    if not text:
+
+        return None
+
+    try:
+
+        from gTTS import gTTS
+
+        filepath = "output.mp3"
+
+        tts = gTTS(text)
+
+        tts.save(filepath)
+
+        return filepath
+
+    except Exception as e:
+
+        print(f"gTTS failed: {e}")
+
+        return None
+
+
 
 # ---------------------------
-# Gradio UI
+
+# Gradio Interface
+
 # ---------------------------
+
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## ✍🏾 Wilson Handwritten OCR")
-
-    model_choice = gr.Radio(
-        choices=list(MODEL_PATHS.keys()),
-        value=list(MODEL_PATHS.keys())[0],
-        label="Select OCR Model",
-    )
-
-    query_input = gr.Textbox(
-        label="Custom Prompt (optional)",
-        placeholder="Leave empty for RAW structured output",
-    )
-
-    image_input = gr.Image(type="pil", label="Upload Image (desktop/mobile)")
-
-    # Buttons first
-    extract_btn = gr.Button("📤 Extract RAW Text", variant="primary")
-    raw_output = gr.Textbox(
-        label="📜 RAW Structured Output (exact as written)",
-        lines=18,
-        show_copy_button=True,
-    )
-    pdf_file = gr.File(label="PDF File")
-    word_file = gr.File(label="Word File")
-    audio_file = gr.File(label="Audio File")
-
-    with gr.Accordion("⚙️ Advanced Options", open=False):
-        max_new_tokens = gr.Slider(1, 2048, value=MAX_NEW_TOKENS_DEFAULT, step=1, label="Max new tokens")
-        temperature = gr.Slider(0.1, 2.0, value=0.1, step=0.05, label="Temperature")
-        top_p = gr.Slider(0.05, 1.0, value=1.0, step=0.05, label="Top-p (nucleus)")
-        top_k = gr.Slider(0, 1000, value=0, step=1, label="Top-k")
-        repetition_penalty = gr.Slider(0.8, 2.0, value=1.0, step=0.05, label="Repetition penalty")
-
-    # Extract text
-    extract_btn.click(
-        fn=ocr_image_gpu,
-        inputs=[image_input, model_choice, query_input,
-                max_new_tokens, temperature, top_p, top_k, repetition_penalty],
-        outputs=[raw_output]
-    )
-
-    # Export buttons
-    pdf_btn = gr.Button("⬇️ Download as PDF")
-    pdf_btn.click(fn=save_as_pdf, inputs=[raw_output], outputs=[pdf_file])
-
-    word_btn = gr.Button("⬇️ Download as Word")
-    word_btn.click(fn=save_as_word, inputs=[raw_output], outputs=[word_file])
-
-    audio_btn = gr.Button("🔊 Download as Audio")
-    audio_btn.click(fn=save_as_audio, inputs=[raw_output], outputs=[audio_file])
-
-    # Clear button
-    clear_btn = gr.Button("🧹 Clear")
-    clear_btn.click(
-        fn=lambda: ("", None, "", MAX_NEW_TOKENS_DEFAULT, 0.1, 1.0, 0, 1.0),
-        outputs=[raw_output, image_input, query_input,
-                 max_new_tokens, temperature, top_p, top_k, repetition_penalty],
-    )
+
+    gr.Markdown("## ✍🏾 wilson Handwritten OCR ")
+
+
+
+    model_choice = gr.Radio(
+
+        choices=list(MODEL_PATHS.keys()),
+
+        value=list(MODEL_PATHS.keys())[0],
+
+        label="Select OCR Model",
+
+    )
+
+
+
+    with gr.Tab("🖼 Image Inference"):
+
+        query_input = gr.Textbox(
+
+            label="Custom Prompt (optional)",
+
+            placeholder="Leave empty for RAW structured output",
+
+        )
+
+
+
+        # Upload + Webcam (Gradio 4.x uses `sources`)
+
+        image_input = gr.Image(
+
+            type="pil",
+
+            label="Upload / Capture Handwritten Image",
+
+            sources=["upload", "webcam"],
+
+        )
+
+
+
+        with gr.Accordion("⚙️ Advanced Options", open=False):
+
+            max_new_tokens = gr.Slider(1, 2048, value=MAX_NEW_TOKENS_DEFAULT, step=1, label="Max new tokens")
+
+            temperature = gr.Slider(0.1, 2.0, value=0.1, step=0.05, label="Temperature")
+
+            top_p = gr.Slider(0.05, 1.0, value=1.0, step=0.05, label="Top-p (nucleus)")
+
+            top_k = gr.Slider(0, 1000, value=0, step=1, label="Top-k")
+
+            repetition_penalty = gr.Slider(0.8, 2.0, value=1.0, step=0.05, label="Repetition penalty")
+
+
+
+        with gr.Row():
+
+            extract_btn = gr.Button("📤 Extract RAW Text", variant="primary")
+
+            clear_btn = gr.Button("🧹 Clear")
+
+
+
+        raw_output = gr.Textbox(
+
+            label="📜 RAW Structured Output (exact as written)",
+
+            lines=18,
+
+            show_copy_button=True,
+
+        )
+
+
+
+        with gr.Row():
+
+            pdf_btn = gr.Button("⬇️ Download as PDF")
+
+            word_btn = gr.Button("⬇️ Download as Word")
+
+            audio_btn = gr.Button("🔊 Download as Audio")
+
+
+
+        pdf_file = gr.File(label="PDF File")
+
+        word_file = gr.File(label="Word File")
+
+        audio_file = gr.File(label="Audio File")
+
+
+
+        extract_btn.click(
+
+            fn=ocr_image,
+
+            inputs=[
+
+                image_input,
+
+                model_choice,
+
+                query_input,
+
+                max_new_tokens,
+
+                temperature,
+
+                top_p,
+
+                top_k,
+
+                repetition_penalty,
+
+            ],
+
+            outputs=[raw_output],
+
+            api_name="ocr_image",
+
+        )
+
+
+
+        pdf_btn.click(fn=save_as_pdf, inputs=[raw_output], outputs=[pdf_file])
+
+        word_btn.click(fn=save_as_word, inputs=[raw_output], outputs=[word_file])
+
+        audio_btn.click(fn=save_as_audio, inputs=[raw_output], outputs=[audio_file])
+
+
+
+        clear_btn.click(
+
+            fn=lambda: ("", None, "", MAX_NEW_TOKENS_DEFAULT, 0.1, 1.0, 0, 1.0),
+
+            outputs=[raw_output, image_input, query_input, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+
+        )
+
+
 
 if __name__ == "__main__":
-    demo.queue(max_size=50).launch(show_error=True)
+
+    # Keep queue for GPU tasks; limit concurrency for stability.
+
+    demo.queue(max_size=50).launch(show_error=True)
+