import torch
import gradio as gr
from PIL import Image
import scipy.io.wavfile as wavfile

# Use a pipeline as a high-level helper
from transformers import pipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize pipelines
caption_image = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large", device=device)
narrator = pipeline("text-to-speech", model="kakao-enterprise/vits-ljs")

def generate_audio(text):
    # Generate the narrated text without max_new_tokens
    narrated_text = narrator(text)
    
    # Save the audio to a WAV file
    audio_path = "output.wav"
    wavfile.write(audio_path, rate=narrated_text["sampling_rate"], data=narrated_text["audio"][0])
    return audio_path

def truncate_text(text, max_length=200):
    return text[:max_length] + '...' if len(text) > max_length else text

def caption_my_image(pil_image):
    # Generate the caption from the image
    semantics = caption_image(images=pil_image)[0]['generated_text']
    
    # Optionally truncate the text
    truncated_semantics = truncate_text(semantics, max_length=200)
    
    # Generate the corresponding audio
    audio_path = generate_audio(truncated_semantics)
    return semantics, audio_path

# Gradio interface
demo = gr.Interface(
    fn=caption_my_image,
    inputs=[gr.Image(label="📸 Select Image", type="pil")],
    outputs=[gr.Textbox(label="📝 Generated Caption"), gr.Audio(label="🔊 Audio Caption")],
    title="🖼️ SM Project: Image Captioning",
    description=(
        "✨ **Welcome to the Image Captioning App!** ✨\n\n"
        "This application will allow you to:\n"
        "1. **Upload an Image** 📷\n"
        "2. **Generate a Caption** 📝\n"
        "3. **Listen to the Audio Caption** 🔊\n\n"
        "👉 **Instructions**:\n"
        "1. Click on the 'Select Image' button to upload an image.\n"
        "2. Wait for the application to generate a caption for your image.\n"
        "3. Enjoy the narrated audio of the caption!\n\n"
        "🎉 **Let's get started!** 🎉\n\n"
        "🖊️ **Created by SHASHWAT MISHRA**"
    ),
)

# Launch the demo
demo.launch()