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()