Fix model recognition error and improve error handling

app.py

@@ -2,7 +2,8 @@
 import gradio as gr
 import torch
 import os
-from transformers import AutoProcessor, SpeechT5ForTextToSpeech, set_seed
+import json
+from transformers import AutoProcessor, AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, SpeechT5ForTextToSpeech, set_seed
 import numpy as np
 from scipy import signal
 import warnings
@@ -12,38 +13,98 @@ warnings.filterwarnings("ignore")
 set_seed(42)
 
 # Definizioni di variabili globali
-MODEL_REPO = "Lorenzob/aurora-1.6b-complete"  # Repository aggiornata con il modello completo
+MODEL_REPO = "Lorenzob/aurora-1.6b-complete"  # Repository principale
+FALLBACK_REPO = "nari-labs/Dia-1.6B"  # Repository di fallback
 SAMPLE_RATE = 24000  # Frequenza di campionamento per il modello TTS
 
 # Cache per il modello e il processor (per evitare di ricaricarli ad ogni richiesta)
 model = None
 processor = None
 
+def fix_model_config(model_path):
+    """Aggiunge il model_type alla configurazione se necessario"""
+    try:
+        config_path = os.path.join(model_path, "config.json")
+        if os.path.exists(config_path):
+            with open(config_path, "r") as f:
+                config = json.load(f)
+            
+            # Aggiungi model_type se mancante
+            if "model_type" not in config:
+                config["model_type"] = "speecht5"
+                print(f"Aggiunto model_type 'speecht5' alla configurazione")
+            
+            # Aggiungi architectures se mancante
+            if "architectures" not in config:
+                config["architectures"] = ["SpeechT5ForTextToSpeech"]
+                print(f"Aggiunto architectures al config")
+            
+            # Salva la configurazione aggiornata
+            with open(config_path, "w") as f:
+                json.dump(config, f, indent=2)
+            
+            print(f"Configurazione aggiornata e salvata in {config_path}")
+            return True
+        else:
+            print(f"File di configurazione non trovato in {model_path}")
+            return False
+    except Exception as e:
+        print(f"Errore nella modifica del config: {e}")
+        return False
+
 def load_model_and_processor():
-    """Carica il modello e il processor solo se non sono già stati caricati"""
+    """Carica il modello e il processor con gestione degli errori avanzata"""
     global model, processor
     
     if model is None or processor is None:
         try:
-            print("📂 Caricamento del modello Aurora-1.6b-complete...")
+            print(f"📂 Tentativo di caricamento del modello da {MODEL_REPO}...")
+            
+            # Prova a caricare il processor
             processor = AutoProcessor.from_pretrained(MODEL_REPO)
+            
+            # Carica il modello specificando esplicitamente la classe
             model = SpeechT5ForTextToSpeech.from_pretrained(
                 MODEL_REPO, 
                 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                 device_map="auto"
             )
-            print("✅ Modello caricato con successo!")
+            print("✅ Modello principale caricato con successo!")
+            
         except Exception as e:
-            print(f"❌ Errore nel caricamento del modello: {e}")
-            # Fallback al modello originale di Dia se il caricamento fallisce
-            print("⚠️ Tentativo di fallback al modello Dia-1.6B...")
-            processor = AutoProcessor.from_pretrained("nari-labs/Dia-1.6B")
-            model = SpeechT5ForTextToSpeech.from_pretrained(
-                "nari-labs/Dia-1.6B",
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                device_map="auto"
-            )
-            print("✅ Modello di fallback caricato con successo!")
+            print(f"❌ Errore nel caricamento del modello principale: {e}")
+            
+            # Prova con il modello di fallback
+            try:
+                print(f"⚠️ Tentativo di caricamento del modello di fallback da {FALLBACK_REPO}...")
+                
+                # Prova a caricare il processor di fallback
+                processor = AutoProcessor.from_pretrained(FALLBACK_REPO)
+                
+                # Carica il modello di fallback specificando esplicitamente la classe
+                model = SpeechT5ForTextToSpeech.from_pretrained(
+                    FALLBACK_REPO,
+                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                    device_map="auto"
+                )
+                print("✅ Modello di fallback caricato con successo!")
+                
+            except Exception as e2:
+                print(f"❌ Errore anche nel caricamento del modello di fallback: {e2}")
+                
+                # Se entrambi i tentativi falliscono, prova con un modello TTS generico ben supportato
+                try:
+                    print("🔄 Tentativo con un modello TTS generico (microsoft/speecht5_tts)...")
+                    processor = AutoProcessor.from_pretrained("microsoft/speecht5_tts")
+                    model = SpeechT5ForTextToSpeech.from_pretrained(
+                        "microsoft/speecht5_tts",
+                        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                        device_map="auto"
+                    )
+                    print("✅ Modello generico caricato con successo!")
+                except Exception as e3:
+                    print(f"❌ Tutti i tentativi di caricamento sono falliti: {e3}")
+                    raise RuntimeError("Impossibile caricare alcun modello TTS")
     
     return model, processor
 
@@ -63,10 +124,13 @@ def text_to_speech(text, language="it", speaker_id=0, speed=1.0, show_log=False)
         # Prepara gli input per il modello
         inputs = processor(
             text=text, 
-            language=language,
             return_tensors="pt"
         )
         
+        # Aggiungi il parametro di lingua se supportato dal processor
+        if "language" in processor.model_input_names:
+            inputs["language"] = language
+        
         # Sposta gli input sul dispositivo di calcolo
         for k, v in inputs.items():
             if hasattr(v, "to"):
@@ -82,7 +146,11 @@ def text_to_speech(text, language="it", speaker_id=0, speed=1.0, show_log=False)
         
         # Genera il speech
         with torch.no_grad():
-            speech = model.generate(**inputs, **gen_params)
+            # Passa speaker_embeddings se disponibili/necessari
+            if hasattr(model, "generate_speech"):
+                speech = model.generate_speech(**inputs, **gen_params)
+            else:
+                speech = model.generate(**inputs, **gen_params)
         
         # Converti il tensore in un array numpy
         speech_array = speech.cpu().numpy().squeeze()
@@ -116,7 +184,7 @@ with gr.Blocks(title="Aurora-1.6b TTS Demo", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
     # 🎙️ Aurora-1.6b 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 con pesi completi.
+    Questa demo utilizza il modello Aurora-1.6b-complete per la sintesi vocale (TTS), un modello fine-tuned basato su Dia-1.6B.
     
     Il modello supporta italiano, inglese, spagnolo, francese e tedesco, ma è stato ottimizzato per l'italiano.
     """)