import torch from transformers import AutoModelForCausalLM, AutoProcessor from PIL import Image import requests import gradio as gr import spaces import subprocess subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True) model_id = "yifeihu/TB-OCR-preview-0.1" DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = AutoModelForCausalLM.from_pretrained( model_id, device_map="cuda", trust_remote_code=True, torch_dtype="auto", attn_implementation='flash_attention_2', load_in_4bit=True ) processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True, num_crops=16 ) def phi_ocr(image): question = "Convert the text to markdown format." prompt_message = [{ 'role': 'user', 'content': f'<|image_1|>\n{question}', }] prompt = processor.tokenizer.apply_chat_template(prompt_message, tokenize=False, add_generation_prompt=True) inputs = processor(prompt, [image], return_tensors="pt").to("cuda") generation_args = { "max_new_tokens": 1024, "temperature": 0.1, "do_sample": False } generate_ids = model.generate(**inputs, eos_token_id=processor.tokenizer.eos_token_id, **generation_args) generate_ids = generate_ids[:, inputs['input_ids'].shape[1]:] response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] response = response.split("")[0] return response @spaces.GPU def process_image(input_image): return phi_ocr(input_image) with gr.Blocks() as demo: gr.Markdown("# OCR with TB-OCR-preview-0.1") gr.Markdown("Upload an image to extract and convert text to markdown format.") gr.Markdown("[Check out the model here](https://huggingface.co/yifeihu/TB-OCR-preview-0.1)") with gr.Row(): input_image = gr.Image(type="pil") output_text = gr.Textbox() input_image.change(fn=process_image, inputs=input_image, outputs=output_text) demo.launch()