import gradio as gr
from transformers import AutoProcessor, AutoModelForCausalLM
from PIL import Image

# Load the model and processor
processor = AutoProcessor.from_pretrained("Aman1212222/Image_to_ballad")
model = AutoModelForCausalLM.from_pretrained("Aman1212222/Image_to_ballad")

# Prediction function
def generate_caption(image):
    # Process the image
    inputs = processor(images=image, return_tensors="pt")
    pixel_values = inputs.pixel_values

    # Generate the caption
    generated_ids = model.generate(pixel_values=pixel_values, max_length=100)
    generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    
    return generated_caption

# Enhanced CSS for WOW factor
css_style = """
body {
    background: linear-gradient(135deg, #0f2027, #203a43, #2c5364);
    font-family: 'Roboto', sans-serif;
    margin: 0;
    padding: 0;
    display: flex;
    justify-content: center;
    align-items: center;
    height: 100vh;
}

.gradio-container {
    backdrop-filter: blur(10px);
    background: rgba(255, 255, 255, 0.1);
    border-radius: 15px;
    padding: 30px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.2);
    color: white;
    text-align: center;
    width: 500px;
}

.gradio-title {
    font-size: 36px;
    font-weight: bold;
    letter-spacing: 1px;
    margin-bottom: 20px;
    color: #ffffff;
    text-shadow: 2px 2px 8px rgba(0, 0, 0, 0.6);
}

.gradio-description {
    font-size: 18px;
    margin-bottom: 30px;
    color: #ffffff;
    text-shadow: 1px 1px 5px rgba(0, 0, 0, 0.4);
}

.gradio-input, .gradio-output {
    background: rgba(255, 255, 255, 0.2);
    border: none;
    border-radius: 10px;
    padding: 15px;
    margin-bottom: 20px;
    transition: transform 0.3s ease;
}

.gradio-input:hover, .gradio-output:hover {
    transform: scale(1.05);
}

.gr-button {
    background-color: #ff5733;
    color: white;
    border-radius: 12px;
    padding: 12px 24px;
    font-size: 18px;
    font-weight: bold;
    cursor: pointer;
    transition: background-color 0.3s ease, transform 0.3s ease;
    box-shadow: 0 6px 12px rgba(0, 0, 0, 0.2);
}

.gr-button:hover {
    background-color: #ff2e00;
    transform: translateY(-4px);
}

footer {
    display: none !important;
}
"""

# Create Gradio Interface
interface = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="✨ Image to Ballad Generator ✨",
    description="Upload an image, and we'll generate a ballad for it with our AI-powered model!",
    theme="default",
    css=css_style
)

# Launch the interface with a public link
interface.launch(share=True)