app.py

import spaces
from snac import SNAC
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import snapshot_download
from dotenv import load_dotenv
load_dotenv()

# Vérifier si CUDA est disponible
device = "cuda" if torch.cuda.is_available() else "cpu"

print("Chargement du modèle SNAC...")
snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
snac_model = snac_model.to(device)

model_name = "canopylabs/3b-fr-ft-research_release"

# Télécharger uniquement la configuration du modèle et les safetensors
snapshot_download(
    repo_id=model_name,
    allow_patterns=[
        "config.json",
        "*.safetensors",
        "model.safetensors.index.json",
    ],
    ignore_patterns=[
        "optimizer.pt",
        "pytorch_model.bin",
        "training_args.bin",
        "scheduler.pt",
        "tokenizer.json",
        "tokenizer_config.json",
        "special_tokens_map.json",
        "vocab.json",
        "merges.txt",
        "tokenizer.*"
    ]
)

model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
model.to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)
print(f"Modèle Orpheus chargé sur {device}")

# Traiter le texte d'entrée
def process_prompt(prompt, voice, tokenizer, device):
    prompt = f"{voice}: {prompt}"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    
    start_token = torch.tensor([[128259]], dtype=torch.int64)  # Début humain
    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.int64)  # Fin du texte, Fin humain
    
    modified_input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)  # SOH SOT Texte EOT EOH
    
    # Pas besoin de padding pour une seule entrée
    attention_mask = torch.ones_like(modified_input_ids)
    
    return modified_input_ids.to(device), attention_mask.to(device)

# Analyser les tokens de sortie en audio
def parse_output(generated_ids):
    token_to_find = 128257
    token_to_remove = 128258
    
    token_indices = (generated_ids == token_to_find).nonzero(as_tuple=True)

    if len(token_indices[1]) > 0:
        last_occurrence_idx = token_indices[1][-1].item()
        cropped_tensor = generated_ids[:, last_occurrence_idx+1:]
    else:
        cropped_tensor = generated_ids

    processed_rows = []
    for row in cropped_tensor:
        masked_row = row[row != token_to_remove]
        processed_rows.append(masked_row)

    code_lists = []
    for row in processed_rows:
        row_length = row.size(0)
        new_length = (row_length // 7) * 7
        trimmed_row = row[:new_length]
        trimmed_row = [t - 128266 for t in trimmed_row]
        code_lists.append(trimmed_row)
        
    return code_lists[0]  # Retourner uniquement le premier pour un seul échantillon

# Redistribuer les codes pour la génération audio
def redistribute_codes(code_list, snac_model):
    device = next(snac_model.parameters()).device  # Obtenir le périphérique du modèle SNAC
    
    layer_1 = []
    layer_2 = []
    layer_3 = []
    for i in range((len(code_list)+1)//7):
        layer_1.append(code_list[7*i])
        layer_2.append(code_list[7*i+1]-4096)
        layer_3.append(code_list[7*i+2]-(2*4096))
        layer_3.append(code_list[7*i+3]-(3*4096))
        layer_2.append(code_list[7*i+4]-(4*4096))
        layer_3.append(code_list[7*i+5]-(5*4096))
        layer_3.append(code_list[7*i+6]-(6*4096))
        
    # Déplacer les tenseurs vers le même périphérique que le modèle SNAC
    codes = [
        torch.tensor(layer_1, device=device).unsqueeze(0),
        torch.tensor(layer_2, device=device).unsqueeze(0),
        torch.tensor(layer_3, device=device).unsqueeze(0)
    ]
    
    audio_hat = snac_model.decode(codes)
    return audio_hat.detach().squeeze().cpu().numpy()  # Toujours retourner un tableau numpy CPU

# Fonction principale de génération
@spaces.GPU()
def generate_speech(text, voice, temperature, top_p, repetition_penalty, max_new_tokens, progress=gr.Progress()):
    if not text.strip():
        return None
    
    try:
        progress(0.1, "Traitement du texte...")
        input_ids, attention_mask = process_prompt(text, voice, tokenizer, device)
        
        progress(0.3, "Génération des tokens de parole...")
        with torch.no_grad():
            generated_ids = model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                temperature=temperature,
                top_p=top_p,
                repetition_penalty=repetition_penalty,
                num_return_sequences=1,
                eos_token_id=128258,
            )
        
        progress(0.6, "Traitement des tokens de parole...")
        code_list = parse_output(generated_ids)
        
        progress(0.8, "Conversion en audio...")
        audio_samples = redistribute_codes(code_list, snac_model)
        
        return (24000, audio_samples)  # Retourner le taux d'échantillonnage et l'audio
    except Exception as e:
        print(f"Erreur lors de la génération de la parole: {e}")
        return None

# Exemples pour l'interface utilisateur
examples = [
    ["Bonjour, je m'appelle Tara, <chuckle> et je suis un modèle de génération de parole qui peut ressembler à une personne.", "tara", 0.6, 0.95, 1.1, 1200],
    ["On m'a aussi appris à comprendre et à produire des éléments paralinguistiques <sigh> comme soupirer, ou <laugh> rire, ou <yawn> bâiller !", "dan", 0.7, 0.95, 1.1, 1200],
    ["J'habite à San Francisco, et j'ai, euh voyons voir, 3 milliards 7 cents... <gasp> bon, disons simplement beaucoup de paramètres.", "leah", 0.6, 0.9, 1.2, 1200],
    ["Parfois, quand je parle trop, j'ai besoin de <cough> m'excuser. <sniffle> Le temps a été assez froid dernièrement.", "leo", 0.65, 0.9, 1.1, 1200],
    ["Parler en public peut être difficile. <groan> Mais avec suffisamment de pratique, n'importe qui peut s'améliorer.", "jess", 0.7, 0.95, 1.1, 1200],
    ["La randonnée était épuisante mais la vue depuis le sommet était absolument à couper le souffle ! <sigh> Ça valait totalement le coup.", "mia", 0.65, 0.9, 1.15, 1200],
    ["Tu as entendu cette blague ? <laugh> Je n'ai pas pu m'arrêter de rire quand je l'ai entendue pour la première fois. <chuckle> C'est toujours drôle.", "zac", 0.7, 0.95, 1.1, 1200],
    ["Après avoir couru le marathon, j'étais tellement fatigué <yawn> et j'avais besoin d'un long repos. <sigh> Mais je me sentais accompli.", "zoe", 0.6, 0.95, 1.1, 1200]
]

# Voix disponibles
VOICES = ["tara", "leah", "jess", "leo", "dan", "mia", "zac", "zoe"]

# Balises émotives disponibles
EMOTIVE_TAGS = ["`<laugh>`", "`<chuckle>`", "`<sigh>`", "`<cough>`", "`<sniffle>`", "`<groan>`", "`<yawn>`", "`<gasp>`"]

# Créer l'interface Gradio
with gr.Blocks(title="Orpheus Text-to-Speech") as demo:
    gr.Markdown(f"""
    # 🎵 [Orpheus French Text-to-Speech](https://github.com/canopyai/Orpheus-TTS)
    Saisissez votre texte ci-dessous et écoutez-le transformé en parole naturelle avec le modèle Orpheus TTS.
    
    ## Conseils pour de meilleurs résultats:
    - Ajoutez des éléments paralinguistiques comme {", ".join(EMOTIVE_TAGS)} ou `euh` pour une parole plus humaine.
    - Les textes plus longs fonctionnent généralement mieux que les phrases très courtes
    - Augmenter `repetition_penalty` et `temperature` fait parler le modèle plus rapidement.
    """)    
    with gr.Row():
        with gr.Column(scale=3):
            text_input = gr.Textbox(
                label="Texte à prononcer", 
                placeholder="Entrez votre texte ici...",
                lines=5
            )
            voice = gr.Dropdown(
                choices=VOICES, 
                value="tara", 
                label="Voix"
            )
            
            with gr.Accordion("Paramètres avancés", open=False):
                temperature = gr.Slider(
                    minimum=0.1, maximum=1.5, value=0.6, step=0.05,
                    label="Température", 
                    info="Des valeurs plus élevées (0.7-1.0) créent une parole plus expressive mais moins stable"
                )
                top_p = gr.Slider(
                    minimum=0.1, maximum=1.0, value=0.95, step=0.05,
                    label="Top P", 
                    info="Seuil d'échantillonnage du noyau"
                )
                repetition_penalty = gr.Slider(
                    minimum=1.0, maximum=2.0, value=1.1, step=0.05,
                    label="Pénalité de répétition", 
                    info="Des valeurs plus élevées découragent les motifs répétitifs"
                )
                max_new_tokens = gr.Slider(
                    minimum=100, maximum=2000, value=1200, step=100,
                    label="Longueur maximale", 
                    info="Longueur maximale de l'audio généré (en tokens)"
                )
            
            with gr.Row():
                submit_btn = gr.Button("Générer la parole", variant="primary")
                clear_btn = gr.Button("Effacer")
                
        with gr.Column(scale=2):
            audio_output = gr.Audio(label="Parole générée", type="numpy")
            
    # Configuration des exemples
    gr.Examples(
        examples=examples,
        inputs=[text_input, voice, temperature, top_p, repetition_penalty, max_new_tokens],
        outputs=audio_output,
        fn=generate_speech,
        cache_examples=True,
    )
    
    # Configuration des gestionnaires d'événements
    submit_btn.click(
        fn=generate_speech,
        inputs=[text_input, voice, temperature, top_p, repetition_penalty, max_new_tokens],
        outputs=audio_output
    )
    
    clear_btn.click(
        fn=lambda: (None, None),
        inputs=[],
        outputs=[text_input, audio_output]
    )

# Lancer l'application
if __name__ == "__main__":
    demo.queue().launch(share=False, ssr_mode=False)