import gradio as gr
from PIL import Image
from transformers import pipeline
import torch
from diffusers import AudioLDMPipeline

# =========================
# COMPROBACIÓN DE DISPOSITIVO
# =========================
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Usando dispositivo: {device}")

# =========================
# MODELOS
# =========================

# Imagen -> Texto
modelo_imagen_texto = pipeline(
    "image-to-text",
    model="Salesforce/blip-image-captioning-base",
    device=0 if device=="cuda" else -1
)

# Texto -> Audio (AudioLDMPipeline)
pipe_audio = AudioLDMPipeline.from_pretrained(
    "cvssp/audioldm-s-full-v2",
    torch_dtype=torch.float16 if device=="cuda" else torch.float32
).to(device)

# =========================
# FUNCIÓN PRINCIPAL
# =========================
def imagen_a_audio(imagen):
    """
    1. Convierte la imagen en descripción de texto.
    2. Convierte la descripción en audio usando AudioLDMPipeline.
    """
    # Imagen -> Texto
    resultado = modelo_imagen_texto(Image.fromarray(imagen))
    descripcion = resultado[0]["generated_text"]
    print(f"Descripción generada: {descripcion}")

    # Texto -> Audio
    audio_result = pipe_audio(
        prompt=descripcion,
        num_inference_steps=50,
        audio_length_in_s=5.0
    )

    # El resultado contiene .audios
    audio_array = audio_result.audios[0]
    
    # AudioLDMPipeline usa 22050 Hz como sampling rate por defecto
    sampling_rate = 22050

    return descripcion, (sampling_rate, audio_array)

# =========================
# INTERFAZ GRADIO
# =========================
demo = gr.Interface(
    fn=imagen_a_audio,
    inputs=gr.Image(type="numpy", label="Sube una imagen"),
    outputs=[
        gr.Textbox(label="Descripción de la imagen"),
        gr.Audio(label="Audio generado")
    ],
    title="Imagen a Audio para accesibilidad",
    description=(
        "Esta aplicación analiza una imagen y genera una descripción textual. "
        "Luego crea un audio a partir de esa descripción usando AudioLDMPipeline, "
        "para ayudar a personas con discapacidad visual a percibir el contenido visual."
    ),
)

demo.launch(share=True)