Spaces:
Runtime error
Runtime error
File size: 607 Bytes
980de81 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 |
import gradio as gr
from transformers import AutoProcessor, AutoModelForCausalLM
from PIL import Image
import torch
model = AutoModelForCausalLM.from_pretrained("llava-hf/llava-1.5-7b-hf", torch_dtype=torch.float16).to("cuda")
processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
def chat(image, prompt):
inputs = processor(prompt, images=image, return_tensors="pt").to("cuda")
output = model.generate(**inputs, max_new_tokens=50)
return processor.tokenizer.decode(output[0], skip_special_tokens=True)
gr.Interface(fn=chat, inputs=["image", "text"], outputs="text").launch()
|