Update README.md

README.md

@@ -57,35 +57,100 @@ You can run inference using the demo space: [Orpheus TTS Spanish Fine-Tuned](htt
 To run inference locally with full control:
 
 ```pythonpython
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 from snac import SNAC
 
-base_model = AutoModelForCausalLM.from_pretrained("canopylabs/3b-es_it-pretrain-research_release", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
-tokenizer = AutoTokenizer.from_pretrained("canopylabs/3b-es_it-pretrain-research_release")
-snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz")
-
-prompt = "alloy (intense_fear_dread_apprehension_and_horror): Estoy atrapado, por favor ayúdame."
-input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-output = model.generate(input_ids)
-# Postprocess generated tokens (simplified)
-audio_tokens = output[0].tolist()
-# Trim to multiple of 7, subtract offset, and decode
-trimmed = [t - 128266 for t in audio_tokens if t >= 128266]
-
+# --- Minimal config ---
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+BASE  = "canopylabs/3b-es_it-pretrain-research_release"
+LORA  = "sirekist98/orpheustts_spanish_finetuned"
+SNAC_ID = "hubertsiuzdak/snac_24khz"
+
+VOICE = "alloy"
+EMOTION_ID = "intense_fear_dread_apprehension_horror_terror_panic"
+TEXT = "Estoy atrapado, por favor ayúdame."
+prompt = f"{VOICE} ({EMOTION_ID}): {TEXT}"
+
+# --- Load models ---
+tokenizer  = AutoTokenizer.from_pretrained(BASE)
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE,
+    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32
+)
+model      = PeftModel.from_pretrained(base_model, LORA).to(device).eval()
+snac_model = SNAC.from_pretrained(SNAC_ID).to(device)
+
+# --- Prepare input (same as your Space) ---
+input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+start_tok = torch.tensor([[128259]], dtype=torch.long).to(device)
+end_toks  = torch.tensor([[128009, 128260]], dtype=torch.long).to(device)
+
+input_ids = torch.cat([start_tok, input_ids, end_toks], dim=1)
+MAX_LEN   = 4260
+pad_len   = MAX_LEN - input_ids.shape[1]
+pad       = torch.full((1, pad_len), 128263, dtype=torch.long).to(device)
+input_ids = torch.cat([pad, input_ids], dim=1)
+attention_mask = torch.cat(
+    [torch.zeros((1, pad_len), dtype=torch.long),
+     torch.ones((1, input_ids.shape[1] - pad_len), dtype=torch.long)],
+    dim=1
+).to(device)
+
+# --- Generate ---
+generated = 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-process (find 128257, remove 128258, multiple of 7, subtract 128266) ---
+AUDIO_TOKEN_OFFSET = 128266
+token_to_find      = 128257
+token_to_remove    = 128258
+
+idxs = (generated == token_to_find).nonzero(as_tuple=True)
+cropped = generated[:, idxs[1][-1].item() + 1:] if len(idxs[1]) > 0 else generated
+cleaned = cropped[cropped != token_to_remove]
+codes   = cleaned[: (len(cleaned) // 7) * 7].tolist()
+codes   = [int(t) - AUDIO_TOKEN_OFFSET for t in codes]
+
+# --- SNAC decode (same layout as your Space) ---
 layer_1, layer_2, layer_3 = [], [], []
-for i in range(len(trimmed) // 7):
-    layer_1.append(trimmed[7*i])
-    layer_2.append(trimmed[7*i+1])
-    layer_3.extend(trimmed[7*i+2:7*i+4])
-    layer_2.append(trimmed[7*i+4])
-    layer_3.extend(trimmed[7*i+5:7*i+7])
-
+for i in range((len(codes) + 1) // 7):
+    b = 7 * i
+    if b + 6 >= len(codes):
+        break
+    layer_1.append(codes[b + 0])
+    layer_2.append(codes[b + 1] - 4096)
+    layer_3.append(codes[b + 2] - 2 * 4096)
+    layer_3.append(codes[b + 3] - 3 * 4096)
+    layer_2.append(codes[b + 4] - 4 * 4096)
+    layer_3.append(codes[b + 5] - 5 * 4096)
+    layer_3.append(codes[b + 6] - 6 * 4096)
+
+dev_snac = snac_model.quantizer.quantizers[0].codebook.weight.device
 layers = [
-    torch.tensor(layer_1).unsqueeze(0).to(snac_model.device),
-    torch.tensor(layer_2).unsqueeze(0).to(snac_model.device),
-    torch.tensor(layer_3).unsqueeze(0).to(snac_model.device),
+    torch.tensor(layer_1).unsqueeze(0).to(dev_snac),
+    torch.tensor(layer_2).unsqueeze(0).to(dev_snac),
+    torch.tensor(layer_3).unsqueeze(0).to(dev_snac),
 ]
-audio = snac_model.decode(layers).squeeze().cpu().numpy()
+
+with torch.no_grad():
+    audio = snac_model.decode(layers).squeeze().cpu().numpy()
+
+# 'audio' is the 24kHz waveform.
+# Optional:
+# from scipy.io.wavfile import write as write_wav
+# write_wav("output.wav", 24000, audio)
 ```
 
 ---