import gradio as gr
from transformers import pipeline

try:
    captioning_pipeline = pipeline(
        "image-to-text", 
        model="nlpconnect/vit-gpt2-image-captioning", 
        device=-1  # -1 CPU, 0 GPU
    )
except Exception as e:
    captioning_pipeline = None
    print(f"Error loading model: {e}")

# Inference Function
def generate_caption(image):
    if image is None:
        return "Please upload an image to generate a caption."
    
    if captioning_pipeline is None:
        return "Model failed to load. Check logs for details."

    try:
        results = captioning_pipeline(image, max_new_tokens=50)  # Set a limit on caption length
        caption = results[0]['generated_text']
        return caption
    except Exception as e:
        return f"An error occurred during generation: {e}"

# Gradio Interface Definition
example_paths = [
    ["examples/dog_park.jpg"],
    ["examples/city_scene.png"],
    ["examples/beach_sunset.jpg"]
]

iface = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Textbox(label="Generated Caption"),
    title="Custom Image Caption Generator (Hugging Face Space)",
    description="Upload an image and have a generative AI model describe what it sees.",
    examples=example_paths,
    allow_flagging="auto"
)

if __name__ == "__main__":
    iface.launch()