Spaces:
Build error
Build error
| import gradio as gr | |
| import torch | |
| import transformers | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from PIL import Image | |
| import warnings | |
| # disable some warnings | |
| transformers.logging.set_verbosity_error() | |
| transformers.logging.disable_progress_bar() | |
| warnings.filterwarnings('ignore') | |
| # set device | |
| torch.set_default_device('cuda') # or 'cpu' | |
| model_name = 'cognitivecomputations/dolphin-vision-7b' | |
| # create model | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| torch_dtype=torch.float16, | |
| device_map='auto', | |
| trust_remote_code=True) | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| model_name, | |
| trust_remote_code=True) | |
| def inference(prompt, image): | |
| messages = [ | |
| {"role": "user", "content": f'<image>\n{prompt}'} | |
| ] | |
| text = tokenizer.apply_chat_template( | |
| messages, | |
| tokenize=False, | |
| add_generation_prompt=True | |
| ) | |
| text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')] | |
| input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0) | |
| image_tensor = model.process_images([image], model.config).to(dtype=model.dtype) | |
| # generate | |
| output_ids = model.generate( | |
| input_ids, | |
| images=image_tensor, | |
| max_new_tokens=2048, | |
| use_cache=True)[0] | |
| return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip() | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt_input = gr.Textbox(label="Prompt", placeholder="Describe this image in detail") | |
| image_input = gr.Image(label="Image", type="pil") | |
| submit_button = gr.Button("Submit") | |
| with gr.Column(): | |
| output_text = gr.Textbox(label="Output") | |
| submit_button.click(fn=inference, inputs=[prompt_input, image_input], outputs=output_text) | |
| demo.launch() |