import gradio as gr
from transformers import pipeline

model_id = "bczhou/tiny-llava-v1-hf"
pipe = pipeline("image-to-text", model=model_id)

def generate_text(prompt, image):
    # Generate text description
    prompt = f"USER: <image>\n{prompt}\nASSISTANT:"
    text_description = pipe(images=image, prompt=prompt,generate_kwargs={"max_new_tokens": 200}) # Batch for efficiency
    return text_description[0]['generated_text'].split("\nASSISTANT:")[-1]

# Create Gradio interface
inputs = [gr.Textbox(label="Input Text"), gr.Image(type="pil")]  # Input for uploading an image
outputs = gr.Textbox(label="Generated Text")  # Output for displaying the text

# Create the interface
interface = gr.Interface(fn=generate_text, inputs=inputs, outputs=outputs, title="LLaVa Image-to-Text")

# Launch the interface
interface.launch()