import gradio as gr
from PIL import Image
from image_captioning import caption_image_unconditional, caption_image_conditional

# Gradio interface for unconditional image captioning
def gradio_unconditional_image_caption(image):
    return caption_image_unconditional(image)

# Gradio interface for conditional image captioning
def gradio_conditional_image_caption(image, text):
    return caption_image_conditional(image, text)

# Define the Gradio interface using the updated syntax
iface_unconditional = gr.Interface(fn=gradio_unconditional_image_caption,
                                   inputs=gr.Image(type="pil"),
                                   outputs="text",
                                   title="Unconditional Image Captioning",
                                   description="Upload an image to get an unconditional caption generated by the BLIP model.")

iface_conditional = gr.Interface(fn=gradio_conditional_image_caption,
                                 inputs=[gr.Image(type="pil"), gr.Textbox(lines=1, placeholder="Enter a prompt")],
                                 outputs="text",
                                 title="Conditional Image Captioning",
                                 description="Upload an image and enter a prompt to get a conditional caption generated by the BLIP model.")

# Combine both interfaces into a single Gradio app
app = gr.TabbedInterface([iface_unconditional, iface_conditional], ["Unconditional", "Conditional"])

# Launch the app
app.launch()