Update app.py

app.py

@@ -1,658 +1,149 @@
+# app.py — HTR Space (Compact Version)
 
-
-
-
-import os
-
-import time
-
+import os, 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 transformers import (
-
-    AutoProcessor,
-
-    AutoModelForImageTextToText,
-
-    Qwen2_5_VLForConditionalGeneration,
-
-)
-
-
-
-# ---------------------------
-
-# Models
-
-# ---------------------------
-
+# ---------------- Models ----------------
 MODEL_PATHS = {
-
-    "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,
-
-    ),
-
+    "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"
 
-
-
-# ---------------------------
-
-# Preload models at startup
-
-# ---------------------------
-
-_loaded_processors = {}
-
-_loaded_models = {}
-
-
-
+_loaded_processors, _loaded_models = {}, {}
 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}")
-
-
-
-# ---------------------------
-
-# Warmup (GPU)
-
-# ---------------------------
-
+    try:
+        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], _loaded_models[name] = processor, model
+        print(f"✅ {name} ready.")
+    except Exception as e:
+        print(f"⚠️ Failed to load {name}: {e}")
+
+# ---------------- GPU Warmup ----------------
 @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}"
-
-
-
-# ---------------------------
-
-# Helpers
-
-# ---------------------------
-
+    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."}]}]
+        chat_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) if tokenizer and hasattr(tokenizer, "apply_chat_template") else "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}"
+
+# ---------------- 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")
-
-
+    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")
+    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)
-
-
+    for obj in [processor, tokenizer, getattr(model, "tokenizer", None)]:
+        try: return obj.batch_decode(output_ids, skip_special_tokens=True)[0]
+        except Exception: pass
+    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
-
-# ---------------------------
-
+    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 ----------------
 @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()
-
-
+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, model, tokenizer = _loaded_processors[model_choice], _loaded_models[model_choice], getattr(_loaded_processors[model_choice], "tokenizer", None)
+    prompt = _default_prompt(query)
+    batch = _build_inputs(processor, tokenizer, image, prompt).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)
+    return _decode_text(model, processor, tokenizer, output_ids).replace("<|im_end|>", "").strip()
+
+# ---------------- Export Helpers ----------------
+def _safe_text(text: str) -> str: 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
+    doc = SimpleDocTemplate("output.pdf")
+    flowables = [Paragraph(t, getSampleStyleSheet()["Normal"]) for t in text.splitlines() if t != ""]
+    if not flowables: flowables = [Paragraph(" ", getSampleStyleSheet()["Normal"])]
+    doc.build(flowables)
+    return "output.pdf"
 
 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
-
-
-
-# gTTS uses Google TTS (requires outbound internet). Wrap in try/except so Space doesn't crash.
+    text = _safe_text(text)
+    if not text: return None
+    doc = Document()
+    for line in text.splitlines(): doc.add_paragraph(line)
+    doc.save("output.docx")
+    return "output.docx"
 
 def save_as_audio(text):
+    text = _safe_text(text)
+    if not text: return None
+    try: from gTTS import gTTS; tts = gTTS(text); tts.save("output.mp3"); return "output.mp3"
+    except Exception as e: print(f"gTTS failed: {e}"); return None
 
-    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 Interface
-
-# ---------------------------
-
+# ---------------- 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",
-
-    )
-
-
-
-    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],
-
-        )
-
-
+    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")
+        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")
+        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)
+        pdf_btn = gr.Button("⬇️ Download as PDF")
+        word_btn = gr.Button("⬇️ Download as Word")
+        audio_btn = gr.Button("🔊 Download as Audio")
+        pdf_file, word_file, audio_file = gr.File(label="PDF File"), gr.File(label="Word 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__":
-
-    # Keep queue for GPU tasks; limit concurrency for stability.
-
-    demo.queue(max_size=50).launch(show_error=True)
-
+    demo.queue(max_size=50).launch(show_error=True)