app.py

import gradio as gr
import torch
import whisperx
import json
import os
import tempfile
from datetime import datetime
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import warnings
import gc
import psutil
import time
warnings.filterwarnings("ignore")

# === CONFIGURAÇÕES GLOBAIS ===
LANGUAGE = "pt"
TERMO_FIXO = ["CETOX", "CETOX31", "WhisperX", "VSL", "AI", "IA", "CPA", "CPM", "ROI", "ROAS"]
CORREÇÕES_ESPECÍFICAS = {
    "setox": "CETOX",
    "setox31": "CETOX 31",
    "SETOX": "CETOX",
    "SETOX31": "CETOX 31",
    "Setox": "CETOX",
    "Setox31": "CETOX 31",
    "vsl": "VSL",
    "VSl": "VSL",
    "vSL": "VSL"
}
MODEL_NAME = "unicamp-dl/ptt5-base-portuguese-vocab"

# Configurações por modelo (otimizadas para VSL de 13min)
MODEL_CONFIGS = {
    "large-v3": {
        "display_name": "🚀 Large-v3 (Máxima Precisão)",
        "description": "Melhor modelo disponível - ideal para VSL profissional",
        "score_minimo": 0.25,
        "batch_size": 4,
        "chunk_size": 30,
        "beam_size": 5,
        "best_of": 5,
        "temperature": 0.0,
        "recommended": True
    },
    "large-v2": {
        "display_name": "⚡ Large-v2 (Alta Precisão)",
        "description": "Excelente qualidade com boa velocidade",
        "score_minimo": 0.3,
        "batch_size": 6,
        "chunk_size": 30,
        "beam_size": 5,
        "best_of": 3,
        "temperature": 0.0,
        "recommended": False
    },
    "medium": {
        "display_name": "🏃 Medium (Rápido)",
        "description": "Boa qualidade, processamento mais rápido",
        "score_minimo": 0.35,
        "batch_size": 8,
        "chunk_size": 30,
        "beam_size": 3,
        "best_of": 3,
        "temperature": 0.1,
        "recommended": False
    },
    "turbo": {
        "display_name": "⚡ Turbo (Ultra Rápido)",
        "description": "Processamento mais rápido para testes",
        "score_minimo": 0.4,
        "batch_size": 12,
        "chunk_size": 30,
        "beam_size": 1,
        "best_of": 1,
        "temperature": 0.2,
        "recommended": False
    }
}

# === SETUP DISPOSITIVO ===
device = "cuda" if torch.cuda.is_available() else "cpu"
compute_type = "float16" if device == "cuda" else "int8"

# === MODELOS GLOBAIS (CACHE) ===
whisper_models = {}
align_model = None
metadata = None
corretor = None
corretor_disponivel = False

def get_system_info():
    """Retorna informações do sistema"""
    if torch.cuda.is_available():
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
        return f"{gpu_name} ({gpu_memory:.1f}GB)"
    else:
        ram = psutil.virtual_memory().total / 1024**3
        return f"CPU ({ram:.1f}GB RAM)"

def inicializar_modelos(modelo_selecionado, progress=gr.Progress()):
    """Inicializa os modelos necessários"""
    global whisper_models, align_model, metadata, corretor, corretor_disponivel
    
    try:
        config = MODEL_CONFIGS[modelo_selecionado]
        
        progress(0.1, desc=f"🔄 Carregando {config['display_name']}...")
        
        # Carregar WhisperX se não estiver em cache
        if modelo_selecionado not in whisper_models:
            whisper_models[modelo_selecionado] = whisperx.load_model(
                modelo_selecionado, 
                device, 
                compute_type=compute_type, 
                language=LANGUAGE,
                asr_options={
                    "beam_size": config["beam_size"],
                    "best_of": config["best_of"],
                    "temperature": config["temperature"],
                    "condition_on_previous_text": True,
                    "word_timestamps": True,
                    "prepend_punctuations": "\"'"¿([{-",
                    "append_punctuations": "\"'.。,，!！?？:：")]}、",
                    "vad_filter": True,
                    "vad_parameters": dict(min_silence_duration_ms=500)
                }
            )
        
        progress(0.3, desc="🎯 Carregando alinhamento temporal...")
        if align_model is None:
            align_model, metadata = whisperx.load_align_model(
                language_code=LANGUAGE, 
                device=device
            )
        
        progress(0.5, desc="📝 Carregando corretor PTT5...")
        if not corretor_disponivel:
            try:
                tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
                model_corr = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
                corretor = pipeline(
                    "text2text-generation", 
                    model=model_corr, 
                    tokenizer=tokenizer, 
                    device=0 if device == "cuda" else -1,
                    batch_size=4
                )
                corretor_disponivel = True
            except Exception as e:
                print(f"Correção desativada: {e}")
                corretor_disponivel = False
        
        progress(1.0, desc="✅ Modelos carregados!")
        
        system_info = get_system_info()
        return f"✅ **{config['display_name']} carregado!**\n\n🖥️ **Sistema:** {system_info}\n🎯 **Otimizado para:** VSL de 13 minutos"
        
    except Exception as e:
        return f"❌ Erro: {str(e)}"

def corrigir_palavra(palavra):
    """Corrige palavra com regras específicas para VSL"""
    if not palavra or not palavra.strip():
        return palavra
    
    palavra_limpa = palavra.strip()
    
    # Correções específicas para VSL
    if palavra_limpa in CORREÇÕES_ESPECÍFICAS:
        return CORREÇÕES_ESPECÍFICAS[palavra_limpa]
    
    # Não corrigir termos técnicos, números, URLs
    if (palavra_limpa.upper() in [t.upper() for t in TERMO_FIXO] or 
        palavra_limpa.isnumeric() or 
        len(palavra_limpa) <= 2 or
        "www." in palavra_limpa.lower() or
        "@" in palavra_limpa):
        return palavra_limpa
    
    if not corretor_disponivel:
        return palavra_limpa.capitalize()
    
    try:
        entrada = f"corrigir gramática: {palavra_limpa.lower()}"
        saida = corretor(entrada, max_length=50, do_sample=False, num_beams=2)[0]["generated_text"]
        resultado = saida.strip()
        return resultado.capitalize() if resultado else palavra_limpa.capitalize()
    except:
        return palavra_limpa.capitalize()

def processar_audio(audio_file, modelo_selecionado, progress=gr.Progress()):
    """Processa áudio com modelo selecionado"""
    if audio_file is None:
        return None, "❌ Faça upload do áudio da VSL."
    
    if modelo_selecionado not in MODEL_CONFIGS:
        return None, "❌ Modelo inválido selecionado."
    
    config = MODEL_CONFIGS[modelo_selecionado]
    start_time = time.time()
    
    try:
        # Verificar se modelo está carregado
        progress(0.05, desc="🔧 Verificando modelos...")
        if modelo_selecionado not in whisper_models:
            inicializar_modelos(modelo_selecionado)
        
        # Carregar áudio
        progress(0.1, desc="🎵 Carregando VSL...")
        audio = whisperx.load_audio(audio_file)
        duracao = len(audio) / 16000
        
        if duracao > 900:  # 15 minutos
            return None, f"⚠️ Áudio muito longo ({duracao/60:.1f}min). Máximo recomendado: 15min"
        
        progress(0.2, desc=f"🎤 Transcrevendo com {config['display_name']}...")
        
        # Transcrever com configurações otimizadas
        result = whisper_models[modelo_selecionado].transcribe(
            audio, 
            batch_size=config["batch_size"],
            chunk_size=config["chunk_size"],
            condition_on_previous_text=True,
            language=LANGUAGE
        )
        
        progress(0.6, desc="🎯 Alinhamento temporal de precisão...")
        
        # Alinhamento com configurações para VSL
        aligned = whisperx.align(
            result["segments"], 
            align_model, 
            metadata, 
            audio, 
            device,
            return_char_alignments=False,
            interpolate_method="linear",
            extend_duration=0.1
        )
        
        progress(0.8, desc="📝 Aplicando correções para VSL...")
        
        # Processar palavras
        resultado = []
        total_palavras = len(aligned.get("word_segments", []))
        palavras_processadas = 0
        
        for i, word in enumerate(aligned.get("word_segments", [])):
            if i % 15 == 0:
                progress(0.8 + (i / total_palavras) * 0.15, 
                        desc=f"📝 Processando {i+1}/{total_palavras} palavras")
            
            # Filtros otimizados para VSL
            if (word.get("score", 0) < config["score_minimo"] or 
                not word.get("word", "").strip() or
                len(word.get("word", "").strip()) < 1):
                continue
            
            palavra_original = word["word"].strip()
            palavra_corrigida = corrigir_palavra(palavra_original)
            palavras_processadas += 1
            
            resultado.append({
                "word": palavra_corrigida,
                "original": palavra_original,
                "start": round(word["start"], 3),
                "end": round(word["end"], 3),
                "score": round(word.get("score", 0), 3),
                "confidence": "high" if word.get("score", 0) > 0.8 else "medium" if word.get("score", 0) > 0.6 else "low"
            })
        
        progress(0.95, desc="💾 Gerando JSON final...")
        
        # Criar output otimizado para VSL
        processing_time = time.time() - start_time
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        output = {
            "metadata": {
                "timestamp": timestamp,
                "tipo_conteudo": "VSL",
                "duracao_audio": round(duracao, 2),
                "tempo_processamento": round(processing_time, 2),
                "velocidade_processamento": round(duracao / processing_time, 2),
                "total_words": len(resultado),
                "arquivo_original": os.path.basename(audio_file),
                "modelo_whisper": f"WhisperX {config['display_name']}",
                "modelo_correcao": MODEL_NAME if corretor_disponivel else "Sem correção",
                "configuracao": {
                    "score_minimo": config["score_minimo"],
                    "batch_size": config["batch_size"],
                    "beam_size": config["beam_size"],
                    "temperature": config["temperature"]
                },
                "sistema": get_system_info(),
                "otimizado_para": "VSL de até 15 minutos"
            },
            "words": resultado,
            "estatisticas": {
                "palavras_detectadas": len(resultado),
                "palavras_alta_confianca": len([w for w in resultado if w["confidence"] == "high"]),
                "palavras_media_confianca": len([w for w in resultado if w["confidence"] == "medium"]),
                "palavras_baixa_confianca": len([w for w in resultado if w["confidence"] == "low"]),
                "score_medio": round(sum(w["score"] for w in resultado) / len(resultado) if resultado else 0, 3),
                "score_minimo": round(min((w["score"] for w in resultado), default=0), 3),
                "score_maximo": round(max((w["score"] for w in resultado), default=0), 3),
                "densidade_palavras": round(len(resultado) / duracao * 60, 1),  # palavras por minuto
                "correções_aplicadas": sum(1 for w in resultado if w["word"] != w["original"])
            },
            "segmentos_temporais": [
                {
                    "inicio": f"{int(i*60//60):02d}:{int(i*60%60):02d}",
                    "palavras": len([w for w in resultado if i*60 <= w["start"] < (i+1)*60])
                }
                for i in range(int(duracao//60) + 1)
            ]
        }
        
        # Salvar arquivo
        temp_file = tempfile.NamedTemporaryFile(
            mode='w', 
            suffix=f'_VSL_transcricao_{timestamp}.json', 
            delete=False,
            encoding='utf-8'
        )
        
        json.dump(output, temp_file, ensure_ascii=False, indent=2)
        temp_file.close()
        
        # Limpeza de memória
        if device == "cuda":
            torch.cuda.empty_cache()
        gc.collect()
        
        progress(1.0, desc="✅ VSL transcrita com sucesso!")
        
        # Resumo otimizado
        resumo = f"""
✅ **VSL TRANSCRITA COM SUCESSO!**

🎯 **Modelo:** {config['display_name']}
⏱️ **Tempo:** {processing_time:.1f}s ({round(duracao/processing_time, 1)}x velocidade real)
🎵 **Duração:** {duracao/60:.1f} minutos

📊 **Resultados:**
- **{len(resultado)} palavras** detectadas
- **{output['estatisticas']['palavras_alta_confianca']} alta confiança** (score > 0.8)
- **{output['estatisticas']['densidade_palavras']} palavras/min**
- **{output['estatisticas']['correções_aplicadas']} correções** aplicadas

🎯 **Qualidade:**
- **Score médio:** {output['estatisticas']['score_medio']}
- **Precisão temporal:** ±100ms
- **Correções VSL:** CETOX, VSL automáticas

📥 **JSON pronto para download!**
        """
        
        return temp_file.name, resumo
        
    except Exception as e:
        error_msg = f"❌ Erro no processamento: {str(e)}"
        print(error_msg)
        return None, error_msg

def criar_interface():
    """Interface Gradio otimizada para VSL"""
    with gr.Blocks(
        title="🎤 Transcritor VSL Pro - WhisperX",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px;
            margin: auto;
        }
        .model-card {
            border: 2px solid #e1e5e9;
            border-radius: 8px;
            padding: 16px;
            margin: 8px 0;
        }
        .recommended {
            border-color: #10b981;
            background: linear-gradient(135deg, #ecfdf5 0%, #f0fdf4 100%);
        }
        """
    ) as demo:
        
        gr.Markdown("""
        # 🎤 Transcritor VSL Pro - WhisperX
        
        **Transcrição profissional para VSL com precisão temporal máxima**
        
        ✨ **Otimizado especialmente para:**
        - 🎯 **VSL de até 15 minutos**
        - 📺 **Conteúdo de marketing digital**
        - ⏱️ **Timestamps precisos palavra por palavra**
        - 🔧 **Correções automáticas (CETOX, VSL)**
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 🚀 Escolha do Modelo")
                
                # Seletor de modelo com descrições
                modelo_opcoes = []
                modelo_valores = []
                for key, config in MODEL_CONFIGS.items():
                    label = config['display_name']
                    if config['recommended']:
                        label += " ⭐ RECOMENDADO"
                    modelo_opcoes.append(label)
                    modelo_valores.append(key)
                
                modelo_selecionado = gr.Dropdown(
                    choices=list(zip(modelo_valores, modelo_opcoes)),
                    value="large-v3",
                    label="Modelo WhisperX",
                    info="Large-v3 recomendado para VSL profissional"
                )
                
                # Info do modelo selecionado
                with gr.Row():
                    modelo_info = gr.Markdown("""
                    **🚀 Large-v3 (Recomendado)**
                    - Máxima precisão para VSL
                    - Melhor detecção de palavras
                    - Timestamps ultra-precisos
                    """)
                
                gr.Markdown("### 📤 Upload da VSL")
                audio_input = gr.Audio(
                    label="Selecione o áudio da VSL (máx. 15min)",
                    type="filepath",
                    format="wav"
                )
                
                with gr.Row():
                    init_btn = gr.Button(
                        "🔧 Carregar Modelo", 
                        variant="secondary",
                        scale=1
                    )
                    processar_btn = gr.Button(
                        "🚀 Transcrever VSL", 
                        variant="primary",
                        scale=2
                    )
                
            with gr.Column(scale=1):
                gr.Markdown("### 📊 Status & Progresso")
                status_output = gr.Markdown("🟡 Selecione o modelo e faça upload da VSL...")
                
                gr.Markdown("### 💾 Download")
                file_output = gr.File(
                    label="📄 JSON da transcrição VSL",
                    interactive=False
                )
                
                # Sistema info
                system_info_display = gr.Markdown(f"🖥️ **Sistema:** {get_system_info()}")
        
        # Atualizar info do modelo
        def atualizar_info_modelo(modelo):
            if modelo in MODEL_CONFIGS:
                config = MODEL_CONFIGS[modelo]
                return f"""
                **{config['display_name']}**
                {config['description']}
                
                📊 **Configurações:**
                - Score mínimo: {config['score_minimo']}
                - Batch size: {config['batch_size']}
                - Beam size: {config['beam_size']}
                """
            return "Selecione um modelo"
        
        modelo_selecionado.change(
            fn=atualizar_info_modelo,
            inputs=[modelo_selecionado],
            outputs=[modelo_info]
        )
        
        # Eventos
        init_btn.click(
            fn=inicializar_modelos,
            inputs=[modelo_selecionado],
            outputs=[status_output]
        )
        
        processar_btn.click(
            fn=processar_audio,
            inputs=[audio_input, modelo_selecionado],
            outputs=[file_output, status_output]
        )
        
        # Informações técnicas
        with gr.Accordion("ℹ️ Especificações Técnicas", open=False):
            gr.Markdown(f"""
            ### 🔧 Configurações por Modelo
            
            | Modelo | Precisão | Velocidade | Uso Recomendado |
            |--------|----------|------------|-----------------|
            | **Large-v3** ⭐ | Máxima | Moderada | VSL profissional |
            | **Large-v2** | Alta | Boa | VSL geral |
            | **Medium** | Boa | Rápida | Testes rápidos |
            | **Turbo** | Básica | Ultra-rápida | Rascunhos |
            
            ### 🎯 Otimizações para VSL:
            - **VAD Filter:** Remove silêncios longos
            - **Chunks de 30s:** Processamento otimizado
            - **Correções específicas:** CETOX, VSL, termos de marketing
            - **Densidade de palavras:** Análise por minuto
            - **Confiança por palavra:** High/Medium/Low
            
            ### 📊 JSON de Saída:
            - Metadata completa da VSL
            - Timestamps precisos (±100ms)
            - Estatísticas de qualidade
            - Segmentação temporal
            - Análise de densidade
            
            **Sistema atual:** {get_system_info()}
            """)
    
    return demo

# === EXECUÇÃO ===
if __name__ == "__main__":
    print("🎤 Transcritor VSL Pro - WhisperX")
    print(f"🖥️  Sistema: {get_system_info()}")
    print("🎯 Otimizado para VSL de até 15 minutos")
    
    demo = criar_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        quiet=False,
        show_tips=True
    )