app.py

import os
import torch
from transformers import AutoModelForCausalLM, AutoProcessor, TextIteratorStreamer
from PIL import Image
import gradio as gr
from threading import Thread
import logging
import sys

# --- Configuration ---
CONCURRENCY_LIMIT = 1
DEVICE = "cpu"
DTYPE = torch.float32 

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stderr)]
)
logger = logging.getLogger("TachiwinOCR")

# Set CPU threads
torch.set_num_threads(os.cpu_count() or 4)

PROMPTS = {
    "ocr": "OCR:",
    "table": "Table Recognition:",
    "formula": "Formula Recognition:",
    "chart": "Chart Recognition:",
}

# --- Global Model Loading ---
# We load the model globally so it persists across requests. 
# No need for a custom Manager class.
model_path = "tachiwin/PaddleOCR-VL-Tachiwin-BF16"

try:
    logger.info(f"Loading processor from {model_path}...")
    processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
    
    logger.info(f"Loading model from {model_path}...")
    model = AutoModelForCausalLM.from_pretrained(
        model_path, 
        trust_remote_code=True, 
        torch_dtype=DTYPE
    ).to(DEVICE).eval()
    logger.info("Model loaded successfully.")
except Exception as e:
    logger.error(f"Failed to load model: {e}")
    raise e

def inference(img):
    """
    Process image with OCR and Stream the extracted text.
    """
    if img is None:
        yield "Please upload an image."
        return

    # Basic cleanup
    if isinstance(img, str):
        image = Image.open(img).convert("RGB")
    else:
        image = Image.fromarray(img).convert("RGB")

    task = "ocr"
    
    # Prepare inputs
    messages = [
        {"role": "user",
         "content": [
                {"type": "image"},
                {"type": "text", "text": PROMPTS[task]},
            ]
        }
    ]
    
    try:
        text_prompt = processor.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

        inputs = processor(
            image,
            text_prompt,
            add_special_tokens=False,
            return_tensors="pt",
        ).to(DEVICE)
        
        # Initialize Streamer
        streamer = TextIteratorStreamer(
            processor.tokenizer, 
            skip_prompt=True, 
            skip_special_tokens=True
        )

        # Generation Arguments
        generation_kwargs = dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=256,
            min_new_tokens=2,
            do_sample=True,
            use_cache=True,
            temperature=1.5,
            min_p=0.1,
        )

        logger.info("Starting model generation...")
        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()

        generated_text = ""
        status_indicator = " 🦡 ...generating ⏳"
        yield "🦡 The text will start showing shortly... ⏳"
        for new_text in streamer:
            if generated_text == "":
                # Todo: remove the removal of first character when the model is fixed
                generated_text = new_text[1:] 
            else:
                generated_text += new_text
            # Yielding here updates the Gradio textbox in real-time
            logger.info("yielding: " + generated_text)
            yield generated_text + status_indicator
        
        yield generated_text

    except Exception as e:
        import traceback
        error_detail = traceback.format_exc()
        logger.error(error_detail)
        yield f"Error during OCR processing:\n\n```\n{error_detail}\n```"


# --- Interface Setup ---

title = '🌎 Tachiwin OCR for the Indigenous Languages of Mexico 🦡'

description = '''
### PaddleOCR-VL Fine-tuned for the 68 Indigenous Languages of Mexico

This model represents a **world first in tech access and linguistic rights**, specifically trained to recognize 
the diverse character and glyph repertoire of Mexico's 68 indigenous languages.

**How to use:** Simply upload an image containing text in any Mexican indigenous language, and the model will 
detect and recognize the text.

### Warning: as this free demonstrator space uses only CPU, even small image could take up to 5 minutes, so be patient.

🔗 [TachiwinOCR Model](https://huggingface.co/tachiwin/PaddleOCR-VL-Tachiwin-BF16)
🔗 [TachiwinOCR Training Code](https://github.com/ljcamargo/tachiwin_paddleocrvl_finetuning)
'''

examples = [
    ['cco.jpg'],
    ['cnt.jpg'],
    ['maj.jpg'],
    ['mir.jpg'],
    ['ote.jpg'],
    ['otm.jpg'],
    ['lac.jpg'],
]

example_labels = """
### Example Images:
| Image | Language | Text |
|-------|----------|-------------|
| cco.jpg | Comaltepec Chinantec | jo̱ dsʉꞌ i̱ dseaˋ íˋ cajɨ́ɨmˉbre uíiꞌ˜ e dseeˉ e caꞌéerˋ do, co̱ꞌ cajíimˉ búꞌˆ quiáꞌrˉ írˋ, co̱ꞌ i̱ búꞌˆ do caféꞌˋreꞌ laco̱ꞌ féꞌˋ dseabˋ, |
| cnt.jpg | Tepetotutla Chiantec | JMƗG₄ JË₁CA₂TÓ'₂ BÁ₄ LA₂ CHONG₂ JNIOG₄. MA₂NEI'₂ BÁ₄ 'NIA'₂, JÁ₅ JAN₂ I₂'ŊIA₅₄ CRISTO. RË₂NË́₃ NË́₃, JUƗN₅ BÁ₄ I₂MA₂CA₂RË₃JNIÁ₂ I₂'ŊIA₅₄ CRISTO. JAUN₂ BÁ₄ LË₃ NE₄ JNIOG₄ A₂JA₂QUIÁN₃ JMƗG₄ JË₁CA₂TÓ'₂. I₂CA₂'UƗN₂ JË₄ QUIÁN₂ JNIOG₄ BÁ₄ 'ÉI₂. DSÓN'₂ BÁ₄ DSAU₅, |
| maj.jpg | Mazatec, Jalapa de Díaz | Kui xi já maña̱ xi ngakjá ku̱a̱kúya ni xi ts'e̱ Nti̱a̱ná. Kj'a̱í ni xi ku̱a̱kúyanu̱u, kui xi ts'i̱ínkatsúnnu̱u. Najmi ts'i̱ínkie yjoho̱ nga Nda̱ Nti̱a̱ná xi ts'asjejihi̱n. B'a̱ ts'ín ki̱tsa̱ ts'i̱ín nibánehe̱ ra̱ yjoho̱ nga n'e̱kje. Nkjin xi i̱ncha ts'i̱ín ni xi i̱ncha ts'ín jóo̱, ni xi tu̱ subahá maná. |
| mir.jpg | Isthmus Mixe | Cab jaduhṉ yhahixøꞌøy coo jaꞌa naam̱dägøꞌøbä tiúnät wiindsǿṉ maa jaꞌa Diostøjcän, coo jaduhṉ ñäꞌä niguiumayǿøjät. |
| otm.jpg | Eastern Highland Otomi | ma'ueque ma mbʉihʉ. Nɛ gätho gahʉ dyʉ mbäją gahʉ bi 'dac ma ts |
| lac.jpg | Lacandon | wa quin chen u'yicob a t'ʌnex, wa yʌn in wu'yicob a ba' cu ya'aric C'uj? Tin t'ʌn, mʌ' in wu'yicob a t'ʌnex, yʌn in wu'yicob a ba' cu ya'aric C'uj. Yʌn in man in wa'aricob a ba' caj in wirajob yejer a ba' caj in wu'yajob ―baxuc tu ya'araj Pedro ti' u jach ts'urirob. Jeroj tune', chich t'ʌn Pedro yejer Juan ten u jach ts'urirob u winiquirob judío, caj ts'oquij caj cha'b u binob ten u jach ts'urirob. |
"""

css = """
.output_image, .input_image {height: 40rem !important; width: 100% !important;} 
.output_markdown {min-height: 30rem !important;}
.output_markdown p, .output_markdown {
    font-size: 1.3rem !important; 
    line-height: 1.6 !important;
}
"""

demo = gr.Interface(
    fn=inference,
    inputs=gr.Image(type='filepath', label='Input'),
    outputs=gr.Markdown(label='Output', elem_classes="output_markdown"),
    title=title,
    description=description,
    examples=examples,
    cache_examples=False,
    css=css,
    concurrency_limit=CONCURRENCY_LIMIT,
    article=f"""
    {example_labels}
    
    ### About Tachiwin 🦡
    
    **Tachiwin** (from Totonac - "Language") is dedicated to bridging 
    the digital divide for indigenous languages of Mexico through AI technology.
    
    ### Supported Language Families
    
    **Uto-Aztecan:** Náhuatl, Yaqui, Mayo, Huichol, Tepehuán, Tarahumara  
    **Mayan:** Maya, Tzeltal, Tzotzil, Chol, Tojolabal, Q'anjob'al, Mam  
    **Oto-Manguean:** Zapoteco, Mixteco, Otomí, Mazateco, Chinanteco, Triqui  
    **Totonac-Tepehua:** Totonaco, Tepehua  
    **Mixe-Zoque:** Mixe, Zoque, Popoluca  
    **Other:** Purépecha, Huave, Seri, Kickapoo, Kiliwa  
    
    ...covering all 68 officially recognized indigenous languages of Mexico.
    
    ---
    
    Made with ❤️ for linguistic diversity and indigenous rights 🦡

    """
)

if __name__ == "__main__":
    demo.queue().launch(debug=True)