Deep Alignment: Merged LoRA weights and synced token configs

app.py

@@ -1,32 +1,57 @@
 import os
+import time
 import gradio as gr
 import torch
 import numpy as np
 import cv2
 from PIL import Image
 from peft import PeftModel
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from transformers import AutoTokenizer, TrOCRProcessor, ViTImageProcessor, VisionEncoderDecoderModel
 from cnn_model import CharacterClassifier
 
-# --- CONFIGURATION ---
+# --- CONSTANTS (From local trocr_engine.py) ---
 BASE_MODEL_ID = "paudelanil/trocr-devanagari-2"
 ADAPTER_ID = "manishw10/devgen-trocr-devanagari-lora"
 CNN_MODEL_PATH = "devanagari-cnn-classifier.pt"
+SPECIAL_TOKEN_NAMES = ("bos_token_id", "cls_token_id", "eos_token_id", "pad_token_id", "sep_token_id")
 
 IS_SPACE = "SPACE_ID" in os.environ
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# --- INITIALIZATION ---
-print(f"System: Restoring Original DevGen Logic...")
-processor = TrOCRProcessor.from_pretrained(BASE_MODEL_ID)
+# --- MODEL LOADING (Mirrored from TrOCREngine.__init__) ---
+print(f"System: Aligning with local TrOCREngine...")
+
+# 1. Load Processor
+try:
+    processor = TrOCRProcessor.from_pretrained(BASE_MODEL_ID)
+except Exception:
+    image_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+    processor = TrOCRProcessor(image_processor=image_processor, tokenizer=tokenizer)
+
+# 2. Load and Config Model
 base_model = VisionEncoderDecoderModel.from_pretrained(BASE_MODEL_ID)
-model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
+base_model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
+base_model.config.pad_token_id = processor.tokenizer.pad_token_id
+base_model.config.eos_token_id = processor.tokenizer.sep_token_id
+base_model.config.vocab_size = base_model.config.decoder.vocab_size
+
+# 3. Apply and MERGE LoRA (Critical for consistency)
+peft_model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
+try:
+    model = peft_model.merge_and_unload()
+    print("System: LoRA weights merged successfully (1:1 with local).")
+except Exception as e:
+    print(f"System: Merge failed, using wrapper: {e}")
+    model = peft_model
+
 model.to(device)
 model.eval()
 
+# 4. Load CNN
 cnn_engine = CharacterClassifier(model_path=CNN_MODEL_PATH, device=device)
 
-# --- ORIGINAL ROUTING LOGIC (1:1 from image_router.py) ---
+# --- ORIGINAL ROUTING (Mirrored from image_router.py) ---
 def _flood_fill(binary, visited, start_y, start_x, h, w):
     stack = [(start_y, start_x)]
     size = 0
@@ -36,10 +61,10 @@ def _flood_fill(binary, visited, start_y, start_x, h, w):
             continue
         visited[y, x] = True
         size += 1
-        stack.extend([(y + 1, x), (y - 1, x), (y, x + 1), (y, x - 1)])
+        stack.extend([(y+1, x), (y-1, x), (y, x+1), (y, x-1)])
     return size
 
-def _count_blobs(binary, min_size=10):
+def count_blobs(binary, min_size=10):
     h, w = binary.shape
     visited = np.zeros_like(binary, dtype=bool)
     count = 0
@@ -51,7 +76,7 @@ def _count_blobs(binary, min_size=10):
                     count += 1
     return count
 
-def original_classify_input(image):
+def classify_input_type(image):
     gray = image.convert("L")
     arr = np.array(gray)
     threshold = min(arr.mean() * 0.75, 200)
@@ -60,15 +85,15 @@ def original_classify_input(image):
     rows = np.any(binary, axis=1)
     cols = np.any(binary, axis=0)
     if not rows.any() or not cols.any():
-        return "character", "no_ink"
+        return "character", 0.5, 0.0, 0
     
     rmin, rmax = np.where(rows)[0][[0, -1]]
     cmin, cmax = np.where(cols)[0][[0, -1]]
     w, h = cmax - cmin + 1, rmax - rmin + 1
     aspect_ratio = w / max(h, 1)
-    blob_count = _count_blobs(binary, min_size=max(binary.size * 0.001, 10))
+    blob_count = count_blobs(binary, min_size=max(binary.size * 0.001, 10))
     
-    # decision logic (exact copy)
+    # DECISION LOGIC (1:1 with local image_router.py)
     is_character = True
     if aspect_ratio > 2.5: is_character = False
     elif aspect_ratio > 1.8 and blob_count >= 3: is_character = False
@@ -77,48 +102,59 @@ def original_classify_input(image):
     elif blob_count == 1 and aspect_ratio < 1.5: is_character = True
     elif aspect_ratio > 1.6: is_character = False
     
-    return "character" if is_character else "word", f"AR: {aspect_ratio:.2f}, Blobs: {blob_count}"
+    return ("character" if is_character else "word"), aspect_ratio, blob_count
 
-# --- MAIN INFERENCE ---
+# --- PREDICT ---
 def predict(image, manual_mode):
-    if image is None: return None, "Upload image.", ""
+    if image is None: return None, "Upload an image.", "", ""
     
+    # 1. Routing
     if manual_mode == "Automatic":
-        mode, reason = original_classify_input(image)
+        mode, ar, bc = classify_input_type(image)
+        routing_status = f"Auto: {mode.upper()} (AR: {ar:.2f}, Blobs: {bc})"
     else:
-        mode, reason = manual_mode.lower(), "Manual"
-    
-    if mode == "character" and cnn_engine.available:
-        # Use exact CNN preprocessing from CharacterClassifier
-        result = cnn_engine.predict(image)
-        return result["text"], f"Original Logic: {mode.upper()} ({reason})", "CNN Classifier"
-    else:
-        # Standard TrOCR inference
-        pixel_values = processor(image.convert("RGB"), return_tensors="pt").pixel_values.to(device)
-        with torch.no_grad():
-            gen_ids = model.base_model.generate(
-                pixel_values=pixel_values,
-                num_beams=4,
-                max_new_tokens=64,
-                decoder_start_token_id=model.config.decoder_start_token_id
-            )
-            text = processor.batch_decode(gen_ids, skip_special_tokens=True)[0]
-        return text, f"Original Logic: {mode.upper()} ({reason})", "TrOCR + LoRA"
+        mode = manual_mode.lower()
+        routing_status = f"Manual: {mode.upper()}"
+        
+    try:
+        if mode == "character" and cnn_engine.available:
+            result = cnn_engine.predict(image)
+            return result["text"], routing_status, "CNN Classifier", ""
+        else:
+            # Word Recognition
+            pixel_values = processor(image.convert("RGB"), return_tensors="pt").pixel_values.to(device)
+            with torch.no_grad():
+                outputs = model.generate(
+                    pixel_values,
+                    num_beams=4,
+                    max_length=128,
+                    early_stopping=True
+                )
+            text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+            return text, routing_status, "TrOCR + LoRA", ""
+    except Exception as e:
+        return f"Error: {str(e)}", "Inference Failed", "None", ""
 
 # --- UI ---
-with gr.Blocks(theme=gr.themes.Default(), css=".gradio-container {background: #0f172a; color: white;}") as demo:
-    gr.Markdown("# 🕉️ DevGen OCR (Original Logic)")
-    with gr.Row():
+CSS = """
+.gradio-container { background: #0f172a; color: white; font-family: 'Inter', sans-serif; }
+.panel { background: rgba(30, 41, 59, 0.8); border-radius: 20px; padding: 20px; border: 1px solid #334155; }
+.result-box { font-size: 2.5rem !important; font-weight: bold; color: #818cf8; text-align: center; }
+"""
+
+with gr.Blocks(css=CSS, theme=gr.themes.Default()) as demo:
+    gr.Markdown("# 🕉️ DevGen OCR — Professional Engine")
+    with gr.Row(elem_classes="panel"):
         with gr.Column():
-            inp = gr.Image(type="pil", label="Input")
-            mode = gr.Radio(["Automatic", "Word", "Character"], value="Automatic", label="Mode")
-            btn = gr.Button("Recognize", variant="primary")
+            input_img = gr.Image(type="pil", label="Input Handwriting")
+            mode_sel = gr.Radio(["Automatic", "Word", "Character"], value="Automatic", label="Mode")
+            run_btn = gr.Button("Recognize", variant="primary")
         with gr.Column():
-            out = gr.Textbox(label="Result")
-            status = gr.Markdown("Ready.")
-            eng = gr.Textbox(label="Engine", interactive=False)
+            output_text = gr.Textbox(label="Recognition Result", elem_classes="result-box")
+            status_lbl = gr.Markdown("Engine ready.")
+            engine_lbl = gr.Textbox(label="Model Used", interactive=False)
 
-    btn.click(predict, [inp, mode], [out, status, eng])
+    run_btn.click(predict, [input_img, mode_sel], [output_text, status_lbl, engine_lbl])
 
 if __name__ == "__main__":
     demo.launch()