artyomxyz/pix2struct-docmatix

1import subprocess2subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)3 4import os5 6import gradio as gr7from huggingface_hub import snapshot_download8import spaces9import torch10from transformers import T5TokenizerFast11 12from pix2struct.modeling import Pix2StructModel13from pix2struct.processing import extract_patches14from pix2struct.inference import ask_generator, generate, DocumentQueries, DocumentQuery15 16 17hub_token = os.environ.get('HUB_TOKEN')18model_path = snapshot_download('artyomxyz/pix2struct-docmatix', use_auth_token=hub_token)19 20model = Pix2StructModel.load(model_path)21model.eval()22model = model.to('cuda')23tokenizer = T5TokenizerFast.from_pretrained('google/pix2struct-base')24 25@spaces.GPU26def ask(image, questions):27    questions = questions.split('\n')28    documents = [29        DocumentQueries(30            meta=None,31            patches=extract_patches([image]),32            queries=[33                DocumentQuery(34                    meta=None,35                    generator=ask_generator(tokenizer, question)36                )37                for question in questions38            ]39        )40    ]41    with torch.inference_mode():42        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):43            result = generate(model, documents, device='cuda')44    return '\n'.join([q.output for q in result[0].queries])45 46demo = gr.Interface(47    fn=ask,48    inputs=[49        gr.Image(type='numpy'),50        gr.Textbox(label="Questions (one question per line)"),51    ],52    outputs='text'53)54demo.launch()55