app.py

import gradio as gr
import torch
import os
import json
import numpy as np
from scipy import signal
import warnings
import requests
import tempfile
import shutil
from pathlib import Path
import traceback
warnings.filterwarnings("ignore")

# Imposta un seed per la riproducibilità
torch.manual_seed(42)

# Definizioni di variabili globali
MODEL_REPO = "Lorenzob/aurora-1.6b-complete"  # Repository del modello
CACHE_DIR = "./model_cache"  # Directory per la cache del modello
SAMPLE_RATE = 24000  # Frequenza di campionamento

# Cache per componenti del modello
processor = None
model = None
speaker_embeddings_cache = {}

def download_file(url, save_path):
    """Scarica un file da un URL"""
    response = requests.get(url, stream=True)
    response.raise_for_status()
    
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    with open(save_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            f.write(chunk)
    
    return save_path

def get_speaker_embeddings(speaker_id=0):
    """Ottieni gli speaker embeddings"""
    global speaker_embeddings_cache
    
    # Correggi l'indice dello speaker (gli embeddings disponibili sono numerati da 01 a 24)
    speaker_id = max(1, min(10, speaker_id + 1))
    
    if speaker_id in speaker_embeddings_cache:
        print(f"Usando embeddings in cache per speaker {speaker_id}")
        return speaker_embeddings_cache[speaker_id]
    
    try:
        url = f"https://huggingface.co/datasets/Matthijs/cmu-arctic-xvectors/resolve/main/cmu_us_{speaker_id:02d}_xvector.pt"
        tmp_dir = os.path.join(CACHE_DIR, "speakers")
        os.makedirs(tmp_dir, exist_ok=True)
        tmp_file = os.path.join(tmp_dir, f"speaker_{speaker_id:02d}.pt")
        
        if not os.path.exists(tmp_file):
            print(f"Scaricamento embeddings per speaker {speaker_id}...")
            download_file(url, tmp_file)
        
        print(f"Caricamento embeddings per speaker {speaker_id} da {tmp_file}")
        speaker_embeddings = torch.load(tmp_file)
        speaker_embeddings_cache[speaker_id] = speaker_embeddings
        return speaker_embeddings
    except Exception as e:
        print(f"Errore nel caricamento embeddings per speaker {speaker_id}: {e}")
        print("Utilizzo embeddings predefiniti")
        default_embeddings = torch.zeros((1, 512))
        speaker_embeddings_cache[speaker_id] = default_embeddings
        return default_embeddings

def fix_aurora_config():
    """Scarica e corregge la configurazione di Aurora"""
    config_url = f"https://huggingface.co/{MODEL_REPO}/resolve/main/config.json"
    local_config_path = os.path.join(CACHE_DIR, "config.json")
    
    # Crea la directory cache se non esiste
    os.makedirs(CACHE_DIR, exist_ok=True)
    
    try:
        # Scarica il file di configurazione
        print(f"Scaricamento della configurazione da {config_url}...")
        download_file(config_url, local_config_path)
        
        # Leggi il file di configurazione
        with open(local_config_path, "r") as f:
            config = json.load(f)
        
        # Modifica la configurazione per SpeechT5
        config["model_type"] = "speecht5"
        if "architectures" not in config or not config["architectures"]:
            config["architectures"] = ["SpeechT5ForTextToSpeech"]
        
        # Salva la configurazione modificata
        with open(local_config_path, "w") as f:
            json.dump(config, f, indent=2)
        
        print(f"Configurazione aggiornata salvata in {local_config_path}")
        return local_config_path
    except Exception as e:
        print(f"Errore nella configurazione del modello: {e}")
        return None

def load_aurora_model_manually():
    """Carica manualmente il modello Aurora-1.6b-complete"""
    global processor, model
    
    # Se il modello è già caricato, ritorna
    if model is not None and processor is not None:
        return model, processor
    
    try:
        print("🔄 Caricamento manuale del modello Aurora-1.6b-complete...")
        
        # Prima correggi la configurazione
        config_path = fix_aurora_config()
        if not config_path:
            raise ValueError("Impossibile correggere la configurazione del modello")
        
        # Importa le classi dopo la correzione della configurazione
        from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
        
        # Scarica i file del processor
        processor_files = {
            "tokenizer_config.json": f"https://huggingface.co/{MODEL_REPO}/resolve/main/tokenizer_config.json",
            "tokenizer.json": f"https://huggingface.co/{MODEL_REPO}/resolve/main/tokenizer.json",
            "special_tokens_map.json": f"https://huggingface.co/{MODEL_REPO}/resolve/main/special_tokens_map.json"
        }
        
        for filename, url in processor_files.items():
            local_path = os.path.join(CACHE_DIR, filename)
            if not os.path.exists(local_path):
                try:
                    print(f"Scaricamento di {filename}...")
                    download_file(url, local_path)
                except Exception as e:
                    print(f"Errore nel download di {filename}: {e}")
        
        # Carica il processor dalla directory cache
        try:
            # Usa il config modificato
            print("Tentativo di caricamento del processor dalla cache...")
            processor = SpeechT5Processor.from_pretrained(CACHE_DIR)
            print("Processor caricato con successo!")
        except Exception as e:
            print(f"Errore nel caricamento del processor: {e}")
            # Fallback a microsoft
            print("Utilizzo processor di Microsoft come fallback...")
            processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
        
        # Scarica i pesi del modello
        model_file = "model.safetensors"
        model_url = f"https://huggingface.co/{MODEL_REPO}/resolve/main/{model_file}"
        local_model_path = os.path.join(CACHE_DIR, model_file)
        
        if not os.path.exists(local_model_path):
            print(f"Scaricamento dei pesi del modello da {model_url}...")
            download_file(model_url, local_model_path)
        
        # Carica il modello usando il config modificato
        print("Caricamento del modello con config modificato...")
        model = SpeechT5ForTextToSpeech.from_pretrained(
            CACHE_DIR,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto"
        )
        
        print("✅ Modello e processor caricati con successo!")
        return model, processor
        
    except Exception as e:
        print(f"❌ Errore nel caricamento manuale di Aurora: {e}")
        traceback_str = traceback.format_exc()
        print(f"Traceback completo:\n{traceback_str}")
        
        # Se tutto fallisce, utilizza il modello Microsoft
        try:
            from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
            print("⚠️ Fallback al modello Microsoft...")
            processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
            model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
            print("✅ Modello Microsoft caricato come fallback")
        except Exception as e2:
            print(f"❌ Anche il fallback è fallito: {e2}")
            raise RuntimeError("Impossibile caricare alcun modello TTS")
        
        return model, processor

def text_to_speech(text, language="it", speaker_id=0, speed=1.0, show_log=True):
    """Converte testo in voce utilizzando Aurora-1.6b-complete"""
    if not text.strip():
        return None, "Per favore, inserisci del testo da convertire in voce."
    
    # Log di debug
    if show_log:
        print(f"Richiesta TTS ricevuta: '{text}' (Lingua: {language}, Speaker: {speaker_id}, Velocità: {speed})")
    
    try:
        # Carica il modello e il processor
        model, processor = load_aurora_model_manually()
        
        # Controlla se stiamo usando il modello Microsoft (fallback)
        is_microsoft_model = "microsoft" in str(type(model))
        
        # Ottieni gli speaker embeddings
        speaker_emb = get_speaker_embeddings(speaker_id)
        
        if is_microsoft_model:
            # Usa il modello Microsoft
            if show_log:
                print("Utilizzo del modello Microsoft SpeechT5 (fallback)...")
            
            # Crea input IDs dal testo
            inputs = processor(text=text, return_tensors="pt")
            
            # Genera l'audio
            with torch.no_grad():
                speech = model.generate_speech(
                    inputs["input_ids"], 
                    speaker_emb
                )
            
            # Imposta la frequenza di campionamento
            sample_rate = 16000  # Microsoft usa 16kHz
        else:
            # Usa il modello Aurora
            if show_log:
                print("Utilizzo del modello Aurora-1.6b-complete...")
            
            # Prepara gli input - IMPORTANTE: non includiamo 'language' che non è supportato
            inputs = processor(
                text=text,
                return_tensors="pt"
            )
            
            # Sposta gli input sul dispositivo di calcolo
            for k, v in inputs.items():
                if hasattr(v, "to"):
                    inputs[k] = v.to(model.device)
            
            # Sposta gli speaker embeddings sul dispositivo di calcolo
            if hasattr(model, "device"):
                speaker_emb = speaker_emb.to(model.device)
            
            # Mostra i dettagli degli inputs per debug
            if show_log:
                print(f"Input keys: {list(inputs.keys())}")
                print(f"Speaker embeddings shape: {speaker_emb.shape}")
            
            # Genera il speech usando generate_speech
            with torch.no_grad():
                if show_log:
                    print("Chiamata a model.generate_speech()...")
                speech = model.generate_speech(
                    inputs["input_ids"],
                    speaker_emb
                )
            
            # Imposta la frequenza di campionamento
            sample_rate = SAMPLE_RATE
        
        # Converti il tensore in un array numpy
        speech_array = speech.cpu().numpy()
        
        # Applica il controllo della velocità
        if speed != 1.0:
            # Usa scipy.signal per ricampionare l'audio e cambiare la velocità
            speech_array = signal.resample(speech_array, int(len(speech_array) / speed))
        
        if show_log:
            print(f"✅ Audio generato con successo! Lunghezza: {len(speech_array)} campioni")
        
        return (sample_rate, speech_array), None
    except Exception as e:
        error_msg = f"Errore nella generazione dell'audio: {str(e)}"
        traceback_str = traceback.format_exc()
        detailed_error = f"{error_msg}\n\nTraceback dettagliato:\n{traceback_str}"
        print(f"❌ {detailed_error}")
        return None, detailed_error

# Esempi predefiniti per l'interfaccia
examples = [
    ["Ciao, mi chiamo Aurora e sono un assistente vocale italiano.", "it", 0, 1.0, True],
    ["Hello, my name is Aurora and I'm an Italian voice assistant.", "en", 1, 1.0, True],
    ["Hola, me llamo Aurora y soy un asistente de voz italiano.", "es", 2, 1.0, True],
    ["La vita è bella e il sole splende nel cielo azzurro.", "it", 3, 1.0, True],
    ["Mi piace viaggiare e scoprire nuove città e culture.", "it", 4, 1.2, True],
    ["L'intelligenza artificiale sta trasformando il modo in cui interagiamo con i computer e con il mondo che ci circonda.", "it", 5, 0.9, True]
]

# Definizione dell'interfaccia Gradio
with gr.Blocks(title="Aurora-1.6b-complete TTS Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎙️ Aurora-1.6b-complete Text-to-Speech Demo
    
    Questa demo utilizza il modello **Aurora-1.6b-complete** per la sintesi vocale (TTS), un modello fine-tuned basato su Dia-1.6B.
    
    Puoi selezionare diversi stili di voce cambiando lo Speaker ID.
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="Testo da convertire in voce",
                placeholder="Inserisci qui il testo da convertire...",
                lines=5,
                value="Ciao, sono Aurora, un assistente vocale italiano basato su intelligenza artificiale."
            )
            
            with gr.Row():
                language_input = gr.Dropdown(
                    choices=["it", "en", "es", "fr", "de"], 
                    label="Lingua", 
                    value="it",
                    info="La lingua del testo (attualmente ignorata dal modello)"
                )
                speaker_input = gr.Slider(
                    label="Speaker ID",
                    value=0,
                    minimum=0,
                    maximum=9,
                    step=1,
                    info="ID dello speaker (0-9, ogni ID ha caratteristiche vocali diverse)"
                )
                speed_input = gr.Slider(
                    minimum=0.5, 
                    maximum=1.5, 
                    value=1.0, 
                    step=0.1, 
                    label="Velocità",
                    info="Valori più bassi = voce più lenta, valori più alti = voce più veloce"
                )
            
            debug_input = gr.Checkbox(label="Mostra log di debug", value=True)
            
            submit_btn = gr.Button("Genera Audio", variant="primary")
            
        with gr.Column(scale=1):
            audio_output = gr.Audio(label="Audio generato", show_share_button=True)
            error_output = gr.Textbox(label="Messaggi di errore", visible=True, lines=6)
    
    # Esempi
    gr.Examples(
        examples=examples,
        inputs=[text_input, language_input, speaker_input, speed_input, debug_input],
        outputs=[audio_output, error_output],
        fn=text_to_speech,
        cache_examples=True,
    )
    
    # Info aggiuntive
    gr.Markdown("""
    ## 📝 Note sull'utilizzo
    
    - Il modello funziona meglio con frasi di lunghezza media (fino a 20-30 parole)
    - Puoi cambiare lo Speaker ID per ottenere voci con caratteristiche diverse
    - La velocità di generazione dipende dalle risorse disponibili sul server
    - Il checkbox "Mostra log di debug" è utile per diagnosticare eventuali problemi
    
    ## 🔗 Crediti
    
    - [Lorenzob/aurora-1.6b-complete](https://huggingface.co/Lorenzob/aurora-1.6b-complete) (modello completo)
    - [nari-labs/Dia-1.6B](https://huggingface.co/nari-labs/Dia-1.6B) (modello base originale)
    - [CMU Arctic XVectors](https://huggingface.co/datasets/Matthijs/cmu-arctic-xvectors) (speaker embeddings)
    """)
    
    # Configurazione degli eventi
    submit_btn.click(
        fn=text_to_speech,
        inputs=[text_input, language_input, speaker_input, speed_input, debug_input],
        outputs=[audio_output, error_output],
    )

# Precarica il modello all'avvio
print("Inizializzazione del modello Aurora-1.6b-complete...")
try:
    load_aurora_model_manually()
except Exception as e:
    print(f"Errore nell'inizializzazione: {e}")

# Avvia l'interfaccia
demo.launch()