import gradio as gr
from transformers import pipeline
import sys # Import sys for stderr printing

# Load the image-to-text pipeline
# We'll use a common pre-trained model for image captioning: "nlpconnect/vit-gpt2-image-captioning"
# This model requires the 'transformers' library and its dependencies (like torch)
# It also requires the Pillow library for image handling.
try:
    print("Attempting to load image-to-text model...")
    # Specify the task and model name
    image_to_text_pipeline = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
    print("Image-to-text model loaded successfully.")
    model_loaded = True
except Exception as e:
    print(f"Error loading model: {e}", file=sys.stderr)
    print("Please ensure you have 'transformers', 'torch' (or tensorflow), and 'Pillow' installed.", file=sys.stderr)
    image_to_text_pipeline = None # Set to None if loading fails
    model_loaded = False


def generate_text_from_image(image):
    """
    Generates text (caption) from an input image using the loaded model.
    """
    if not model_loaded or not image_to_text_pipeline:
        return "Model not loaded. Please check the application logs for errors."

    if image is None:
        return "Please upload an image."

    print("Received image for processing.")

    try:
        # The pipeline expects a PIL Image object or a file path
        # Gradio's Image input component returns a PIL Image by default
        print("Calling image-to-text pipeline...")
        # The output is typically a list of dictionaries, e.g., [{'generated_text': 'a cat sitting on a couch'}]
        output = image_to_text_pipeline(image)
        print(f"Pipeline raw output: {output}")

        # Extract the generated text
        if output and isinstance(output, list) and len(output) > 0 and 'generated_text' in output[0]:
            generated_text = output[0]['generated_text']
            print(f"Generated text: {generated_text}")
            return generated_text
        else:
            print(f"Unexpected pipeline output format: {output}", file=sys.stderr)
            return "Could not generate text. Unexpected output format from model."

    except Exception as e:
        print(f"An error occurred during image processing: {e}", file=sys.stderr)
        # Print the traceback for more detailed error info
        import traceback
        traceback.print_exc(file=sys.stderr)
        return f"An error occurred during image processing: {e}"


# Create the Gradio interface
# Input is an Image component, Output is a Textbox component
if model_loaded: # Only create interface if model loaded successfully
    interface = gr.Interface(
        fn=generate_text_from_image, # The function to run
        inputs=gr.Image(type="pil", label="Upload Image"), # Input component (Image)
        outputs=gr.Textbox(label="Generated Text"), # Output component (Textbox)
        title="Image to Text Generator", # Title of the app
        description="Upload an image and get a text description generated by a Hugging Face model." # Description
    )

    # Launch the Gradio interface
    if __name__ == "__main__":
        print("Launching Gradio interface...")
        interface.launch()
        print("Gradio interface launched.")
else:
    print("Gradio interface not launched due to model loading error.", file=sys.stderr)