Update app.py

app.py

@@ -0,0 +1,137 @@
+from spaces import GPU
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+from snac import SNAC
+import gradio as gr
+
+# Config
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+base_model_id = "canopylabs/3b-es_it-pretrain-research_release"
+lora_model_id = "sirekist98/spanish_conversational_tts"
+snac_model_id = "hubertsiuzdak/snac_24khz"
+
+# Load models
+tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_id,
+    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
+)
+model = PeftModel.from_pretrained(base_model, lora_model_id)
+model = model.to(device)
+model.eval()
+
+snac_model = SNAC.from_pretrained(snac_model_id).to(device)
+
+# Speakers (sin emociones)
+speakers = [
+    "Alex", "Carmen", "Daniel", "Diego", "Hugo", "Lucía", "María", "Pablo", "Sofía"
+]
+
+# Helper to decode tokens to audio
+
+def decode_snac(code_list):
+    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))
+
+    # Obtener dispositivo del primer codebook
+    device_snac = snac_model.quantizer.quantizers[0].codebook.weight.device
+
+    layers = [
+        torch.tensor(layer_1).unsqueeze(0).to(device_snac),
+        torch.tensor(layer_2).unsqueeze(0).to(device_snac),
+        torch.tensor(layer_3).unsqueeze(0).to(device_snac),
+    ]
+
+    with torch.no_grad():
+        audio = snac_model.decode(layers).squeeze().cpu().numpy()
+    return audio
+
+
+# Inference (sin emociones)
+@GPU
+def tts(prompt, speaker):
+    # Estructura de prompt: "<SPEAKER>: <texto>"
+    full_prompt = f"{speaker}: {prompt}"
+
+    input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids.to(device)
+
+    # Tokens especiales (iguales que tu versión anterior)
+    start_token = torch.tensor([[128259]], dtype=torch.long).to(device)
+    end_tokens = torch.tensor([[128009, 128260]], dtype=torch.long).to(device)
+
+    input_ids = torch.cat([start_token, input_ids, end_tokens], dim=1)
+
+    # Padding fijo a 4260 para que encaje con el entrenamiento
+    padding_len = max(0, 4260 - input_ids.shape[1])
+    if padding_len > 0:
+        pad = torch.full((1, padding_len), 128263, dtype=torch.long).to(device)
+        input_ids = torch.cat([pad, input_ids], dim=1)
+        attention_mask = torch.cat([
+            torch.zeros((1, padding_len), dtype=torch.long),
+            torch.ones((1, input_ids.shape[1]-padding_len), dtype=torch.long)
+        ], dim=1).to(device)
+    else:
+        attention_mask = torch.ones_like(input_ids, dtype=torch.long).to(device)
+
+    with torch.no_grad():
+        generated_ids = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=1200,
+            do_sample=True,
+            temperature=0.6,
+            top_p=0.95,
+            repetition_penalty=1.1,
+            num_return_sequences=1,
+            eos_token_id=128258,
+            use_cache=True,
+        )
+
+    # Post-procesado: recortar desde el último token 128257 y limpiar 128258
+    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 = generated_ids[:, last_occurrence_idx+1:]
+    else:
+        cropped = generated_ids
+
+    cleaned = cropped[cropped != token_to_remove]
+
+    # Asegurar múltiplos de 7 y ajustar offset SNAC
+    trimmed = cleaned[: (len(cleaned) // 7) * 7]
+    trimmed = [int(t) - 128266 for t in trimmed]
+
+    audio = decode_snac(trimmed)
+    return (24000, audio)
+
+
+# Gradio UI (simple: texto + speaker)
+with gr.Blocks() as demo:
+    gr.Markdown("# 🗣️ Orpheus Spanish TTS — sin emociones\nSelecciona un *speaker* y escribe el texto.")
+
+    with gr.Row():
+        with gr.Column():
+            text_input = gr.Textbox(label="Texto", placeholder="Escribe aquí el texto a locutar")
+            speaker_dropdown = gr.Dropdown(choices=speakers, value=speakers[0], label="Speaker")
+            submit_btn = gr.Button("Generar audio")
+        with gr.Column():
+            audio_output = gr.Audio(label="Audio generado")
+
+    submit_btn.click(
+        fn=tts,
+        inputs=[text_input, speaker_dropdown],
+        outputs=audio_output,
+    )
+
+if __name__ == "__main__":
+    demo.launch()