Spaces:
Runtime error
Runtime error
File size: 732 Bytes
f27b541 7f99746 f27b541 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
import gradio as gr
from transformers import pipeline, BitsAndBytesConfig
import torch
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16
)
model_id = "llava-hf/llava-1.5-7b-hf"
pipe = pipeline("image-to-text",
model=model_id,
model_kwargs={"quantization_config": quantization_config}
)
def launch(image, prompt):
prompt = f"USER: <image>\n{prompt}\nASSISTANT:"
outputs = pipe(image, prompt=prompt, generate_kwargs={"max_new_tokens": 200})
return out[0]['generated_text']
iface = gr.Interface(launch,
inputs=[gr.Image(type='pil'), 'text'],
outputs="text")
iface.launch() |