app.py

import subprocess
import os
import shutil 
from pathlib import Path
from datasets import load_dataset, Audio
import gradio as gr
import time

# ----------------------------------------------------------------------
# 1. FUNCIONES DE DESCARGA Y PROCESAMIENTO
# ----------------------------------------------------------------------

def download_clip(
    video_identifier,
    output_filename,
    start_time,
    end_time,
    num_attempts=3,
    url_base='https://www.youtube.com/watch?v='
):
    """Descarga un clip de YouTube usando yt-dlp con mejor manejo de errores."""
    
    # Comando mejorado con más opciones de compatibilidad
    command = [
        "yt-dlp",
        "--quiet",
        "--force-keyframes-at-cuts",
        "--no-warnings",
        "-x",
        "--audio-format", "wav",
        "-f", "bestaudio/best",
        "--no-check-certificate",  # Evita problemas de SSL
        "--prefer-insecure",
        "--socket-timeout", "30",
        "--retries", "3",
        "-o", output_filename,
        "--download-sections", f"*{start_time}-{end_time}",
        f"{url_base}{video_identifier}"
    ]

    for attempt in range(num_attempts):
        try:
            result = subprocess.run(
                command, 
                capture_output=True, 
                text=True,
                timeout=60  # Timeout de 60 segundos
            )
            
            # Verificar si el archivo existe y tiene tamaño > 0
            if os.path.exists(output_filename) and os.path.getsize(output_filename) > 0:
                return True, 'Downloaded successfully'
            
            if attempt < num_attempts - 1:
                time.sleep(2)  # Esperar antes de reintentar
                
        except subprocess.TimeoutExpired:
            print(f"Timeout en intento {attempt + 1} para {video_identifier}")
            if attempt < num_attempts - 1:
                time.sleep(2)
        except Exception as e:
            print(f"Error en intento {attempt + 1} para {video_identifier}: {str(e)}")
            if attempt < num_attempts - 1:
                time.sleep(2)
    
    # Si llegamos aquí, la descarga falló
    error_msg = f"Failed after {num_attempts} attempts"
    return False, error_msg


def main(
    data_dir: str,
    sampling_rate: int = 44100,
    limit: int = None,
    num_proc: int = 1,
    writer_batch_size: int = 1000,
):
    """
    Descarga los clips de MusicCaps, procesa el dataset y maneja la limpieza.
    """

    data_dir_path = Path(data_dir)
    
    # Limpieza de disco
    if data_dir_path.exists():
        shutil.rmtree(data_dir_path)
        print(f"Limpiando el directorio {data_dir} para liberar espacio.")
        
    data_dir_path.mkdir(exist_ok=True, parents=True)

    print("Iniciando carga del dataset MusicCaps...")
    ds = load_dataset('google/MusicCaps', split='train')
    
    if limit is not None:
        print(f"Limitando a {limit} ejemplos")
        ds = ds.select(range(limit))

    def process(example):
        outfile_path = str(data_dir_path / f"{example['ytid']}.wav")
        
        status = False
        if not os.path.exists(outfile_path):
            status, log = download_clip(
                example['ytid'],
                outfile_path,
                example['start_s'],
                example['end_s'],
            )
            if not status:
                print(f"❌ Falló descarga de {example['ytid']}: {log}")
        else:
            # Si el archivo existe, verificar que tenga contenido
            status = os.path.getsize(outfile_path) > 0

        example['audio'] = outfile_path if status else ""
        example['download_status'] = status
        return example

    # Descargar y actualizar el dataset
    ds = ds.map(
        process,
        num_proc=num_proc,
        writer_batch_size=writer_batch_size,
        keep_in_memory=False
    )
    
    return ds

# ----------------------------------------------------------------------
# 2. CARGA DEL DATASET AL INICIO Y FILTRADO
# ----------------------------------------------------------------------

print("=" * 60)
print("INICIANDO APLICACIÓN MUSICCAPS DEMO")
print("=" * 60)

# Intentar con más ejemplos para aumentar probabilidad de éxito
ds = main('./music_data', num_proc=1, limit=30)  # Aumentado a 30
print(f"Dataset cargado. Total de ejemplos antes del filtro: {len(ds)}")

# Filtrar ejemplos válidos
ds = ds.filter(lambda example: example['download_status'] == True)
print(f"✅ Ejemplos válidos después del filtro: {len(ds)}")

# Si no hay ejemplos válidos, intentar con otro rango
if len(ds) == 0:
    print("⚠️  Primer intento falló. Intentando con ejemplos 100-130...")
    full_ds = load_dataset('google/MusicCaps', split='train')
    ds = full_ds.select(range(100, 130))
    
    data_dir_path = Path('./music_data')
    def process(example):
        outfile_path = str(data_dir_path / f"{example['ytid']}.wav")
        status = False
        if not os.path.exists(outfile_path):
            status, log = download_clip(
                example['ytid'],
                outfile_path,
                example['start_s'],
                example['end_s'],
            )
        example['audio'] = outfile_path if status else ""
        example['download_status'] = status
        return example
    
    ds = ds.map(process, num_proc=1, keep_in_memory=False)
    ds = ds.filter(lambda example: example['download_status'] == True)
    print(f"✅ Ejemplos válidos en segundo intento: {len(ds)}")

# Aplicar tipo Audio solo si hay ejemplos válidos
if len(ds) > 0:
    ds = ds.cast_column('audio', Audio(sampling_rate=44100))
    print(f"🎵 Dataset preparado con {len(ds)} ejemplos")

# ----------------------------------------------------------------------
# 3. INTERFAZ GRADIO
# ----------------------------------------------------------------------

def get_example(idx):
    """Función que Gradio llama al mover el slider."""
    idx = int(idx)
    if idx >= len(ds):
        idx = len(ds) - 1
    
    ex = ds[idx]
    return ex['audio']['path'], ex['caption']

# Crear interfaz solo si hay datos
if len(ds) > 0:
    demo = gr.Interface(
        get_example,
        inputs=gr.Slider(
            0, 
            len(ds) - 1, 
            value=0, 
            step=1, 
            label=f"Índice del Ejemplo (0 a {len(ds) - 1})" 
        ),
        outputs=[
            gr.Audio(label="🎵 Muestra de Audio"), 
            gr.Textbox(label="📝 Descripción (Caption)", lines=3)
        ],
        title="🎵 MusicCaps Dataset Explorer",
        description=f"Explora {len(ds)} clips de música con sus descripciones. Mueve el slider para cambiar de ejemplo.",
        live=True,
    )
    
    print("🚀 Lanzando aplicación Gradio...")
    demo.launch()
else:
    print("=" * 60)
    print("❌ ERROR CRÍTICO")
    print("=" * 60)
    print("No se pudo descargar ningún ejemplo válido del dataset.")
    print("Posibles causas:")
    print("1. Problemas de conectividad en Hugging Face Spaces")
    print("2. YouTube bloqueando las descargas")
    print("3. Videos no disponibles o privados")
    print("\nSugerencias:")
    print("- Verificar que yt-dlp esté actualizado")
    print("- Intentar con un VPN o proxy")
    print("- Usar un dataset pre-descargado")
    print("=" * 60)
    
    # Crear una interfaz mínima de error
    demo = gr.Interface(
        lambda: ("⚠️ No hay datos disponibles", "No se pudieron descargar clips de audio"),
        inputs=None,
        outputs=[gr.Textbox(label="Estado"), gr.Textbox(label="Mensaje")],
        title="❌ MusicCaps Demo - Error",
        description="No se pudieron descargar ejemplos del dataset. Verifica los logs para más detalles."
    )
    demo.launch()